stemedb/crates/stemedb-rpc/src/server.rs

//! gRPC server implementation for the sync service.
//!
//! This module provides the server-side handlers for sync operations.
//! The actual storage and sync logic is injected via traits to allow
//! flexible deployment configurations.
//!
//! # Example
//!
//! ```ignore
//! use stemedb_rpc::server::{SyncServiceHandler, SyncStorage};
//! use tonic::transport::Server;
//!
//! let storage = MyStorage::new(...);
//! let handler = SyncServiceHandler::new(storage);
//!
//! Server::builder()
//!     .add_service(SyncServiceServer::new(handler))
//!     .serve(addr)
//!     .await?;
//! ```

use crate::proto::sync_service_server::SyncService;
use crate::proto::{
    AssertionData, FetchRequest, FetchResponse, GetLeavesRequest, GetLeavesResponse, GossipRequest,
    GossipResponse, PingRequest, PingResponse, RootExchangeRequest, RootExchangeResponse,
};
use async_trait::async_trait;
use std::sync::Arc;
use tonic::{Request, Response, Status};
use tracing::{debug, info, instrument, warn};

/// Backend storage interface for sync operations.
///
/// Implement this trait to connect the sync service to your storage layer.
#[async_trait]
pub trait SyncStorage: Send + Sync + 'static {
    /// Store an assertion received via gossip.
    ///
    /// Returns Ok(true) if stored, Ok(false) if already existed.
    async fn store_gossip_assertion(
        &self,
        hash: [u8; 32],
        data: Vec<u8>,
        hlc_time: u64,
        hlc_counter: u32,
        hlc_node_id: [u8; 16],
    ) -> Result<bool, String>;

    /// Get the current Merkle root and assertion count.
    async fn get_merkle_state(&self) -> Result<(Option<[u8; 32]>, u64), String>;

    /// Fetch assertions by hash.
    ///
    /// Returns (hash, data) pairs for assertions that exist.
    async fn fetch_assertions(
        &self,
        hashes: Vec<[u8; 32]>,
    ) -> Result<Vec<([u8; 32], Vec<u8>)>, String>;

    /// Get this node's ID and assertion count for ping response.
    async fn get_node_info(&self) -> Result<([u8; 16], u64), String>;

    /// Get all Merkle tree leaf hashes.
    ///
    /// Returns up to `max_leaves` hashes (0 = no limit, capped at 10000).
    async fn get_leaves(&self, max_leaves: u64) -> Result<(Vec<[u8; 32]>, bool), String>;
}

/// gRPC service handler for sync operations.
pub struct SyncServiceHandler<S> {
    storage: Arc<S>,
}

impl<S: SyncStorage> SyncServiceHandler<S> {
    /// Create a new sync service handler with the given storage backend.
    pub fn new(storage: Arc<S>) -> Self {
        Self { storage }
    }
}

#[async_trait]
impl<S: SyncStorage> SyncService for SyncServiceHandler<S> {
    #[instrument(skip(self, request), fields(hash_len = request.get_ref().assertion_hash.len()))]
    async fn gossip(
        &self,
        request: Request<GossipRequest>,
    ) -> Result<Response<GossipResponse>, Status> {
        let req = request.into_inner();

        // Validate hash length
        if req.assertion_hash.len() != 32 {
            return Err(Status::invalid_argument(format!(
                "assertion_hash must be 32 bytes, got {}",
                req.assertion_hash.len()
            )));
        }

        // Validate HLC node ID length
        if req.hlc_node_id.len() != 16 {
            return Err(Status::invalid_argument(format!(
                "hlc_node_id must be 16 bytes, got {}",
                req.hlc_node_id.len()
            )));
        }

        let mut hash = [0u8; 32];
        hash.copy_from_slice(&req.assertion_hash);

        let mut hlc_node_id = [0u8; 16];
        hlc_node_id.copy_from_slice(&req.hlc_node_id);

        debug!(hash = %hex::encode(&hash[..8]), "Received gossip");

        match self
            .storage
            .store_gossip_assertion(
                hash,
                req.assertion_data,
                req.hlc_time,
                req.hlc_counter,
                hlc_node_id,
            )
            .await
        {
            Ok(stored) => {
                if stored {
                    info!(hash = %hex::encode(&hash[..8]), "Stored gossiped assertion");
                } else {
                    debug!(hash = %hex::encode(&hash[..8]), "Duplicate gossip (already stored)");
                }
                Ok(Response::new(GossipResponse { accepted: true, error: String::new() }))
            }
            Err(e) => {
                warn!(error = %e, "Failed to store gossiped assertion");
                Ok(Response::new(GossipResponse { accepted: false, error: e }))
            }
        }
    }

    #[instrument(skip(self, request), fields(assertion_count = request.get_ref().assertion_count))]
    async fn exchange_roots(
        &self,
        request: Request<RootExchangeRequest>,
    ) -> Result<Response<RootExchangeResponse>, Status> {
        let req = request.into_inner();

        // Validate root length if provided
        if !req.merkle_root.is_empty() && req.merkle_root.len() != 32 {
            return Err(Status::invalid_argument(format!(
                "merkle_root must be 32 bytes if provided, got {}",
                req.merkle_root.len()
            )));
        }

        let (local_root, local_count) =
            self.storage.get_merkle_state().await.map_err(Status::internal)?;

        let remote_root: Option<[u8; 32]> = if req.merkle_root.len() == 32 {
            let mut root = [0u8; 32];
            root.copy_from_slice(&req.merkle_root);
            Some(root)
        } else {
            None
        };

        let roots_match = match (&local_root, &remote_root) {
            (Some(local), Some(remote)) => local == remote,
            (None, None) => true,
            _ => false,
        };

        debug!(
            local_count,
            remote_count = req.assertion_count,
            roots_match,
            "Exchanged Merkle roots"
        );

        Ok(Response::new(RootExchangeResponse {
            merkle_root: local_root.map(|r| r.to_vec()).unwrap_or_default(),
            assertion_count: local_count,
            roots_match,
        }))
    }

    #[instrument(skip(self, request), fields(hash_count = request.get_ref().hashes.len()))]
    async fn fetch_assertions(
        &self,
        request: Request<FetchRequest>,
    ) -> Result<Response<FetchResponse>, Status> {
        let req = request.into_inner();

        // Limit request size to prevent abuse
        const MAX_HASHES: usize = 1000;
        if req.hashes.len() > MAX_HASHES {
            return Err(Status::invalid_argument(format!(
                "Too many hashes requested: {} > {}",
                req.hashes.len(),
                MAX_HASHES
            )));
        }

        // Convert and validate hashes
        let mut hashes = Vec::with_capacity(req.hashes.len());
        for (i, hash_bytes) in req.hashes.iter().enumerate() {
            if hash_bytes.len() != 32 {
                return Err(Status::invalid_argument(format!(
                    "hash[{}] must be 32 bytes, got {}",
                    i,
                    hash_bytes.len()
                )));
            }
            let mut hash = [0u8; 32];
            hash.copy_from_slice(hash_bytes);
            hashes.push(hash);
        }

        let results = self.storage.fetch_assertions(hashes).await.map_err(Status::internal)?;

        debug!(requested = req.hashes.len(), found = results.len(), "Fetched assertions");

        let assertions = results
            .into_iter()
            .map(|(hash, data)| AssertionData { hash: hash.to_vec(), data })
            .collect();

        Ok(Response::new(FetchResponse { assertions }))
    }

    #[instrument(skip(self, _request))]
    async fn ping(&self, _request: Request<PingRequest>) -> Result<Response<PingResponse>, Status> {
        let (node_id, assertion_count) =
            self.storage.get_node_info().await.map_err(Status::internal)?;

        debug!(assertion_count, "Responding to ping");

        Ok(Response::new(PingResponse { node_id: node_id.to_vec(), assertion_count }))
    }

    #[instrument(skip(self, request), fields(max_leaves = request.get_ref().max_leaves))]
    async fn get_leaves(
        &self,
        request: Request<GetLeavesRequest>,
    ) -> Result<Response<GetLeavesResponse>, Status> {
        let req = request.into_inner();

        let (leaves, truncated) =
            self.storage.get_leaves(req.max_leaves).await.map_err(Status::internal)?;

        debug!(leaf_count = leaves.len(), truncated, "Returning Merkle leaves");

        Ok(Response::new(GetLeavesResponse {
            leaves: leaves.into_iter().map(|l| l.to_vec()).collect(),
            truncated,
        }))
    }
}

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

    /// Mock storage for testing.
    struct MockStorage {
        node_id: [u8; 16],
        assertion_count: u64,
    }

    #[async_trait]
    impl SyncStorage for MockStorage {
        async fn store_gossip_assertion(
            &self,
            _hash: [u8; 32],
            _data: Vec<u8>,
            _hlc_time: u64,
            _hlc_counter: u32,
            _hlc_node_id: [u8; 16],
        ) -> Result<bool, String> {
            Ok(true)
        }

        async fn get_merkle_state(&self) -> Result<(Option<[u8; 32]>, u64), String> {
            Ok((Some([1u8; 32]), self.assertion_count))
        }

        async fn fetch_assertions(
            &self,
            hashes: Vec<[u8; 32]>,
        ) -> Result<Vec<([u8; 32], Vec<u8>)>, String> {
            // Return mock data for each hash
            Ok(hashes.into_iter().map(|h| (h, vec![1, 2, 3])).collect())
        }

        async fn get_node_info(&self) -> Result<([u8; 16], u64), String> {
            Ok((self.node_id, self.assertion_count))
        }

        async fn get_leaves(&self, max_leaves: u64) -> Result<(Vec<[u8; 32]>, bool), String> {
            let all_leaves = vec![[1u8; 32], [2u8; 32], [3u8; 32]];
            if max_leaves > 0 && (max_leaves as usize) < all_leaves.len() {
                Ok((all_leaves.into_iter().take(max_leaves as usize).collect(), true))
            } else {
                Ok((all_leaves, false))
            }
        }
    }

    #[tokio::test]
    async fn test_ping() {
        let storage = Arc::new(MockStorage { node_id: [42u8; 16], assertion_count: 100 });
        let handler = SyncServiceHandler::new(storage);

        let request = Request::new(PingRequest { node_id: vec![1u8; 16] });
        let response = handler.ping(request).await.expect("ping should succeed");

        assert_eq!(response.get_ref().node_id, vec![42u8; 16]);
        assert_eq!(response.get_ref().assertion_count, 100);
    }

    #[tokio::test]
    async fn test_gossip_invalid_hash_length() {
        let storage = Arc::new(MockStorage { node_id: [1u8; 16], assertion_count: 0 });
        let handler = SyncServiceHandler::new(storage);

        let request = Request::new(GossipRequest {
            assertion_hash: vec![1u8; 16], // Wrong length
            assertion_data: vec![],
            hlc_time: 0,
            hlc_counter: 0,
            hlc_node_id: vec![1u8; 16],
        });

        let result = handler.gossip(request).await;
        assert!(result.is_err());
        assert_eq!(result.err().map(|e| e.code()), Some(tonic::Code::InvalidArgument));
    }

    #[tokio::test]
    async fn test_fetch_too_many_hashes() {
        let storage = Arc::new(MockStorage { node_id: [1u8; 16], assertion_count: 0 });
        let handler = SyncServiceHandler::new(storage);

        let request = Request::new(FetchRequest {
            hashes: vec![vec![0u8; 32]; 1001], // More than MAX_HASHES
        });

        let result = handler.fetch_assertions(request).await;
        assert!(result.is_err());
        assert_eq!(result.err().map(|e| e.code()), Some(tonic::Code::InvalidArgument));
    }
}