cp_tor/

wire.rs

//! Length-prefixed wire format for Canon search protocol messages.
//!
//! All messages use a 4-byte big-endian u32 length prefix followed by
//! a CBOR-encoded payload. Maximum message size is 1MB.

use crate::error::{Result, TorError};
use crate::types::MAX_MESSAGE_SIZE;
use serde::{de::DeserializeOwned, Serialize};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Encode a message to length-prefixed CBOR wire format.
pub fn encode<T: Serialize>(msg: &T) -> Result<Vec<u8>> {
    let mut cbor_buf = Vec::new();
    ciborium::into_writer(msg, &mut cbor_buf)
        .map_err(|e| TorError::Serialization(e.to_string()))?;

    if cbor_buf.len() > MAX_MESSAGE_SIZE as usize {
        return Err(TorError::MessageTooLarge {
            size: cbor_buf.len(),
            max: MAX_MESSAGE_SIZE as usize,
        });
    }

    let len = cbor_buf.len() as u32;
    let mut wire = Vec::with_capacity(4 + cbor_buf.len());
    wire.extend_from_slice(&len.to_be_bytes());
    wire.extend_from_slice(&cbor_buf);
    Ok(wire)
}

/// Decode a message from length-prefixed CBOR wire format.
pub fn decode<T: DeserializeOwned>(data: &[u8]) -> Result<T> {
    if data.len() < 4 {
        return Err(TorError::Serialization(
            "Message too short for length prefix".to_string(),
        ));
    }

    let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;

    if len > MAX_MESSAGE_SIZE as usize {
        return Err(TorError::MessageTooLarge {
            size: len,
            max: MAX_MESSAGE_SIZE as usize,
        });
    }

    if data.len() < 4 + len {
        return Err(TorError::Serialization(format!(
            "Incomplete message: expected {} bytes, got {}",
            4 + len,
            data.len()
        )));
    }

    ciborium::from_reader(&data[4..4 + len])
        .map_err(|e| TorError::Serialization(format!("CBOR decode failed: {e}")))
}

/// Write a length-prefixed CBOR message to an async writer.
pub async fn write_message<W, T>(writer: &mut W, msg: &T) -> Result<()>
where
    W: AsyncWrite + Unpin,
    T: Serialize,
{
    let wire = encode(msg)?;
    writer.write_all(&wire).await.map_err(TorError::Io)?;
    writer.flush().await.map_err(TorError::Io)?;
    Ok(())
}

/// Read a length-prefixed CBOR message from an async reader.
pub async fn read_message<R, T>(reader: &mut R) -> Result<T>
where
    R: AsyncRead + Unpin,
    T: DeserializeOwned,
{
    // Read 4-byte length prefix
    let mut len_buf = [0u8; 4];
    reader
        .read_exact(&mut len_buf)
        .await
        .map_err(TorError::Io)?;
    let len = u32::from_be_bytes(len_buf) as usize;

    // Zero-length message is a keepalive probe
    if len == 0 {
        return Err(TorError::Keepalive);
    }

    if len > MAX_MESSAGE_SIZE as usize {
        return Err(TorError::MessageTooLarge {
            size: len,
            max: MAX_MESSAGE_SIZE as usize,
        });
    }

    // Read payload
    let mut payload = vec![0u8; len];
    reader
        .read_exact(&mut payload)
        .await
        .map_err(TorError::Io)?;

    ciborium::from_reader(payload.as_slice())
        .map_err(|e| TorError::Serialization(format!("CBOR decode failed: {e}")))
}

/// Send a keepalive probe (zero-length message).
pub async fn write_keepalive<W: AsyncWrite + Unpin>(writer: &mut W) -> Result<()> {
    writer
        .write_all(&0u32.to_be_bytes())
        .await
        .map_err(TorError::Io)?;
    writer.flush().await.map_err(TorError::Io)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{RemoteSearchResult, SearchRequest, SearchResponse, SearchStatus};

    #[test]
    fn test_encode_decode_roundtrip() {
        let msg = vec![1u32, 2, 3, 4, 5];
        let wire = encode(&msg).unwrap();

        // First 4 bytes are length prefix
        let len = u32::from_be_bytes([wire[0], wire[1], wire[2], wire[3]]);
        assert_eq!(len as usize, wire.len() - 4);

        let decoded: Vec<u32> = decode(&wire).unwrap();
        assert_eq!(decoded, msg);
    }

    #[test]
    fn test_encode_search_request() {
        let req = SearchRequest {
            request_id: [1u8; 16],
            query_embedding: vec![100, -200, 300],
            query_text: None,
            max_results: 10,
            include_proofs: false,
            model_hash: [2u8; 32],
            timestamp: 1000,
            signature: [3u8; 64],
            public_key: [4u8; 32],
        };

        let wire = encode(&req).unwrap();
        let decoded: SearchRequest = decode(&wire).unwrap();
        assert_eq!(decoded.request_id, req.request_id);
        assert_eq!(decoded.query_embedding, req.query_embedding);
    }

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

        let wire = encode(&resp).unwrap();
        let decoded: SearchResponse = decode(&wire).unwrap();
        assert_eq!(decoded.results.len(), 1);
        assert_eq!(decoded.results[0].chunk_text, "hello");
    }

    #[test]
    fn test_message_too_large() {
        // Create a message larger than 1MB
        let big = vec![0u8; MAX_MESSAGE_SIZE as usize + 100];
        let result = encode(&big);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_truncated() {
        let msg = vec![1u32, 2, 3];
        let wire = encode(&msg).unwrap();

        // Truncate the payload
        let result: std::result::Result<Vec<u32>, _> = decode(&wire[..wire.len() - 2]);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_empty() {
        let result: std::result::Result<Vec<u32>, _> = decode(&[]);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_async_write_read_roundtrip() {
        let req = SearchRequest {
            request_id: [5u8; 16],
            query_embedding: vec![1, 2, 3, 4, 5],
            query_text: Some("test".to_string()),
            max_results: 5,
            include_proofs: true,
            model_hash: [6u8; 32],
            timestamp: 3000,
            signature: [7u8; 64],
            public_key: [8u8; 32],
        };

        // Write to buffer
        let mut buf = Vec::new();
        write_message(&mut buf, &req).await.unwrap();

        // Read back
        let mut cursor = &buf[..];
        let decoded: SearchRequest = read_message(&mut cursor).await.unwrap();
        assert_eq!(decoded.request_id, req.request_id);
        assert_eq!(decoded.query_text, Some("test".to_string()));
    }

    #[tokio::test]
    async fn test_keepalive_roundtrip() {
        let mut buf = Vec::new();
        write_keepalive(&mut buf).await.unwrap();

        assert_eq!(buf.len(), 4);
        assert_eq!(buf, vec![0, 0, 0, 0]);

        // Reading a keepalive should return Keepalive error
        let mut cursor = &buf[..];
        let result: std::result::Result<SearchRequest, _> = read_message(&mut cursor).await;
        assert!(matches!(result, Err(TorError::Keepalive)));
    }

    #[tokio::test]
    async fn test_multiple_messages_on_stream() {
        let msg1 = "hello".to_string();
        let msg2 = "world".to_string();
        let msg3 = vec![1u32, 2, 3];

        let mut buf = Vec::new();
        write_message(&mut buf, &msg1).await.unwrap();
        write_message(&mut buf, &msg2).await.unwrap();
        write_message(&mut buf, &msg3).await.unwrap();

        let mut cursor = &buf[..];
        let d1: String = read_message(&mut cursor).await.unwrap();
        let d2: String = read_message(&mut cursor).await.unwrap();
        let d3: Vec<u32> = read_message(&mut cursor).await.unwrap();

        assert_eq!(d1, "hello");
        assert_eq!(d2, "world");
        assert_eq!(d3, vec![1, 2, 3]);
    }
}