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

//! Compact mesh envelope using truncated 16-byte addresses.
//!
//! [`MeshEnvelopeV2`] is a bandwidth-optimized envelope format for constrained
//! links (LoRa, serial). It uses [`MeshAddress`] (16 bytes) instead of full
//! 32-byte public keys, saving 32 bytes per envelope.
//!
//! Full public keys are exchanged during the announce phase and cached in the
//! routing table. The envelope only needs addresses for routing.

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

use crate::address::MeshAddress;
use crate::identity::MeshIdentity;

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

/// Version byte for envelope format detection.
const ENVELOPE_V2_VERSION: u8 = 0x02;

/// Priority levels for mesh routing.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[repr(u8)]
pub enum Priority {
    /// Lowest priority (announce, telemetry).
    Low = 0,
    /// Normal priority (regular messages).
    Normal = 1,
    /// High priority (important messages).
    High = 2,
    /// Emergency priority (always forwarded first).
    Emergency = 3,
}

impl Default for Priority {
    fn default() -> Self {
        Self::Normal
    }
}

impl From<u8> for Priority {
    fn from(v: u8) -> Self {
        match v {
            0 => Self::Low,
            1 => Self::Normal,
            2 => Self::High,
            3 => Self::Emergency,
            _ => Self::Normal,
        }
    }
}

/// Compact mesh envelope with 16-byte truncated addresses.
///
/// # Wire overhead
///
/// - Version: 1 byte
/// - Flags: 1 byte (priority: 2 bits, reserved: 6 bits)
/// - ID: 16 bytes (truncated from 32)
/// - Sender: 16 bytes
/// - Recipient: 16 bytes (or 0 for broadcast)
/// - TTL: 2 bytes (u16, max ~18 hours)
/// - Hop count: 1 byte
/// - Max hops: 1 byte
/// - Timestamp: 4 bytes (u32, seconds since epoch mod 2^32)
/// - Signature: 64 bytes
/// - Payload: variable
///
/// **Total fixed overhead: ~122 bytes** (vs ~174 for V1 with full keys)
/// Savings: ~52 bytes per envelope
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MeshEnvelopeV2 {
    /// Format version (0x02 for V2).
    pub version: u8,
    /// Flags byte: bits 0-1 = priority, bits 2-7 reserved.
    pub flags: u8,
    /// 16-byte truncated content ID (for deduplication).
    pub id: [u8; 16],
    /// 16-byte truncated sender address.
    pub sender_addr: MeshAddress,
    /// 16-byte truncated recipient address (BROADCAST for all).
    pub recipient_addr: MeshAddress,
    /// Encrypted payload (opaque to mesh layer).
    pub payload: Vec<u8>,
    /// Time-to-live in seconds (u16, max 65535 = ~18 hours).
    pub ttl_secs: u16,
    /// Current hop count.
    pub hop_count: u8,
    /// Maximum hops before drop.
    pub max_hops: u8,
    /// Unix timestamp (seconds, truncated to u32).
    pub timestamp: u32,
    /// Ed25519 signature (64 bytes, stored as Vec for serde compatibility).
    pub signature: Vec<u8>,
}

impl MeshEnvelopeV2 {
    /// Create and sign a new compact mesh envelope.
    pub fn new(
        identity: &MeshIdentity,
        recipient_addr: MeshAddress,
        payload: Vec<u8>,
        ttl_secs: u16,
        max_hops: u8,
        priority: Priority,
    ) -> Self {
        let sender_addr = MeshAddress::from_public_key(&identity.public_key());
        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() & 0xFFFF_FFFF) as u32;

        let id = Self::compute_id(
            &sender_addr,
            &recipient_addr,
            &payload,
            ttl_secs,
            max_hops,
            timestamp,
        );

        let flags = (priority as u8) & 0x03;

        let mut envelope = Self {
            version: ENVELOPE_V2_VERSION,
            flags,
            id,
            sender_addr,
            recipient_addr,
            payload,
            ttl_secs,
            hop_count,
            max_hops,
            timestamp,
            signature: Vec::new(),
        };

        let signable = envelope.signable_bytes();
        let sig = identity.sign(&signable);
        envelope.signature = sig.to_vec();

        envelope
    }

    /// Create for broadcast (recipient = all zeros).
    pub fn broadcast(
        identity: &MeshIdentity,
        payload: Vec<u8>,
        ttl_secs: u16,
        max_hops: u8,
        priority: Priority,
    ) -> Self {
        Self::new(identity, MeshAddress::BROADCAST, payload, ttl_secs, max_hops, priority)
    }

    /// Compute the 16-byte truncated content ID.
    fn compute_id(
        sender_addr: &MeshAddress,
        recipient_addr: &MeshAddress,
        payload: &[u8],
        ttl_secs: u16,
        max_hops: u8,
        timestamp: u32,
    ) -> [u8; 16] {
        let mut hasher = Sha256::new();
        hasher.update(sender_addr.as_bytes());
        hasher.update(recipient_addr.as_bytes());
        hasher.update(payload);
        hasher.update(ttl_secs.to_le_bytes());
        hasher.update([max_hops]);
        hasher.update(timestamp.to_le_bytes());
        let hash = hasher.finalize();
        let mut id = [0u8; 16];
        id.copy_from_slice(&hash[..16]);
        id
    }

    /// Bytes to sign/verify (excludes signature and hop_count).
    fn signable_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(64 + self.payload.len());
        buf.push(self.version);
        buf.push(self.flags);
        buf.extend_from_slice(&self.id);
        buf.extend_from_slice(self.sender_addr.as_bytes());
        buf.extend_from_slice(self.recipient_addr.as_bytes());
        buf.extend_from_slice(&self.payload);
        buf.extend_from_slice(&self.ttl_secs.to_le_bytes());
        buf.push(self.max_hops);
        buf.extend_from_slice(&self.timestamp.to_le_bytes());
        buf
    }

    /// Verify the signature using the sender's full public key.
    ///
    /// The caller must have the sender's full key (from announce/routing table).
    pub fn verify_with_key(&self, sender_public_key: &[u8; 32]) -> bool {
        // First check that the address matches the key
        if !self.sender_addr.matches_key(sender_public_key) {
            return false;
        }
        // Signature must be exactly 64 bytes
        let sig: [u8; 64] = match self.signature.as_slice().try_into() {
            Ok(s) => s,
            Err(_) => return false,
        };
        let signable = self.signable_bytes();
        quicprochat_core::IdentityKeypair::verify_raw(sender_public_key, &signable, &sig).is_ok()
    }

    /// Get the priority level.
    pub fn priority(&self) -> Priority {
        Priority::from(self.flags & 0x03)
    }

    /// Check if broadcast (recipient is all zeros).
    pub fn is_broadcast(&self) -> bool {
        self.recipient_addr.is_broadcast()
    }

    /// Check if expired.
    pub fn is_expired(&self) -> bool {
        let now = (SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() & 0xFFFF_FFFF) as u32;
        // Handle u32 wraparound (every ~136 years)
        let elapsed = now.wrapping_sub(self.timestamp);
        elapsed > self.ttl_secs as u32
    }

    /// Can this envelope be forwarded?
    pub fn can_forward(&self) -> bool {
        self.hop_count < self.max_hops && !self.is_expired()
    }

    /// Create a forwarded copy with hop_count incremented.
    pub fn forwarded(&self) -> Self {
        let mut copy = self.clone();
        copy.hop_count = copy.hop_count.saturating_add(1);
        copy
    }

    /// Serialize to compact CBOR.
    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.
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        let env: Self = ciborium::from_reader(bytes)?;
        if env.version != ENVELOPE_V2_VERSION {
            anyhow::bail!("unexpected envelope version: {}", env.version);
        }
        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_key = [0xBBu8; 32];
        let recipient_addr = MeshAddress::from_public_key(&recipient_key);

        let env = MeshEnvelopeV2::new(
            &id,
            recipient_addr,
            b"hello compact".to_vec(),
            3600,
            5,
            Priority::Normal,
        );

        assert_eq!(env.version, ENVELOPE_V2_VERSION);
        assert_eq!(env.hop_count, 0);
        assert!(env.verify_with_key(&id.public_key()));
        assert!(!env.is_expired());
        assert!(env.can_forward());
    }

    #[test]
    fn broadcast_envelope() {
        let id = test_identity();
        let env = MeshEnvelopeV2::broadcast(
            &id,
            b"announcement".to_vec(),
            300,
            8,
            Priority::Low,
        );

        assert!(env.is_broadcast());
        assert_eq!(env.priority(), Priority::Low);
        assert!(env.verify_with_key(&id.public_key()));
    }

    #[test]
    fn forwarded_still_verifies() {
        let id = test_identity();
        let env = MeshEnvelopeV2::new(
            &id,
            MeshAddress::from_bytes([0xCC; 16]),
            b"forward me".to_vec(),
            3600,
            5,
            Priority::High,
        );

        let fwd = env.forwarded();
        assert_eq!(fwd.hop_count, 1);
        assert!(fwd.verify_with_key(&id.public_key()));

        let fwd2 = fwd.forwarded();
        assert_eq!(fwd2.hop_count, 2);
        assert!(fwd2.verify_with_key(&id.public_key()));
    }

    #[test]
    fn cbor_roundtrip() {
        let id = test_identity();
        let env = MeshEnvelopeV2::new(
            &id,
            MeshAddress::from_bytes([0xDD; 16]),
            b"roundtrip test".to_vec(),
            1800,
            4,
            Priority::Emergency,
        );

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

        assert_eq!(env.id, restored.id);
        assert_eq!(env.sender_addr, restored.sender_addr);
        assert_eq!(env.recipient_addr, restored.recipient_addr);
        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_eq!(env.priority(), Priority::Emergency);
    }

    #[test]
    fn measure_v2_overhead() {
        let id = test_identity();
        let recipient_addr = MeshAddress::from_bytes([0xEE; 16]);

        println!("=== MeshEnvelopeV2 Wire Overhead (CBOR) ===");

        // Empty payload
        let env_empty = MeshEnvelopeV2::new(&id, recipient_addr, vec![], 3600, 5, Priority::Normal);
        let wire_empty = env_empty.to_wire();
        println!("Payload   0B: wire {} bytes (overhead: {} bytes)", wire_empty.len(), wire_empty.len());
        let v2_overhead = wire_empty.len();

        // Compare to V1
        let v1_env = crate::envelope::MeshEnvelope::new(
            &id,
            &[0xEE; 32],
            vec![],
            3600,
            5,
        );
        let v1_wire = v1_env.to_wire();
        println!("V1 empty:     {} bytes", v1_wire.len());
        println!("V2 savings:   {} bytes ({:.1}%)", 
            v1_wire.len() - v2_overhead,
            ((v1_wire.len() - v2_overhead) as f64 / v1_wire.len() as f64) * 100.0);

        // 10-byte payload
        let env_10 = MeshEnvelopeV2::new(&id, recipient_addr, b"hello mesh".to_vec(), 3600, 5, Priority::Normal);
        let wire_10 = env_10.to_wire();
        println!("Payload  10B: wire {} bytes", wire_10.len());

        // 100-byte payload
        let env_100 = MeshEnvelopeV2::new(&id, recipient_addr, vec![0x42; 100], 3600, 5, Priority::Normal);
        let wire_100 = env_100.to_wire();
        println!("Payload 100B: wire {} bytes", wire_100.len());

        // V2 should be smaller than V1 due to truncated addresses
        // With CBOR field names, actual overhead is higher than theoretical minimum
        // (~336 bytes for V2 vs ~410 for V1 = ~18% savings)
        assert!(v2_overhead < v1_wire.len(), "V2 should be smaller than V1");
        let savings_pct = ((v1_wire.len() - v2_overhead) as f64 / v1_wire.len() as f64) * 100.0;
        assert!(savings_pct > 10.0, "V2 should save at least 10% vs V1");
        println!("Actual V2 savings: {:.1}%", savings_pct);
    }

    #[test]
    fn wrong_key_fails_verification() {
        let id = test_identity();
        let env = MeshEnvelopeV2::new(
            &id,
            MeshAddress::from_bytes([0xFF; 16]),
            b"verify me".to_vec(),
            3600,
            5,
            Priority::Normal,
        );

        // Wrong key should fail
        let wrong_key = [0x42u8; 32];
        assert!(!env.verify_with_key(&wrong_key));

        // Correct key should pass
        assert!(env.verify_with_key(&id.public_key()));
    }

    #[test]
    fn priority_levels() {
        let id = test_identity();

        for prio in [Priority::Low, Priority::Normal, Priority::High, Priority::Emergency] {
            let env = MeshEnvelopeV2::new(
                &id,
                MeshAddress::BROADCAST,
                b"prio test".to_vec(),
                60,
                3,
                prio,
            );
            assert_eq!(env.priority(), prio);
        }
    }
}