cp_tor/

types.rs

//! CP-013 protocol data structures for live anonymous search.

use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;

/// Fixed Canon search protocol port.
pub const CANON_PORT: u16 = 9735;

/// Maximum wire message size (1MB).
pub const MAX_MESSAGE_SIZE: u32 = 1_048_576;

/// Peer registration expiry (72 hours in seconds).
pub const PEER_EXPIRY_SECS: i64 = 72 * 3600;

/// Peer registration renewal interval (24 hours in seconds).
pub const PEER_RENEWAL_SECS: i64 = 24 * 3600;

/// Circuit pool size.
pub const CIRCUIT_POOL_SIZE: usize = 8;

/// Keepalive interval in seconds.
pub const KEEPALIVE_INTERVAL_SECS: u64 = 300;

/// Circuit rotation interval in seconds.
pub const CIRCUIT_ROTATION_SECS: u64 = 1800;

/// Live search timeout in milliseconds.
pub const SEARCH_TIMEOUT_MS: u64 = 2000;

/// Maximum results per search.
pub const MAX_RESULTS: u8 = 20;

/// Rate limit: requests per minute per requester key.
pub const RATE_LIMIT_PER_MIN: u32 = 10;

/// RRF K parameter.
pub const RRF_K: u32 = 60;

/// Number of peers to query per search.
pub const SEARCH_FANOUT: usize = 3;

/// Maximum random jitter before dispatch (milliseconds).
pub const DISPATCH_JITTER_MS: u64 = 200;

// ============================================================================
// Peer Registration (Section 13)
// ============================================================================

/// Capabilities advertised by a peer node.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerCapabilities {
    /// Maximum concurrent queries this peer will handle.
    pub max_concurrent_queries: u8,
    /// Maximum results per query this peer will return.
    pub max_results: u8,
    /// Number of indexed chunks on this peer.
    pub chunk_count: u64,
    /// Whether this peer supports Merkle inclusion proofs.
    pub supports_proofs: bool,
}

impl Default for PeerCapabilities {
    fn default() -> Self {
        Self {
            max_concurrent_queries: 4,
            max_results: 20,
            chunk_count: 0,
            supports_proofs: true,
        }
    }
}

/// Peer registration uploaded to Arweave for discovery.
///
/// Each node publishes its onion address, capabilities, and topic coverage.
/// Registrations expire after 72 hours and should be renewed every 24 hours.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerRegistration {
    /// v3 onion address (56 characters, no .onion suffix).
    pub onion_address: String,
    /// Node ID: `BLAKE3(public_key)`[0..16].
    pub node_id: [u8; 16],
    /// Ed25519 public key of this node.
    pub public_key: [u8; 32],
    /// Advertised capabilities.
    pub capabilities: PeerCapabilities,
    /// Topic coverage tags.
    pub topics: Vec<String>,
    /// BLAKE3 hash of the embedding model weights used by this node.
    pub embedding_model: [u8; 32],
    /// Registration timestamp (Unix milliseconds).
    pub timestamp: i64,
    /// Ed25519 signature over BLAKE3(CBOR(preceding fields)).
    #[serde(with = "BigArray")]
    pub signature: [u8; 64],
}

impl PeerRegistration {
    /// Compute the signing bytes: BLAKE3 of CBOR of all fields except signature.
    pub fn signing_bytes(&self) -> [u8; 32] {
        let signable = PeerRegistrationSignable {
            onion_address: &self.onion_address,
            node_id: &self.node_id,
            public_key: &self.public_key,
            capabilities: &self.capabilities,
            topics: &self.topics,
            embedding_model: &self.embedding_model,
            timestamp: self.timestamp,
        };
        let mut buf = Vec::new();
        ciborium::into_writer(&signable, &mut buf).expect("CBOR serialization cannot fail");
        *blake3::hash(&buf).as_bytes()
    }

    /// Check if this registration has expired relative to the given time.
    /// Also rejects registrations with timestamps more than 60s in the future.
    pub fn is_expired(&self, now_ms: i64) -> bool {
        if self.timestamp > now_ms + 60_000 {
            return true;
        }
        let age_secs = (now_ms - self.timestamp) / 1000;
        age_secs > PEER_EXPIRY_SECS
    }
}

/// Signable portion of `PeerRegistration` (excludes signature).
#[derive(Serialize)]
struct PeerRegistrationSignable<'a> {
    onion_address: &'a str,
    node_id: &'a [u8; 16],
    public_key: &'a [u8; 32],
    capabilities: &'a PeerCapabilities,
    topics: &'a Vec<String>,
    embedding_model: &'a [u8; 32],
    timestamp: i64,
}

// ============================================================================
// Search Protocol (Section 14)
// ============================================================================

/// Search request sent from querier to peer over Tor.
///
/// Uses ephemeral session keys (not the node's registered identity) to
/// prevent peers from linking queries to registered node identities.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchRequest {
    /// Random request identifier.
    pub request_id: [u8; 16],
    /// Query embedding vector (i16 quantized, 1536 dimensions).
    pub query_embedding: Vec<i16>,
    /// Optional query text for lexical search.
    pub query_text: Option<String>,
    /// Maximum results to return (capped at 20).
    pub max_results: u8,
    /// Whether to include Merkle inclusion proofs.
    pub include_proofs: bool,
    /// BLAKE3 hash of the embedding model used.
    pub model_hash: [u8; 32],
    /// Request timestamp (Unix milliseconds).
    pub timestamp: i64,
    /// Ed25519 signature over BLAKE3(CBOR(preceding fields)).
    #[serde(with = "BigArray")]
    pub signature: [u8; 64],
    /// Ephemeral session public key (NOT the node identity key).
    pub public_key: [u8; 32],
}

impl SearchRequest {
    /// Compute the signing bytes: BLAKE3 of CBOR of all fields except signature.
    pub fn signing_bytes(&self) -> [u8; 32] {
        let signable = SearchRequestSignable {
            request_id: &self.request_id,
            query_embedding: &self.query_embedding,
            query_text: &self.query_text,
            max_results: self.max_results,
            include_proofs: self.include_proofs,
            model_hash: &self.model_hash,
            timestamp: self.timestamp,
            public_key: &self.public_key,
        };
        let mut buf = Vec::new();
        ciborium::into_writer(&signable, &mut buf).expect("CBOR serialization cannot fail");
        *blake3::hash(&buf).as_bytes()
    }
}

#[derive(Serialize)]
struct SearchRequestSignable<'a> {
    request_id: &'a [u8; 16],
    query_embedding: &'a Vec<i16>,
    query_text: &'a Option<String>,
    max_results: u8,
    include_proofs: bool,
    model_hash: &'a [u8; 32],
    timestamp: i64,
    public_key: &'a [u8; 32],
}

/// Status code in a search response.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SearchStatus {
    /// Search completed successfully.
    Ok,
    /// Peer uses a different embedding model.
    ModelMismatch,
    /// Peer is at capacity.
    Overloaded,
    /// Request was malformed or signature verification failed.
    InvalidRequest,
}

/// A single search result from a remote peer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteSearchResult {
    /// Chunk identifier.
    pub chunk_id: [u8; 16],
    /// Chunk text content.
    pub chunk_text: String,
    /// Document path (relative).
    pub document_path: String,
    /// RRF score (CP-012 integer scale).
    pub score: u32,
    /// Optional Merkle inclusion proof against `peer_state_root`.
    pub merkle_proof: Option<Vec<[u8; 32]>>,
}

/// Search response sent from peer back to querier.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResponse {
    /// Echoed request identifier.
    pub request_id: [u8; 16],
    /// Status of the search.
    pub status: SearchStatus,
    /// Search results.
    pub results: Vec<RemoteSearchResult>,
    /// Peer's current Merkle state root.
    pub peer_state_root: [u8; 32],
    /// Time spent searching (milliseconds).
    pub search_latency_ms: u16,
    /// Response timestamp (Unix milliseconds).
    pub timestamp: i64,
    /// Ed25519 signature over BLAKE3(CBOR(preceding fields)).
    #[serde(with = "BigArray")]
    pub signature: [u8; 64],
}

impl SearchResponse {
    /// Compute the signing bytes.
    pub fn signing_bytes(&self) -> [u8; 32] {
        let signable = SearchResponseSignable {
            request_id: &self.request_id,
            status: &self.status,
            results: &self.results,
            peer_state_root: &self.peer_state_root,
            search_latency_ms: self.search_latency_ms,
            timestamp: self.timestamp,
        };
        let mut buf = Vec::new();
        ciborium::into_writer(&signable, &mut buf).expect("CBOR serialization cannot fail");
        *blake3::hash(&buf).as_bytes()
    }
}

#[derive(Serialize)]
struct SearchResponseSignable<'a> {
    request_id: &'a [u8; 16],
    status: &'a SearchStatus,
    results: &'a Vec<RemoteSearchResult>,
    peer_state_root: &'a [u8; 32],
    search_latency_ms: u16,
    timestamp: i64,
}

// ============================================================================
// Result merging types (Section 16)
// ============================================================================

/// Source of a merged search result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ResultSource {
    /// Result came from local graph only.
    Local,
    /// Result came from a remote peer.
    Remote { peer_node_id: [u8; 16] },
    /// Result found in both local and remote.
    Both { peer_node_id: [u8; 16] },
}

/// A merged search result combining local and remote sources.
#[derive(Debug, Clone)]
pub struct MergedSearchResult {
    pub chunk_id: [u8; 16],
    pub chunk_text: String,
    pub document_path: String,
    pub score: f64,
    pub source: ResultSource,
    pub merkle_proof: Option<Vec<[u8; 32]>>,
    pub peer_state_root: Option<[u8; 32]>,
    pub peer_signature: Option<[u8; 64]>,
}

// ============================================================================
// Circuit pool types (Section 15)
// ============================================================================

/// A warm circuit to a specific peer.
#[derive(Debug, Clone)]
pub struct WarmCircuit {
    /// The peer this circuit connects to.
    pub peer: PeerRegistration,
    /// When this circuit was established (Unix ms).
    pub created_at: i64,
    /// When this circuit was last used for a query (Unix ms).
    pub last_used: i64,
    /// When the last keepalive probe was sent (Unix ms).
    pub last_keepalive: i64,
}

impl WarmCircuit {
    /// Check if this circuit needs a keepalive probe.
    pub fn needs_keepalive(&self, now_ms: i64) -> bool {
        let elapsed = (now_ms - self.last_keepalive) / 1000;
        elapsed >= KEEPALIVE_INTERVAL_SECS as i64
    }

    /// Check if this circuit should be rotated.
    pub fn needs_rotation(&self, now_ms: i64) -> bool {
        let elapsed = (now_ms - self.created_at) / 1000;
        elapsed >= CIRCUIT_ROTATION_SECS as i64
    }
}

/// Peer selection score for circuit pool construction.
#[derive(Debug, Clone)]
pub struct PeerScore {
    pub peer: PeerRegistration,
    pub topic_overlap: f64,
    pub rating: f64,
    pub chunk_score: f64,
    pub composite: f64,
}

impl PeerScore {
    /// Compute composite score per CP-013 §15:
    /// 50% topic overlap + 40% rating + 10% chunk count (log2 scale).
    pub fn compute(
        peer: PeerRegistration,
        node_topics: &[String],
        peer_rating: Option<f64>,
    ) -> Self {
        let topic_overlap = if node_topics.is_empty() || peer.topics.is_empty() {
            0.5 // Neutral score when no topics specified
        } else {
            let matching = peer
                .topics
                .iter()
                .filter(|t| node_topics.contains(t))
                .count();
            // Per CP-013 §15: denominator is the node's interest list length
            matching as f64 / node_topics.len() as f64
        };

        let rating = peer_rating.unwrap_or(0.5); // Unrated peers get 0.5

        let chunk_score = if peer.capabilities.chunk_count > 0 {
            (peer.capabilities.chunk_count as f64).log2() / 20.0 // Normalize: log2(1M) ≈ 20
        } else {
            0.0
        };

        let composite = 0.5 * topic_overlap + 0.4 * rating + 0.1 * chunk_score.min(1.0);

        Self {
            peer,
            topic_overlap,
            rating,
            chunk_score,
            composite,
        }
    }
}

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

    fn test_registration() -> PeerRegistration {
        PeerRegistration {
            onion_address: "a".repeat(56),
            node_id: [1u8; 16],
            public_key: [2u8; 32],
            capabilities: PeerCapabilities::default(),
            topics: vec!["science".to_string(), "math".to_string()],
            embedding_model: [3u8; 32],
            timestamp: 1000000,
            signature: [0u8; 64],
        }
    }

    #[test]
    fn test_peer_capabilities_default() {
        let caps = PeerCapabilities::default();
        assert_eq!(caps.max_concurrent_queries, 4);
        assert_eq!(caps.max_results, 20);
        assert_eq!(caps.chunk_count, 0);
        assert!(caps.supports_proofs);
    }

    #[test]
    fn test_peer_registration_signing_bytes_deterministic() {
        let reg = test_registration();
        let bytes1 = reg.signing_bytes();
        let bytes2 = reg.signing_bytes();
        assert_eq!(bytes1, bytes2);
    }

    #[test]
    fn test_peer_registration_signing_bytes_change_on_mutation() {
        let mut reg1 = test_registration();
        let mut reg2 = test_registration();
        reg2.timestamp = 2000000;

        assert_ne!(reg1.signing_bytes(), reg2.signing_bytes());

        reg1.topics.push("physics".to_string());
        assert_ne!(reg1.signing_bytes(), test_registration().signing_bytes());
    }

    #[test]
    fn test_peer_registration_expiry() {
        let reg = PeerRegistration {
            timestamp: 1000,
            ..test_registration()
        };

        // 71 hours later: not expired
        let now_71h = 1000 + 71 * 3600 * 1000;
        assert!(!reg.is_expired(now_71h));

        // 73 hours later: expired
        let now_73h = 1000 + 73 * 3600 * 1000;
        assert!(reg.is_expired(now_73h));
    }

    #[test]
    fn test_search_request_signing_bytes_deterministic() {
        let req = SearchRequest {
            request_id: [1u8; 16],
            query_embedding: vec![100, -200, 300],
            query_text: Some("test query".to_string()),
            max_results: 10,
            include_proofs: true,
            model_hash: [5u8; 32],
            timestamp: 1000000,
            signature: [0u8; 64],
            public_key: [6u8; 32],
        };

        let bytes1 = req.signing_bytes();
        let bytes2 = req.signing_bytes();
        assert_eq!(bytes1, bytes2);
    }

    #[test]
    fn test_search_response_signing_bytes() {
        let resp = SearchResponse {
            request_id: [1u8; 16],
            status: SearchStatus::Ok,
            results: vec![RemoteSearchResult {
                chunk_id: [2u8; 16],
                chunk_text: "test chunk".to_string(),
                document_path: "doc.md".to_string(),
                score: 100,
                merkle_proof: None,
            }],
            peer_state_root: [3u8; 32],
            search_latency_ms: 150,
            timestamp: 1000000,
            signature: [0u8; 64],
        };

        let bytes = resp.signing_bytes();
        assert_eq!(bytes.len(), 32);
    }

    #[test]
    fn test_warm_circuit_keepalive() {
        let circuit = WarmCircuit {
            peer: test_registration(),
            created_at: 0,
            last_used: 0,
            last_keepalive: 0,
        };

        // Before interval: no keepalive needed
        let before = (KEEPALIVE_INTERVAL_SECS as i64 - 10) * 1000;
        assert!(!circuit.needs_keepalive(before));

        // After interval: keepalive needed
        let after = (KEEPALIVE_INTERVAL_SECS as i64 + 10) * 1000;
        assert!(circuit.needs_keepalive(after));
    }

    #[test]
    fn test_warm_circuit_rotation() {
        let circuit = WarmCircuit {
            peer: test_registration(),
            created_at: 0,
            last_used: 0,
            last_keepalive: 0,
        };

        // Before rotation window: no rotation
        let before = (CIRCUIT_ROTATION_SECS as i64 - 10) * 1000;
        assert!(!circuit.needs_rotation(before));

        // After rotation window: rotate
        let after = (CIRCUIT_ROTATION_SECS as i64 + 10) * 1000;
        assert!(circuit.needs_rotation(after));
    }

    #[test]
    fn test_peer_score_computation() {
        let reg = PeerRegistration {
            capabilities: PeerCapabilities {
                chunk_count: 100_000,
                ..PeerCapabilities::default()
            },
            ..test_registration()
        };

        let node_topics = vec![
            "science".to_string(),
            "math".to_string(),
            "physics".to_string(),
        ];
        let score = PeerScore::compute(reg, &node_topics, Some(0.8));

        // topic overlap: 2/3 = 0.667
        assert!(score.topic_overlap > 0.6 && score.topic_overlap < 0.7);
        assert!((score.rating - 0.8).abs() < 1e-6);
        assert!(score.chunk_score > 0.0);
        assert!(score.composite > 0.0 && score.composite < 1.0);
    }

    #[test]
    fn test_peer_score_unrated() {
        let reg = test_registration();
        let score = PeerScore::compute(reg, &[], None);

        assert!((score.rating - 0.5).abs() < 1e-6);
        assert!((score.topic_overlap - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_search_status_serialization() {
        let statuses = vec![
            SearchStatus::Ok,
            SearchStatus::ModelMismatch,
            SearchStatus::Overloaded,
            SearchStatus::InvalidRequest,
        ];

        for status in statuses {
            let mut buf = Vec::new();
            ciborium::into_writer(&status, &mut buf).unwrap();
            let decoded: SearchStatus = ciborium::from_reader(buf.as_slice()).unwrap();
            assert_eq!(decoded, status);
        }
    }

    #[test]
    fn test_cbor_roundtrip_search_request() {
        let req = SearchRequest {
            request_id: [7u8; 16],
            query_embedding: vec![1, 2, 3, -4, -5],
            query_text: Some("what is quantum computing?".to_string()),
            max_results: 10,
            include_proofs: true,
            model_hash: [9u8; 32],
            timestamp: 1234567890,
            signature: [10u8; 64],
            public_key: [11u8; 32],
        };

        let mut buf = Vec::new();
        ciborium::into_writer(&req, &mut buf).unwrap();
        let decoded: SearchRequest = ciborium::from_reader(buf.as_slice()).unwrap();

        assert_eq!(decoded.request_id, req.request_id);
        assert_eq!(decoded.query_embedding, req.query_embedding);
        assert_eq!(decoded.query_text, req.query_text);
        assert_eq!(decoded.max_results, req.max_results);
        assert_eq!(decoded.model_hash, req.model_hash);
        assert_eq!(decoded.public_key, req.public_key);
    }

    #[test]
    fn test_cbor_roundtrip_search_response() {
        let resp = SearchResponse {
            request_id: [1u8; 16],
            status: SearchStatus::Ok,
            results: vec![
                RemoteSearchResult {
                    chunk_id: [2u8; 16],
                    chunk_text: "result one".to_string(),
                    document_path: "docs/a.md".to_string(),
                    score: 500,
                    merkle_proof: Some(vec![[3u8; 32], [4u8; 32]]),
                },
                RemoteSearchResult {
                    chunk_id: [5u8; 16],
                    chunk_text: "result two".to_string(),
                    document_path: "docs/b.md".to_string(),
                    score: 300,
                    merkle_proof: None,
                },
            ],
            peer_state_root: [6u8; 32],
            search_latency_ms: 250,
            timestamp: 9876543210,
            signature: [7u8; 64],
        };

        let mut buf = Vec::new();
        ciborium::into_writer(&resp, &mut buf).unwrap();
        let decoded: SearchResponse = ciborium::from_reader(buf.as_slice()).unwrap();

        assert_eq!(decoded.request_id, resp.request_id);
        assert_eq!(decoded.status, SearchStatus::Ok);
        assert_eq!(decoded.results.len(), 2);
        assert_eq!(decoded.results[0].chunk_text, "result one");
        assert_eq!(decoded.results[1].score, 300);
        assert!(decoded.results[0].merkle_proof.is_some());
        assert!(decoded.results[1].merkle_proof.is_none());
    }

    #[test]
    fn test_cbor_roundtrip_peer_registration() {
        let reg = test_registration();

        let mut buf = Vec::new();
        ciborium::into_writer(&reg, &mut buf).unwrap();
        let decoded: PeerRegistration = ciborium::from_reader(buf.as_slice()).unwrap();

        assert_eq!(decoded.onion_address, reg.onion_address);
        assert_eq!(decoded.node_id, reg.node_id);
        assert_eq!(decoded.public_key, reg.public_key);
        assert_eq!(decoded.topics, reg.topics);
        assert_eq!(decoded.embedding_model, reg.embedding_model);
        assert_eq!(decoded.timestamp, reg.timestamp);
    }
}