quicproquo/crates/quicprochat-p2p/src/envelope.rs

//! Store-and-forward message envelope for mesh routing.
//!
//! A [`MeshEnvelope`] wraps an encrypted payload with routing metadata
//! (sender/recipient keys, TTL, hop count) and an Ed25519 signature for
//! integrity. Envelopes are deduplicated by a SHA-256 content ID.

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::time::{SystemTime, UNIX_EPOCH};

use crate::identity::MeshIdentity;

/// Default maximum hops for mesh forwarding.
const DEFAULT_MAX_HOPS: u8 = 5;

/// A signed, routable message envelope for mesh store-and-forward.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MeshEnvelope {
    /// SHA-256 content ID (for deduplication).
    pub id: [u8; 32],
    /// 32-byte Ed25519 public key of the sender.
    pub sender_key: Vec<u8>,
    /// 32-byte Ed25519 public key of the recipient (empty for broadcast).
    pub recipient_key: Vec<u8>,
    /// Encrypted message body (opaque to the mesh layer).
    pub payload: Vec<u8>,
    /// Time-to-live in seconds from `timestamp`.
    pub ttl_secs: u32,
    /// Current hop count (incremented on each forward).
    pub hop_count: u8,
    /// Maximum allowed hops before the envelope is dropped.
    pub max_hops: u8,
    /// Unix timestamp (seconds) of creation.
    pub timestamp: u64,
    /// Ed25519 signature over all fields except `signature` itself.
    pub signature: Vec<u8>,
}

impl MeshEnvelope {
    /// Create and sign a new mesh envelope.
    pub fn new(
        identity: &MeshIdentity,
        recipient_key: &[u8],
        payload: Vec<u8>,
        ttl_secs: u32,
        max_hops: u8,
    ) -> Self {
        let sender_key = identity.public_key().to_vec();
        let recipient_key = recipient_key.to_vec();
        let hop_count = 0u8;
        let max_hops = if max_hops == 0 {
            DEFAULT_MAX_HOPS
        } else {
            max_hops
        };
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let id = Self::compute_id(
            &sender_key,
            &recipient_key,
            &payload,
            ttl_secs,
            max_hops,
            timestamp,
        );

        let signable = Self::signable_bytes(&id, &sender_key, &recipient_key, &payload, ttl_secs, max_hops, timestamp);
        let signature = identity.sign(&signable).to_vec();

        Self {
            id,
            sender_key,
            recipient_key,
            payload,
            ttl_secs,
            hop_count,
            max_hops,
            timestamp,
            signature,
        }
    }

    /// Compute the content ID from the immutable envelope fields.
    pub fn compute_id(
        sender_key: &[u8],
        recipient_key: &[u8],
        payload: &[u8],
        ttl_secs: u32,
        max_hops: u8,
        timestamp: u64,
    ) -> [u8; 32] {
        let mut hasher = Sha256::new();
        hasher.update(sender_key);
        hasher.update(recipient_key);
        hasher.update(payload);
        hasher.update(ttl_secs.to_le_bytes());
        hasher.update([max_hops]);
        hasher.update(timestamp.to_le_bytes());
        hasher.finalize().into()
    }

    /// Assemble the byte string that is signed / verified.
    ///
    /// `hop_count` is intentionally excluded: forwarding nodes increment it
    /// without re-signing, so including it would invalidate the sender's
    /// original signature on every hop.
    fn signable_bytes(
        id: &[u8; 32],
        sender_key: &[u8],
        recipient_key: &[u8],
        payload: &[u8],
        ttl_secs: u32,
        max_hops: u8,
        timestamp: u64,
    ) -> Vec<u8> {
        let mut buf = Vec::with_capacity(32 + sender_key.len() + recipient_key.len() + payload.len() + 13);
        buf.extend_from_slice(id);
        buf.extend_from_slice(sender_key);
        buf.extend_from_slice(recipient_key);
        buf.extend_from_slice(payload);
        buf.extend_from_slice(&ttl_secs.to_le_bytes());
        buf.push(max_hops);
        buf.extend_from_slice(&timestamp.to_le_bytes());
        buf
    }

    /// Verify the envelope's Ed25519 signature.
    ///
    /// Returns `true` if the signature is valid and the sender key is a valid
    /// Ed25519 public key.
    pub fn verify(&self) -> bool {
        let sender_key: [u8; 32] = match self.sender_key.as_slice().try_into() {
            Ok(k) => k,
            Err(_) => return false,
        };
        let sig: [u8; 64] = match self.signature.as_slice().try_into() {
            Ok(s) => s,
            Err(_) => return false,
        };
        let signable = Self::signable_bytes(
            &self.id,
            &self.sender_key,
            &self.recipient_key,
            &self.payload,
            self.ttl_secs,
            self.max_hops,
            self.timestamp,
        );
        quicprochat_core::IdentityKeypair::verify_raw(&sender_key, &signable, &sig).is_ok()
    }

    /// Check whether this envelope has expired (TTL elapsed since timestamp).
    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        now.saturating_sub(self.timestamp) > self.ttl_secs as u64
    }

    /// Whether this envelope can be forwarded (not expired and under hop limit).
    pub fn can_forward(&self) -> bool {
        self.hop_count < self.max_hops && !self.is_expired()
    }

    /// Create a forwarded copy with `hop_count` incremented by one.
    ///
    /// The signature remains the sender's original signature — forwarding
    /// nodes do not re-sign.
    pub fn forwarded(&self) -> Self {
        let mut copy = self.clone();
        copy.hop_count = copy.hop_count.saturating_add(1);
        copy
    }

    /// Serialize to compact CBOR binary format (for wire transmission).
    pub fn to_wire(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        ciborium::into_writer(self, &mut buf).expect("CBOR serialization should not fail");
        buf
    }

    /// Deserialize from CBOR binary format.
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        let env: Self = ciborium::from_reader(bytes)?;
        Ok(env)
    }

    /// Deserialize from wire format, trying CBOR first then JSON fallback.
    pub fn from_wire_or_json(bytes: &[u8]) -> anyhow::Result<Self> {
        Self::from_wire(bytes).or_else(|_| Self::from_bytes(bytes))
    }

    /// Serialize to bytes (JSON). Kept for backward compatibility and debugging.
    pub fn to_bytes(&self) -> Vec<u8> {
        // serde_json::to_vec should not fail on a well-formed envelope.
        serde_json::to_vec(self).expect("envelope serialization should not fail")
    }

    /// Deserialize from bytes (JSON). Kept for backward compatibility and debugging.
    pub fn from_bytes(bytes: &[u8]) -> anyhow::Result<Self> {
        let env: Self = serde_json::from_slice(bytes)?;
        Ok(env)
    }
}

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

    fn test_identity() -> MeshIdentity {
        MeshIdentity::generate()
    }

    #[test]
    fn create_and_verify() {
        let id = test_identity();
        let recipient = [0xBBu8; 32];
        let env = MeshEnvelope::new(&id, &recipient, b"hello mesh".to_vec(), 3600, 5);

        assert!(env.verify(), "freshly created envelope must verify");
        assert!(!env.is_expired());
        assert!(env.can_forward());
        assert_eq!(env.hop_count, 0);
        assert_eq!(env.sender_key, id.public_key().to_vec());
        assert_eq!(env.recipient_key, recipient.to_vec());
    }

    #[test]
    fn tampered_payload_fails_verify() {
        let id = test_identity();
        let mut env = MeshEnvelope::new(&id, &[0xCC; 32], b"original".to_vec(), 60, 3);
        env.payload = b"tampered".to_vec();
        assert!(!env.verify(), "tampered envelope must fail verification");
    }

    #[test]
    fn expired_envelope() {
        let id = test_identity();
        let mut env = MeshEnvelope::new(&id, &[0xDD; 32], b"old".to_vec(), 0, 5);
        // Set timestamp to the past so TTL of 0 guarantees expiry.
        env.timestamp = 0;
        assert!(env.is_expired());
        assert!(!env.can_forward());
    }

    #[test]
    fn forward_increments_hop() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[0xEE; 32], b"hop".to_vec(), 3600, 2);
        assert_eq!(env.hop_count, 0);

        let fwd1 = env.forwarded();
        assert_eq!(fwd1.hop_count, 1);
        assert!(fwd1.can_forward());

        let fwd2 = fwd1.forwarded();
        assert_eq!(fwd2.hop_count, 2);
        assert!(!fwd2.can_forward()); // hop_count == max_hops
    }

    #[test]
    fn forwarded_envelope_still_verifies() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[0xAA; 32], b"fwd-verify".to_vec(), 3600, 5);
        assert!(env.verify(), "original must verify");

        let fwd = env.forwarded();
        assert_eq!(fwd.hop_count, 1);
        assert!(fwd.verify(), "forwarded envelope must still verify (hop_count excluded from signature)");

        let fwd2 = fwd.forwarded();
        assert!(fwd2.verify(), "double-forwarded must still verify");
    }

    #[test]
    fn verify_with_wrong_key_fails() {
        let id = test_identity();
        let mut env = MeshEnvelope::new(&id, &[0xBB; 32], b"wrong-key".to_vec(), 3600, 5);
        // Replace sender_key with a different key
        let other = test_identity();
        env.sender_key = other.public_key().to_vec();
        assert!(!env.verify(), "wrong sender key must fail verification");
    }

    #[test]
    fn serialization_roundtrip() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[0xFF; 32], b"roundtrip".to_vec(), 300, 4);
        let bytes = env.to_bytes();
        let restored = MeshEnvelope::from_bytes(&bytes).expect("deserialize");
        assert_eq!(env.id, restored.id);
        assert_eq!(env.payload, restored.payload);
        assert!(restored.verify());
    }

    #[test]
    fn default_max_hops_when_zero() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[0x11; 32], b"defaults".to_vec(), 60, 0);
        assert_eq!(env.max_hops, 5); // DEFAULT_MAX_HOPS
    }

    #[test]
    fn broadcast_envelope_empty_recipient() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[], b"broadcast".to_vec(), 60, 3);
        assert!(env.recipient_key.is_empty());
        assert!(env.verify());
    }

    #[test]
    fn cbor_roundtrip() {
        let id = test_identity();
        let recipient = [0xABu8; 32];
        let env = MeshEnvelope::new(&id, &recipient, b"cbor roundtrip".to_vec(), 3600, 5);

        let wire = env.to_wire();
        let restored = MeshEnvelope::from_wire(&wire).expect("CBOR deserialize");

        assert_eq!(env.id, restored.id);
        assert_eq!(env.sender_key, restored.sender_key);
        assert_eq!(env.recipient_key, restored.recipient_key);
        assert_eq!(env.payload, restored.payload);
        assert_eq!(env.ttl_secs, restored.ttl_secs);
        assert_eq!(env.hop_count, restored.hop_count);
        assert_eq!(env.max_hops, restored.max_hops);
        assert_eq!(env.timestamp, restored.timestamp);
        assert_eq!(env.signature, restored.signature);
        assert!(restored.verify());
    }

    #[test]
    fn cbor_smaller_than_json() {
        let id = test_identity();
        let recipient = [0xCCu8; 32];
        let payload = b"a typical chat message for size comparison testing".to_vec();
        let env = MeshEnvelope::new(&id, &recipient, payload, 3600, 5);

        let wire_len = env.to_wire().len();
        let json_len = env.to_bytes().len();

        println!("CBOR wire size: {wire_len} bytes");
        println!("JSON size:      {json_len} bytes");
        println!("Ratio:          {:.1}x smaller", json_len as f64 / wire_len as f64);

        assert!(
            json_len * 2 > wire_len * 3,
            "CBOR ({wire_len}B) should be materially smaller than JSON ({json_len}B)"
        );
    }

    #[test]
    fn cbor_backward_compat() {
        let id = test_identity();
        let env = MeshEnvelope::new(&id, &[0xDD; 32], b"json compat".to_vec(), 60, 3);

        // Serialize as JSON (old format).
        let json_bytes = env.to_bytes();

        // from_wire_or_json should fall back to JSON parsing.
        let restored = MeshEnvelope::from_wire_or_json(&json_bytes)
            .expect("from_wire_or_json should handle JSON");
        assert_eq!(env.id, restored.id);
        assert_eq!(env.payload, restored.payload);
        assert!(restored.verify());
    }

    #[test]
    fn cbor_from_wire_rejects_garbage() {
        let garbage = [0xFF, 0xFE, 0x00, 0x42, 0x99, 0x01, 0x02, 0x03];
        let result = MeshEnvelope::from_wire(&garbage);
        assert!(result.is_err(), "garbage input must return Err, not panic");
    }
}