synor/crates/synor-mining/src/stratum.rs

//! Stratum protocol for pool mining.
//!
//! Implements the Stratum v2 protocol for communication between miners and pools.
//! This allows for efficient work distribution and share submission.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;

use synor_types::Hash256;

use crate::kheavyhash::KHeavyHash;
use crate::template::BlockTemplate;
use crate::{MiningError, Target};

/// Stratum job for mining.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StratumJob {
    /// Job identifier.
    pub job_id: String,
    /// Block header hash (before nonce).
    pub header_hash: String,
    /// Target for shares.
    pub share_target: String,
    /// Target for blocks.
    pub block_target: String,
    /// Timestamp.
    pub timestamp: u64,
    /// Extra nonce 1 (pool-assigned).
    pub extra_nonce1: String,
    /// Extra nonce 2 size.
    pub extra_nonce2_size: usize,
    /// Clean jobs flag.
    pub clean_jobs: bool,
}

impl StratumJob {
    /// Creates a job from a block template.
    pub fn from_template(template: &BlockTemplate, job_id: String, extra_nonce1: &str) -> Self {
        StratumJob {
            job_id,
            header_hash: hex::encode(&template.header_data),
            share_target: hex::encode(template.target),
            block_target: hex::encode(template.target),
            timestamp: template.timestamp,
            extra_nonce1: extra_nonce1.to_string(),
            extra_nonce2_size: 8,
            clean_jobs: true,
        }
    }
}

/// Share submission from a miner.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ShareSubmission {
    /// Worker name.
    pub worker: String,
    /// Job ID.
    pub job_id: String,
    /// Extra nonce 2.
    pub extra_nonce2: String,
    /// Nonce.
    pub nonce: String,
    /// Timestamp.
    pub timestamp: u64,
}

/// Result of share validation.
#[derive(Clone, Debug)]
pub enum ShareResult {
    /// Valid share that meets share target.
    ValidShare,
    /// Valid block that meets block target.
    ValidBlock { pow_hash: Hash256, nonce: u64 },
    /// Invalid share.
    Invalid(String),
    /// Stale share (old job).
    Stale,
    /// Duplicate share.
    Duplicate,
}

/// Stratum message types.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "method", content = "params")]
pub enum StratumRequest {
    /// Subscribe to mining notifications.
    #[serde(rename = "mining.subscribe")]
    Subscribe { agent: String, protocol: String },

    /// Authorize a worker.
    #[serde(rename = "mining.authorize")]
    Authorize { worker: String, password: String },

    /// Submit a share.
    #[serde(rename = "mining.submit")]
    Submit(ShareSubmission),

    /// Suggest difficulty.
    #[serde(rename = "mining.suggest_difficulty")]
    SuggestDifficulty { difficulty: f64 },

    /// Configure connection.
    #[serde(rename = "mining.configure")]
    Configure { extensions: Vec<String> },
}

/// Stratum response.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StratumResponse {
    /// Request ID.
    pub id: u64,
    /// Result data.
    pub result: Option<serde_json::Value>,
    /// Error if any.
    pub error: Option<StratumError>,
}

/// Stratum error.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StratumError {
    /// Error code.
    pub code: i32,
    /// Error message.
    pub message: String,
}

impl StratumError {
    pub fn unauthorized() -> Self {
        StratumError {
            code: 24,
            message: "Unauthorized worker".into(),
        }
    }

    pub fn invalid_share() -> Self {
        StratumError {
            code: 23,
            message: "Invalid share".into(),
        }
    }

    pub fn stale_share() -> Self {
        StratumError {
            code: 21,
            message: "Stale share".into(),
        }
    }

    pub fn duplicate_share() -> Self {
        StratumError {
            code: 22,
            message: "Duplicate share".into(),
        }
    }

    pub fn low_difficulty() -> Self {
        StratumError {
            code: 23,
            message: "Low difficulty share".into(),
        }
    }
}

/// Stratum notification (server to client).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StratumNotification {
    /// Method name.
    pub method: String,
    /// Parameters.
    pub params: serde_json::Value,
}

/// Connected worker information.
#[derive(Clone, Debug)]
pub struct WorkerInfo {
    /// Worker name.
    pub name: String,
    /// Connected since.
    pub connected_at: u64,
    /// Shares submitted.
    pub shares_submitted: u64,
    /// Valid shares.
    pub shares_valid: u64,
    /// Invalid shares.
    pub shares_invalid: u64,
    /// Current difficulty.
    pub difficulty: f64,
    /// Extra nonce 1.
    pub extra_nonce1: String,
    /// Last share time.
    pub last_share_time: u64,
    /// Hashrate estimate.
    pub hashrate: f64,
}

/// Stratum server for pool mining.
pub struct StratumServer {
    /// Server address.
    address: String,
    /// Current job.
    current_job: RwLock<Option<StratumJob>>,
    /// Connected workers.
    workers: RwLock<HashMap<String, WorkerInfo>>,
    /// Extra nonce counter.
    extra_nonce_counter: AtomicU64,
    /// Job counter.
    job_counter: AtomicU64,
    /// Block found channel.
    block_tx: mpsc::Sender<(Hash256, u64)>,
    /// Share target (for vardiff).
    share_target: RwLock<Target>,
    /// Submitted shares (for dedup).
    submitted_shares: RwLock<HashMap<String, u64>>,
    /// PoW hasher for verification.
    hasher: KHeavyHash,
}

impl StratumServer {
    /// Creates a new stratum server.
    pub fn new(address: String, block_tx: mpsc::Sender<(Hash256, u64)>) -> Self {
        StratumServer {
            address,
            current_job: RwLock::new(None),
            workers: RwLock::new(HashMap::new()),
            extra_nonce_counter: AtomicU64::new(1),
            job_counter: AtomicU64::new(1),
            block_tx,
            share_target: RwLock::new(Target::max()),
            submitted_shares: RwLock::new(HashMap::new()),
            hasher: KHeavyHash::new(),
        }
    }

    /// Updates the current job from a template.
    pub fn update_job(&self, template: &BlockTemplate) {
        let job_id = format!("{:x}", self.job_counter.fetch_add(1, Ordering::Relaxed));
        let job = StratumJob::from_template(template, job_id, "00000000");
        *self.current_job.write() = Some(job);
    }

    /// Gets the current job.
    pub fn current_job(&self) -> Option<StratumJob> {
        self.current_job.read().clone()
    }

    /// Validates a share submission.
    pub fn validate_share(&self, submission: &ShareSubmission) -> ShareResult {
        // Check if job exists
        let job = match self.current_job.read().clone() {
            Some(j) => j,
            None => return ShareResult::Invalid("No active job".into()),
        };

        // Check job ID
        if submission.job_id != job.job_id {
            return ShareResult::Stale;
        }

        // Check for duplicate
        let share_key = format!(
            "{}:{}:{}",
            submission.job_id, submission.extra_nonce2, submission.nonce
        );
        {
            let mut shares = self.submitted_shares.write();
            if shares.contains_key(&share_key) {
                return ShareResult::Duplicate;
            }
            shares.insert(share_key, submission.timestamp);
        }

        // Parse nonce
        let nonce = match u64::from_str_radix(&submission.nonce, 16) {
            Ok(n) => n,
            Err(_) => return ShareResult::Invalid("Invalid nonce".into()),
        };

        // Parse extra_nonce2
        let extra_nonce2 = match hex::decode(&submission.extra_nonce2) {
            Ok(bytes) => bytes,
            Err(_) => return ShareResult::Invalid("Invalid extra_nonce2".into()),
        };

        // Parse header hash from job
        let header_hash = match hex::decode(&job.header_hash) {
            Ok(bytes) => bytes,
            Err(_) => return ShareResult::Invalid("Invalid header hash in job".into()),
        };

        // Parse extra_nonce1 from job
        let extra_nonce1 = match hex::decode(&job.extra_nonce1) {
            Ok(bytes) => bytes,
            Err(_) => return ShareResult::Invalid("Invalid extra_nonce1 in job".into()),
        };

        // Reconstruct header: header_hash + extra_nonce1 + extra_nonce2
        let mut header_data =
            Vec::with_capacity(header_hash.len() + extra_nonce1.len() + extra_nonce2.len());
        header_data.extend_from_slice(&header_hash);
        header_data.extend_from_slice(&extra_nonce1);
        header_data.extend_from_slice(&extra_nonce2);

        // Compute PoW hash
        let pow_hash = self.hasher.hash(&header_data, nonce);

        // Parse share target
        let share_target_bytes: [u8; 32] = match hex::decode(&job.share_target) {
            Ok(bytes) if bytes.len() == 32 => {
                let mut arr = [0u8; 32];
                arr.copy_from_slice(&bytes);
                arr
            }
            _ => return ShareResult::Invalid("Invalid share target".into()),
        };
        let share_target = Target::from_bytes(share_target_bytes);

        // Parse block target
        let block_target_bytes: [u8; 32] = match hex::decode(&job.block_target) {
            Ok(bytes) if bytes.len() == 32 => {
                let mut arr = [0u8; 32];
                arr.copy_from_slice(&bytes);
                arr
            }
            _ => return ShareResult::Invalid("Invalid block target".into()),
        };
        let block_target = Target::from_bytes(block_target_bytes);

        // Check if meets block target (higher priority)
        if block_target.is_met_by(&pow_hash.hash) {
            return ShareResult::ValidBlock {
                pow_hash: pow_hash.hash,
                nonce,
            };
        }

        // Check if meets share target
        if share_target.is_met_by(&pow_hash.hash) {
            return ShareResult::ValidShare;
        }

        // Hash doesn't meet either target
        ShareResult::Invalid("Share does not meet target difficulty".into())
    }

    /// Allocates an extra nonce for a new worker.
    pub fn allocate_extra_nonce(&self) -> String {
        let nonce = self.extra_nonce_counter.fetch_add(1, Ordering::Relaxed);
        format!("{:016x}", nonce)
    }

    /// Registers a worker.
    pub fn register_worker(&self, name: String, extra_nonce1: String) {
        let info = WorkerInfo {
            name: name.clone(),
            connected_at: current_timestamp(),
            shares_submitted: 0,
            shares_valid: 0,
            shares_invalid: 0,
            difficulty: 1.0,
            extra_nonce1,
            last_share_time: 0,
            hashrate: 0.0,
        };
        self.workers.write().insert(name, info);
    }

    /// Updates worker stats after share submission.
    pub fn update_worker_stats(&self, worker: &str, valid: bool) {
        if let Some(info) = self.workers.write().get_mut(worker) {
            info.shares_submitted += 1;
            if valid {
                info.shares_valid += 1;
            } else {
                info.shares_invalid += 1;
            }
            info.last_share_time = current_timestamp();
        }
    }

    /// Gets worker information.
    pub fn get_worker(&self, name: &str) -> Option<WorkerInfo> {
        self.workers.read().get(name).cloned()
    }

    /// Gets all workers.
    pub fn workers(&self) -> Vec<WorkerInfo> {
        self.workers.read().values().cloned().collect()
    }

    /// Runs the stratum server.
    pub async fn run(self: Arc<Self>) -> Result<(), MiningError> {
        let listener = TcpListener::bind(&self.address)
            .await
            .map_err(|e| MiningError::Stratum(e.to_string()))?;

        tracing::info!("Stratum server listening on {}", self.address);

        loop {
            match listener.accept().await {
                Ok((socket, addr)) => {
                    tracing::info!("New stratum connection from {}", addr);
                    let server = Arc::clone(&self);
                    tokio::spawn(async move {
                        if let Err(e) = server.handle_connection(socket).await {
                            tracing::warn!("Stratum connection error: {}", e);
                        }
                    });
                }
                Err(e) => {
                    tracing::error!("Accept error: {}", e);
                }
            }
        }
    }

    async fn handle_connection(self: Arc<Self>, socket: TcpStream) -> Result<(), MiningError> {
        let (reader, mut writer) = socket.into_split();
        let mut reader = BufReader::new(reader);
        let mut line = String::new();

        let mut _worker_name: Option<String> = None;
        let mut authorized = false;

        loop {
            line.clear();
            match reader.read_line(&mut line).await {
                Ok(0) => break, // EOF
                Ok(_) => {
                    // Parse JSON-RPC request
                    let request: serde_json::Value = match serde_json::from_str(&line) {
                        Ok(v) => v,
                        Err(e) => {
                            tracing::warn!("Invalid JSON: {}", e);
                            continue;
                        }
                    };

                    let id = request.get("id").and_then(|v| v.as_u64()).unwrap_or(0);
                    let method = request.get("method").and_then(|v| v.as_str()).unwrap_or("");

                    let response = match method {
                        "mining.subscribe" => {
                            let extra_nonce = self.allocate_extra_nonce();
                            let result =
                                serde_json::json!([[["mining.notify", "1"]], extra_nonce, 8]);
                            StratumResponse {
                                id,
                                result: Some(result),
                                error: None,
                            }
                        }
                        "mining.authorize" => {
                            let params = request.get("params").cloned().unwrap_or_default();
                            let worker =
                                params.get(0).and_then(|v| v.as_str()).unwrap_or("unknown");
                            _worker_name = Some(worker.to_string());
                            authorized = true;
                            self.register_worker(worker.to_string(), self.allocate_extra_nonce());

                            StratumResponse {
                                id,
                                result: Some(serde_json::json!(true)),
                                error: None,
                            }
                        }
                        "mining.submit" => {
                            if !authorized {
                                StratumResponse {
                                    id,
                                    result: None,
                                    error: Some(StratumError::unauthorized()),
                                }
                            } else {
                                // Parse submission
                                let params = request.get("params").cloned().unwrap_or_default();
                                let submission = ShareSubmission {
                                    worker: params
                                        .get(0)
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("")
                                        .to_string(),
                                    job_id: params
                                        .get(1)
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("")
                                        .to_string(),
                                    extra_nonce2: params
                                        .get(2)
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("")
                                        .to_string(),
                                    nonce: params
                                        .get(3)
                                        .and_then(|v| v.as_str())
                                        .unwrap_or("")
                                        .to_string(),
                                    timestamp: current_timestamp(),
                                };

                                let result = self.validate_share(&submission);
                                let (accepted, error) = match result {
                                    ShareResult::ValidShare | ShareResult::ValidBlock { .. } => {
                                        self.update_worker_stats(&submission.worker, true);
                                        (true, None)
                                    }
                                    ShareResult::Stale => {
                                        self.update_worker_stats(&submission.worker, false);
                                        (false, Some(StratumError::stale_share()))
                                    }
                                    ShareResult::Duplicate => {
                                        self.update_worker_stats(&submission.worker, false);
                                        (false, Some(StratumError::duplicate_share()))
                                    }
                                    ShareResult::Invalid(_) => {
                                        self.update_worker_stats(&submission.worker, false);
                                        (false, Some(StratumError::invalid_share()))
                                    }
                                };

                                StratumResponse {
                                    id,
                                    result: Some(serde_json::json!(accepted)),
                                    error,
                                }
                            }
                        }
                        _ => {
                            tracing::debug!("Unknown method: {}", method);
                            continue;
                        }
                    };

                    let response_str = serde_json::to_string(&response)
                        .map_err(|e| MiningError::Stratum(e.to_string()))?;
                    writer
                        .write_all(format!("{}\n", response_str).as_bytes())
                        .await
                        .map_err(|e| MiningError::Stratum(e.to_string()))?;
                }
                Err(e) => {
                    return Err(MiningError::Stratum(e.to_string()));
                }
            }
        }

        Ok(())
    }
}

/// Stratum client for connecting to pools.
pub struct StratumClient {
    /// Pool address.
    pool_address: String,
    /// Worker name.
    worker_name: String,
    /// Password.
    password: String,
    /// Connection state.
    connected: std::sync::atomic::AtomicBool,
    /// Current job.
    current_job: RwLock<Option<StratumJob>>,
    /// Extra nonce 1 from pool.
    extra_nonce1: RwLock<String>,
    /// Extra nonce 2 size.
    extra_nonce2_size: std::sync::atomic::AtomicUsize,
}

impl StratumClient {
    /// Creates a new stratum client.
    pub fn new(pool_address: String, worker_name: String, password: String) -> Self {
        StratumClient {
            pool_address,
            worker_name,
            password,
            connected: std::sync::atomic::AtomicBool::new(false),
            current_job: RwLock::new(None),
            extra_nonce1: RwLock::new(String::new()),
            extra_nonce2_size: std::sync::atomic::AtomicUsize::new(8),
        }
    }

    /// Checks if connected.
    pub fn is_connected(&self) -> bool {
        self.connected.load(Ordering::Relaxed)
    }

    /// Gets the current job.
    pub fn current_job(&self) -> Option<StratumJob> {
        self.current_job.read().clone()
    }

    /// Gets extra nonce 1.
    pub fn extra_nonce1(&self) -> String {
        self.extra_nonce1.read().clone()
    }

    /// Connects to the pool.
    pub async fn connect(&self) -> Result<(), MiningError> {
        let stream = TcpStream::connect(&self.pool_address)
            .await
            .map_err(|e| MiningError::Stratum(e.to_string()))?;

        self.connected.store(true, Ordering::Relaxed);

        // Send subscribe
        let subscribe = serde_json::json!({
            "id": 1,
            "method": "mining.subscribe",
            "params": ["synor-miner/0.1.0", null]
        });

        let (_reader, mut writer) = stream.into_split();
        writer
            .write_all(format!("{}\n", subscribe).as_bytes())
            .await
            .map_err(|e| MiningError::Stratum(e.to_string()))?;

        // Send authorize
        let authorize = serde_json::json!({
            "id": 2,
            "method": "mining.authorize",
            "params": [&self.worker_name, &self.password]
        });
        writer
            .write_all(format!("{}\n", authorize).as_bytes())
            .await
            .map_err(|e| MiningError::Stratum(e.to_string()))?;

        Ok(())
    }

    /// Submits a share.
    pub async fn submit_share(
        &self,
        _job_id: &str,
        _extra_nonce2: u64,
        _nonce: u64,
    ) -> Result<bool, MiningError> {
        // This would send via the open connection
        // Simplified for now
        Ok(true)
    }
}

/// Gets current Unix timestamp in milliseconds.
fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as u64
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_stratum_job() {
        let job = StratumJob {
            job_id: "1".to_string(),
            header_hash: "0".repeat(64),
            share_target: "f".repeat(64),
            block_target: "0".repeat(64),
            timestamp: 1234567890,
            extra_nonce1: "00000001".to_string(),
            extra_nonce2_size: 8,
            clean_jobs: true,
        };

        assert_eq!(job.job_id, "1");
        assert!(job.clean_jobs);
    }

    #[test]
    fn test_worker_info() {
        let info = WorkerInfo {
            name: "worker1".to_string(),
            connected_at: 1000,
            shares_submitted: 100,
            shares_valid: 95,
            shares_invalid: 5,
            difficulty: 1.0,
            extra_nonce1: "00000001".to_string(),
            last_share_time: 2000,
            hashrate: 1000.0,
        };

        assert_eq!(
            info.shares_valid + info.shares_invalid,
            info.shares_submitted
        );
    }

    #[test]
    fn test_stratum_error() {
        let err = StratumError::invalid_share();
        assert_eq!(err.code, 23);
    }

    #[test]
    fn test_share_validation() {
        use tokio::sync::mpsc;

        let (tx, _rx) = mpsc::channel(10);
        let server = StratumServer::new("127.0.0.1:16111".to_string(), tx.clone());

        // Create a job with easy targets
        let job = StratumJob {
            job_id: "1".to_string(),
            header_hash: "00".repeat(32),  // 32 zero bytes
            share_target: "ff".repeat(32), // All 0xff = easy target
            block_target: "00".repeat(32), // All 0x00 = impossible target
            timestamp: 1234567890,
            extra_nonce1: "0000000000000001".to_string(),
            extra_nonce2_size: 8,
            clean_jobs: true,
        };
        *server.current_job.write() = Some(job);

        // Test with no job
        let server2 = StratumServer::new("127.0.0.1:16112".to_string(), tx.clone());
        let submission = ShareSubmission {
            worker: "test".to_string(),
            job_id: "1".to_string(),
            extra_nonce2: "0000000000000001".to_string(),
            nonce: "0000000000000001".to_string(),
            timestamp: 1234567890,
        };
        match server2.validate_share(&submission) {
            ShareResult::Invalid(msg) => assert!(msg.contains("No active job")),
            _ => panic!("Expected Invalid result for no job"),
        }

        // Test stale job
        let stale_submission = ShareSubmission {
            worker: "test".to_string(),
            job_id: "wrong_id".to_string(),
            extra_nonce2: "0000000000000001".to_string(),
            nonce: "0000000000000001".to_string(),
            timestamp: 1234567890,
        };
        match server.validate_share(&stale_submission) {
            ShareResult::Stale => {}
            _ => panic!("Expected Stale result"),
        }

        // Test valid share with easy share target
        let valid_submission = ShareSubmission {
            worker: "test".to_string(),
            job_id: "1".to_string(),
            extra_nonce2: "0000000000000001".to_string(),
            nonce: "0000000000000001".to_string(),
            timestamp: 1234567890,
        };
        let result = server.validate_share(&valid_submission);
        assert!(
            matches!(result, ShareResult::ValidShare),
            "Expected ValidShare, got {:?}",
            result
        );

        // Test duplicate share
        let duplicate_submission = ShareSubmission {
            worker: "test".to_string(),
            job_id: "1".to_string(),
            extra_nonce2: "0000000000000001".to_string(),
            nonce: "0000000000000001".to_string(),
            timestamp: 1234567890,
        };
        match server.validate_share(&duplicate_submission) {
            ShareResult::Duplicate => {}
            _ => panic!("Expected Duplicate result"),
        }
    }
}