cp_tor/

server.rs

//! Search server: handles incoming search requests from remote peers.
//!
//! Reads length-prefixed CBOR `SearchRequests` from any AsyncRead+AsyncWrite
//! stream, performs vector search via the local graph store, and returns
//! signed `SearchResponses`.

use std::sync::Arc;
use std::time::Instant;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

use cp_graph::GraphStore;

use crate::error::{Result, TorError};
use crate::keys::verify_signature;
use crate::rate_limit::RateLimiter;
use crate::types::{
    RemoteSearchResult, SearchRequest, SearchResponse, SearchStatus, MAX_RESULTS, RRF_K,
};
use crate::wire;

/// Configuration for the search server.
pub struct ServerConfig {
    /// The node's Ed25519 signing key seed (for signing responses).
    pub identity_secret: [u8; 32],
    /// The node's Ed25519 public key.
    pub identity_public: [u8; 32],
    /// BLAKE3 hash of the embedding model in use.
    pub model_hash: [u8; 32],
    /// Maximum concurrent queries.
    pub max_concurrent: u8,
    /// Expected embedding dimension (e.g. 1536).
    pub embedding_dim: usize,
}

/// Handle a single incoming search connection.
///
/// Reads a `SearchRequest`, verifies it, performs the search, builds a
/// `SearchResponse`, signs it, and sends it back. This function works
/// with any `AsyncRead` + `AsyncWrite` stream (TCP, Tor `DataStream`, etc.).
///
/// Handles exactly one request-response cycle. For persistent connections
/// that handle multiple messages, use `handle_connection_loop`.
pub async fn handle_connection<S>(
    stream: &mut S,
    graph: &Arc<Mutex<GraphStore>>,
    rate_limiter: &Arc<Mutex<RateLimiter>>,
    config: &ServerConfig,
) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let (mut reader, mut writer) = tokio::io::split(stream);

    // Read the request
    let request: SearchRequest = match wire::read_message(&mut reader).await {
        Ok(req) => req,
        Err(TorError::Keepalive) => {
            debug!("Received keepalive probe");
            return Ok(());
        }
        Err(e) => return Err(e),
    };

    debug!(
        "Received search request {} ({} dims, max_results={})",
        hex::encode(&request.request_id[..4]),
        request.query_embedding.len(),
        request.max_results
    );

    // Verify request signature
    let signing_bytes = request.signing_bytes();
    if let Err(e) = verify_signature(&request.public_key, &signing_bytes, &request.signature) {
        warn!("Invalid request signature: {}", e);
        let response = build_error_response(&request, SearchStatus::InvalidRequest, config);
        wire::write_message(&mut writer, &response).await?;
        return Ok(());
    }

    // Validate request timestamp (replay protection: reject if outside 30s window)
    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;
    if (now_ms - request.timestamp).abs() > 30_000 {
        warn!("Request timestamp outside 30s window");
        let response = build_error_response(&request, SearchStatus::InvalidRequest, config);
        wire::write_message(&mut writer, &response).await?;
        return Ok(());
    }

    // Check rate limit
    let rate_allowed = {
        let mut rl = rate_limiter.lock().await;
        // Periodic cleanup of stale rate limiter entries to prevent unbounded growth
        // from ephemeral session keys. Runs every time we process a request.
        rl.cleanup(now_ms, 120_000);
        rl.check(&request.public_key, now_ms)
    };

    if !rate_allowed {
        debug!(
            "Rate limited requester {}",
            hex::encode(&request.public_key[..4])
        );
        let response = build_error_response(&request, SearchStatus::Overloaded, config);
        wire::write_message(&mut writer, &response).await?;
        return Ok(());
    }

    // Check model compatibility
    if request.model_hash != config.model_hash {
        debug!("Model mismatch from requester");
        let response = build_error_response(&request, SearchStatus::ModelMismatch, config);
        wire::write_message(&mut writer, &response).await?;
        return Ok(());
    }

    // Validate embedding dimensions
    if request.query_embedding.len() != config.embedding_dim {
        warn!(
            "Embedding dimension mismatch: got {}, expected {}",
            request.query_embedding.len(),
            config.embedding_dim
        );
        let response = build_error_response(&request, SearchStatus::InvalidRequest, config);
        wire::write_message(&mut writer, &response).await?;
        return Ok(());
    }

    // Perform search
    let search_start = Instant::now();
    let max_results = request.max_results.min(MAX_RESULTS) as usize;

    let (results, state_root) = {
        let store = graph.lock().await;

        // Convert i16 query embedding to f32 for the HNSW index.
        // Embeddings are quantized with scale 32767 (per CP-010), so we
        // reverse that to get the unit-normalized f32 vector back.
        let query_f32: Vec<f32> = request
            .query_embedding
            .iter()
            .map(|&v| f32::from(v) / 32767.0)
            .collect();

        // Vector search returns Vec<(Uuid, f32)> of (embedding_id, score).
        // The HNSW index stores embedding IDs, so we need to resolve each
        // embedding ID → chunk ID → chunk text and document path.
        let search_hits = store
            .search(&query_f32, max_results)
            .map_err(|e| TorError::InvalidRequest(format!("Search failed: {e}")))?;

        // Build a canonical Merkle tree from all chunk IDs for proof generation.
        // Uses canonical ordering (smaller hash first) so proofs can be verified
        // without directional flags.
        let all_chunk_ids = store
            .get_all_chunk_ids()
            .map_err(|e| TorError::InvalidRequest(format!("Failed to get chunk IDs: {e}")))?;

        let sorted_leaves: Vec<[u8; 32]> = {
            let mut ids = all_chunk_ids.clone();
            ids.sort();
            ids.iter()
                .map(|id| *blake3::hash(id.as_bytes()).as_bytes())
                .collect()
        };
        let chunk_tree_root = canonical_merkle_root(&sorted_leaves);

        // Build index for fast lookup: chunk_id → leaf position
        let mut chunk_id_to_leaf_index: std::collections::HashMap<uuid::Uuid, usize> =
            std::collections::HashMap::new();
        {
            let mut sorted_ids = all_chunk_ids;
            sorted_ids.sort();
            for (i, id) in sorted_ids.iter().enumerate() {
                chunk_id_to_leaf_index.insert(*id, i);
            }
        }

        // Resolve embedding IDs to chunks and document paths
        let mut remote_results = Vec::with_capacity(search_hits.len());
        for (rank, (embedding_id, _similarity)) in search_hits.into_iter().enumerate() {
            // Resolve embedding → chunk
            let Ok(Some(chunk_id)) = store.get_chunk_id_for_embedding(embedding_id) else {
                continue;
            };

            let Ok(Some(chunk)) = store.get_chunk(chunk_id) else {
                continue;
            };

            let doc_path = match store.get_document(chunk.doc_id) {
                Ok(Some(doc)) => doc.path.to_string_lossy().to_string(),
                _ => String::new(),
            };

            // Generate Merkle proof for this chunk
            let proof = chunk_id_to_leaf_index
                .get(&chunk_id)
                .and_then(|&idx| canonical_merkle_proof(&sorted_leaves, idx));

            remote_results.push(RemoteSearchResult {
                chunk_id: *chunk_id.as_bytes(),
                chunk_text: chunk.text,
                document_path: doc_path,
                score: rrf_score(rank),
                merkle_proof: proof,
            });
        }

        (remote_results, chunk_tree_root)
    };

    let search_latency = search_start.elapsed().as_millis() as u16;

    // Build and sign response
    let response = build_ok_response(&request, results, state_root, search_latency, config);

    wire::write_message(&mut writer, &response).await?;

    info!(
        "Responded to search {} with {} results in {}ms",
        hex::encode(&request.request_id[..4]),
        response.results.len(),
        search_latency
    );

    Ok(())
}

/// Handle a persistent connection that may carry multiple messages.
///
/// Loops reading messages from the stream until the remote side disconnects
/// or an unrecoverable error occurs. Keepalive probes are silently handled.
/// This is used by the onion service accept loop for long-lived connections.
pub async fn handle_connection_loop<S>(
    stream: &mut S,
    graph: &Arc<Mutex<GraphStore>>,
    rate_limiter: &Arc<Mutex<RateLimiter>>,
    config: &ServerConfig,
) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    loop {
        match handle_connection(stream, graph, rate_limiter, config).await {
            Ok(()) => {}
            Err(TorError::Io(ref e)) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
                debug!("Peer disconnected");
                return Ok(());
            }
            Err(TorError::Io(ref e)) if e.kind() == std::io::ErrorKind::ConnectionReset => {
                debug!("Peer connection reset");
                return Ok(());
            }
            Err(e) => return Err(e),
        }
    }
}

/// Compute RRF score for a given rank (0-indexed).
fn rrf_score(rank: usize) -> u32 {
    // RRF formula: 1_000_000 / (K + rank + 1)
    // Scale up by 1M for integer precision
    (1_000_000.0 / (f64::from(RRF_K) + rank as f64 + 1.0)) as u32
}

/// Compute a canonical Merkle root from leaf hashes.
///
/// Uses canonical ordering: at each level, the smaller hash goes first.
/// This allows proof verification without directional flags.
fn canonical_merkle_root(leaves: &[[u8; 32]]) -> [u8; 32] {
    if leaves.is_empty() {
        return [0u8; 32];
    }
    if leaves.len() == 1 {
        return leaves[0];
    }
    canonical_merkle_root_recursive(leaves)
}

fn canonical_merkle_root_recursive(hashes: &[[u8; 32]]) -> [u8; 32] {
    if hashes.len() == 1 {
        return hashes[0];
    }

    let mut next_level = Vec::with_capacity(hashes.len().div_ceil(2));
    for chunk in hashes.chunks(2) {
        let mut hasher = blake3::Hasher::new();
        if chunk.len() > 1 {
            if chunk[0] <= chunk[1] {
                hasher.update(&chunk[0]);
                hasher.update(&chunk[1]);
            } else {
                hasher.update(&chunk[1]);
                hasher.update(&chunk[0]);
            }
        } else {
            hasher.update(&chunk[0]);
            hasher.update(&chunk[0]);
        }
        next_level.push(*hasher.finalize().as_bytes());
    }

    canonical_merkle_root_recursive(&next_level)
}

/// Generate a canonical Merkle proof for the leaf at `leaf_index`.
///
/// Returns sibling hashes from leaf to root. The verifier uses canonical
/// ordering (min first) so direction flags are not needed.
fn canonical_merkle_proof(leaves: &[[u8; 32]], leaf_index: usize) -> Option<Vec<[u8; 32]>> {
    if leaf_index >= leaves.len() || leaves.is_empty() {
        return None;
    }
    if leaves.len() == 1 {
        return Some(Vec::new());
    }

    let mut proof = Vec::new();
    let mut level = leaves.to_vec();
    let mut index = leaf_index;

    while level.len() > 1 {
        let sibling_index = if index.is_multiple_of(2) {
            if index + 1 < level.len() {
                index + 1
            } else {
                index
            }
        } else {
            index - 1
        };

        proof.push(level[sibling_index]);

        // Build next level with canonical ordering
        let mut next_level = Vec::with_capacity(level.len().div_ceil(2));
        for chunk in level.chunks(2) {
            let mut hasher = blake3::Hasher::new();
            if chunk.len() > 1 {
                if chunk[0] <= chunk[1] {
                    hasher.update(&chunk[0]);
                    hasher.update(&chunk[1]);
                } else {
                    hasher.update(&chunk[1]);
                    hasher.update(&chunk[0]);
                }
            } else {
                hasher.update(&chunk[0]);
                hasher.update(&chunk[0]);
            }
            next_level.push(*hasher.finalize().as_bytes());
        }

        index /= 2;
        level = next_level;
    }

    Some(proof)
}

fn build_error_response(
    request: &SearchRequest,
    status: SearchStatus,
    config: &ServerConfig,
) -> SearchResponse {
    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;

    let mut response = SearchResponse {
        request_id: request.request_id,
        status,
        results: Vec::new(),
        peer_state_root: [0u8; 32],
        search_latency_ms: 0,
        timestamp: now_ms,
        signature: [0u8; 64],
    };

    let signing_bytes = response.signing_bytes();
    let signing_key = ed25519_dalek::SigningKey::from_bytes(&config.identity_secret);
    response.signature = ed25519_dalek::Signer::sign(&signing_key, &signing_bytes).to_bytes();

    response
}

fn build_ok_response(
    request: &SearchRequest,
    results: Vec<RemoteSearchResult>,
    state_root: [u8; 32],
    search_latency_ms: u16,
    config: &ServerConfig,
) -> SearchResponse {
    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;

    let mut response = SearchResponse {
        request_id: request.request_id,
        status: SearchStatus::Ok,
        results,
        peer_state_root: state_root,
        search_latency_ms,
        timestamp: now_ms,
        signature: [0u8; 64],
    };

    let signing_bytes = response.signing_bytes();
    let signing_key = ed25519_dalek::SigningKey::from_bytes(&config.identity_secret);
    response.signature = ed25519_dalek::Signer::sign(&signing_key, &signing_bytes).to_bytes();

    response
}

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

    #[test]
    fn test_rrf_score_decreasing() {
        let scores: Vec<u32> = (0..10).map(rrf_score).collect();
        for w in scores.windows(2) {
            assert!(w[0] > w[1], "RRF scores should decrease with rank");
        }
    }

    #[test]
    fn test_rrf_score_rank_0() {
        let score = rrf_score(0);
        // 1_000_000 / (60 + 0 + 1) = 1_000_000 / 61 ≈ 16393
        assert_eq!(score, 16393);
    }

    #[test]
    fn test_rrf_score_rank_19() {
        let score = rrf_score(19);
        // 1_000_000 / (60 + 19 + 1) = 1_000_000 / 80 = 12500
        assert_eq!(score, 12500);
    }
}