fips/src/proto/link.rs

//! Link-layer message types: the shared frame catalog and session datagram.

use crate::NodeAddr;
use crate::proto::Error;
use core::fmt;

// ============================================================================
// Link-Layer Message Types
// ============================================================================

/// Link-layer message type identifiers.
///
/// These messages are exchanged between directly connected peers over
/// Noise-encrypted links. All payloads are encrypted with session keys
/// established during the Noise XX handshake.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum LinkMessageType {
    // Forwarding (0x00-0x0F)
    /// Encapsulated session-layer datagram for forwarding.
    /// Payload is opaque to intermediate nodes (end-to-end encrypted).
    SessionDatagram = 0x00,

    // MMP reports (0x01-0x02) — content defined in TASK-2026-0006
    /// Sender-side MMP report (stub).
    SenderReport = 0x01,
    /// Receiver-side MMP report (stub).
    ReceiverReport = 0x02,

    // Tree protocol (0x10-0x1F)
    /// Spanning tree state announcement.
    TreeAnnounce = 0x10,

    // Bloom filter (0x20-0x2F)
    /// Bloom filter reachability update (full or delta).
    FilterAnnounce = 0x20,

    // Discovery (0x30-0x3F)
    /// Request to discover a node's coordinates.
    LookupRequest = 0x30,
    /// Response with target's coordinates.
    LookupResponse = 0x31,

    // Link Control (0x50-0x5F)
    /// Orderly disconnect notification before link closure.
    Disconnect = 0x50,
    /// Periodic heartbeat for link liveness detection.
    /// No payload — the msg_type byte alone is sufficient.
    Heartbeat = 0x51,
}

impl LinkMessageType {
    /// Try to convert from a byte.
    pub fn from_byte(b: u8) -> Option<Self> {
        match b {
            0x00 => Some(LinkMessageType::SessionDatagram),
            0x01 => Some(LinkMessageType::SenderReport),
            0x02 => Some(LinkMessageType::ReceiverReport),
            0x10 => Some(LinkMessageType::TreeAnnounce),
            0x20 => Some(LinkMessageType::FilterAnnounce),
            0x30 => Some(LinkMessageType::LookupRequest),
            0x31 => Some(LinkMessageType::LookupResponse),
            0x50 => Some(LinkMessageType::Disconnect),
            0x51 => Some(LinkMessageType::Heartbeat),
            _ => None,
        }
    }

    /// Convert to a byte.
    pub fn to_byte(self) -> u8 {
        self as u8
    }
}

impl fmt::Display for LinkMessageType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            LinkMessageType::SessionDatagram => "SessionDatagram",
            LinkMessageType::SenderReport => "SenderReport",
            LinkMessageType::ReceiverReport => "ReceiverReport",
            LinkMessageType::TreeAnnounce => "TreeAnnounce",
            LinkMessageType::FilterAnnounce => "FilterAnnounce",
            LinkMessageType::LookupRequest => "LookupRequest",
            LinkMessageType::LookupResponse => "LookupResponse",
            LinkMessageType::Disconnect => "Disconnect",
            LinkMessageType::Heartbeat => "Heartbeat",
        };
        write!(f, "{}", name)
    }
}

// ============================================================================
// Session Datagram (Link-Layer Encapsulation)
// ============================================================================

/// Encapsulated session-layer datagram for multi-hop forwarding.
///
/// This is a link-layer message (type 0x00) that carries session-layer
/// payloads through the mesh. The envelope provides source and destination
/// addressing that transit routers use for forwarding decisions and error
/// routing.
///
/// ## Wire Format (36-byte fixed header)
///
/// | Offset | Field     | Size     | Description                         |
/// |--------|-----------|----------|-------------------------------------|
/// | 0      | msg_type  | 1 byte   | 0x00                                |
/// | 1      | ttl       | 1 byte   | Decremented each hop                |
/// | 2      | path_mtu  | 2 bytes  | Path MTU (LE), min'd at each hop    |
/// | 4      | src_addr  | 16 bytes | Source node_addr                    |
/// | 20     | dest_addr | 16 bytes | Destination node_addr               |
/// | 36     | payload   | variable | Session-layer message               |
///
/// The payload is either end-to-end encrypted (handshake messages, data,
/// reports) or plaintext error signals (CoordsRequired, PathBroken)
/// generated by transit routers.
#[derive(Clone, Debug)]
pub struct SessionDatagram {
    /// Source node address (originator of this datagram).
    /// For data traffic: the source endpoint.
    /// For error signals: the transit router that generated the error.
    pub src_addr: NodeAddr,
    /// Destination node address (for routing decisions).
    pub dest_addr: NodeAddr,
    /// Time-to-live (decremented at each hop, dropped at zero).
    pub ttl: u8,
    /// Path MTU: minimum link MTU along the path so far.
    /// Each forwarding hop applies min(path_mtu, outgoing_link_mtu).
    pub path_mtu: u16,
    /// Session-layer payload (e2e encrypted or plaintext error signal).
    pub payload: Vec<u8>,
}

/// Borrowed view of a session datagram payload.
///
/// This avoids allocating and copying the inner payload when the caller only
/// needs to inspect or locally deliver it.
#[derive(Clone, Copy, Debug)]
pub struct SessionDatagramRef<'a> {
    pub src_addr: NodeAddr,
    pub dest_addr: NodeAddr,
    pub ttl: u8,
    pub path_mtu: u16,
    pub payload: &'a [u8],
}

/// SessionDatagram fixed header size: msg_type(1) + ttl(1) + path_mtu(2) + src_addr(16) + dest_addr(16).
pub const SESSION_DATAGRAM_HEADER_SIZE: usize = 36;

impl SessionDatagram {
    /// Create a new session datagram.
    pub fn new(src_addr: NodeAddr, dest_addr: NodeAddr, payload: Vec<u8>) -> Self {
        Self {
            src_addr,
            dest_addr,
            ttl: 64,
            path_mtu: u16::MAX,
            payload,
        }
    }

    /// Set the TTL.
    pub fn with_ttl(mut self, ttl: u8) -> Self {
        self.ttl = ttl;
        self
    }

    /// Set the path MTU.
    pub fn with_path_mtu(mut self, path_mtu: u16) -> Self {
        self.path_mtu = path_mtu;
        self
    }

    /// Decrement TTL, returning false if exhausted.
    pub fn decrement_ttl(&mut self) -> bool {
        if self.ttl > 0 {
            self.ttl -= 1;
            true
        } else {
            false
        }
    }

    /// Check if the datagram can be forwarded.
    pub fn can_forward(&self) -> bool {
        self.ttl > 0
    }

    /// Encode as link-layer message (msg_type + ttl + path_mtu + src_addr + dest_addr + payload).
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(SESSION_DATAGRAM_HEADER_SIZE + self.payload.len());
        buf.push(LinkMessageType::SessionDatagram.to_byte());
        buf.push(self.ttl);
        buf.extend_from_slice(&self.path_mtu.to_le_bytes());
        buf.extend_from_slice(self.src_addr.as_bytes());
        buf.extend_from_slice(self.dest_addr.as_bytes());
        buf.extend_from_slice(&self.payload);
        buf
    }

    /// Decode from link-layer payload (after msg_type byte has been consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, Error> {
        let view = SessionDatagramRef::decode(payload)?;
        Ok(view.into_owned())
    }
}

impl<'a> SessionDatagramRef<'a> {
    /// Decode a borrowed view from link-layer payload after the msg_type byte.
    pub fn decode(payload: &'a [u8]) -> Result<Self, Error> {
        // ttl(1) + path_mtu(2) + src_addr(16) + dest_addr(16) = 35
        if payload.len() < 35 {
            return Err(Error::MessageTooShort {
                expected: 35,
                got: payload.len(),
            });
        }
        let ttl = payload[0];
        let path_mtu = u16::from_le_bytes([payload[1], payload[2]]);
        let mut src_bytes = [0u8; 16];
        src_bytes.copy_from_slice(&payload[3..19]);
        let mut dest_bytes = [0u8; 16];
        dest_bytes.copy_from_slice(&payload[19..35]);

        Ok(Self {
            src_addr: NodeAddr::from_bytes(src_bytes),
            dest_addr: NodeAddr::from_bytes(dest_bytes),
            ttl,
            path_mtu,
            payload: &payload[35..],
        })
    }

    /// Materialize an owned datagram for forwarding/re-encoding paths.
    pub fn into_owned(self) -> SessionDatagram {
        SessionDatagram {
            src_addr: self.src_addr,
            dest_addr: self.dest_addr,
            ttl: self.ttl,
            path_mtu: self.path_mtu,
            payload: self.payload.to_vec(),
        }
    }
}

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

    // ===== LinkMessageType Tests =====

    #[test]
    fn test_link_message_type_roundtrip() {
        let types = [
            LinkMessageType::TreeAnnounce,
            LinkMessageType::FilterAnnounce,
            LinkMessageType::LookupRequest,
            LinkMessageType::LookupResponse,
            LinkMessageType::SessionDatagram,
            LinkMessageType::Disconnect,
            LinkMessageType::Heartbeat,
        ];

        for ty in types {
            let byte = ty.to_byte();
            let restored = LinkMessageType::from_byte(byte);
            assert_eq!(restored, Some(ty));
        }
    }

    #[test]
    fn test_link_message_type_invalid() {
        assert!(LinkMessageType::from_byte(0xFF).is_none());
        assert!(LinkMessageType::from_byte(0x03).is_none());
        assert!(LinkMessageType::from_byte(0x40).is_none());
    }

    // ===== SessionDatagram Tests =====

    fn make_node_addr(val: u8) -> NodeAddr {
        let mut bytes = [0u8; 16];
        bytes[0] = val;
        NodeAddr::from_bytes(bytes)
    }

    #[test]
    fn test_session_datagram_encode_decode() {
        let src = make_node_addr(0xAA);
        let dest = make_node_addr(0xBB);
        let payload = vec![0x10, 0x00, 0x05, 0x00, 1, 2, 3, 4, 5]; // session payload
        let dg = SessionDatagram::new(src, dest, payload.clone()).with_ttl(32);

        let encoded = dg.encode();
        assert_eq!(encoded[0], 0x00); // msg_type (SessionDatagram)
        assert_eq!(encoded.len(), SESSION_DATAGRAM_HEADER_SIZE + payload.len());

        // Decode (after msg_type)
        let decoded = SessionDatagram::decode(&encoded[1..]).unwrap();
        assert_eq!(decoded.src_addr, src);
        assert_eq!(decoded.dest_addr, dest);
        assert_eq!(decoded.ttl, 32);
        assert_eq!(decoded.payload, payload);
    }

    #[test]
    fn test_session_datagram_ref_decode_borrows_payload() {
        let src = make_node_addr(0xAA);
        let dest = make_node_addr(0xBB);
        let payload = vec![0x10, 0x00, 0x05, 0x00, 1, 2, 3, 4, 5];
        let dg = SessionDatagram::new(src, dest, payload.clone())
            .with_ttl(32)
            .with_path_mtu(1400);

        let encoded = dg.encode();
        let decoded = SessionDatagramRef::decode(&encoded[1..]).unwrap();

        assert_eq!(decoded.src_addr, src);
        assert_eq!(decoded.dest_addr, dest);
        assert_eq!(decoded.ttl, 32);
        assert_eq!(decoded.path_mtu, 1400);
        assert_eq!(decoded.payload, payload.as_slice());
        assert_eq!(
            decoded.payload.as_ptr(),
            encoded[SESSION_DATAGRAM_HEADER_SIZE..].as_ptr()
        );
    }

    #[test]
    #[ignore = "performance benchmark; run explicitly with --ignored --nocapture"]
    fn bench_session_datagram_decode_owned_vs_ref() {
        use std::hint::black_box;
        use std::time::Instant;

        const ITERS: usize = 300_000;

        let src = make_node_addr(0xAA);
        let dest = make_node_addr(0xBB);
        let payload = vec![0x5A; 1200];
        let datagram = SessionDatagram::new(src, dest, payload)
            .with_ttl(32)
            .with_path_mtu(1400);
        let encoded = datagram.encode();
        let link_payload = &encoded[1..];

        let ref_start = Instant::now();
        let mut ref_bytes = 0usize;
        for _ in 0..ITERS {
            let decoded = SessionDatagramRef::decode(black_box(link_payload)).unwrap();
            ref_bytes = ref_bytes.wrapping_add(decoded.payload.len());
            black_box(decoded);
        }
        let ref_elapsed = ref_start.elapsed();

        let owned_start = Instant::now();
        let mut owned_bytes = 0usize;
        for _ in 0..ITERS {
            let decoded = SessionDatagram::decode(black_box(link_payload)).unwrap();
            owned_bytes = owned_bytes.wrapping_add(decoded.payload.len());
            black_box(decoded);
        }
        let owned_elapsed = owned_start.elapsed();

        assert_eq!(ref_bytes, owned_bytes);
        println!(
            "SessionDatagram decode: ref={:.1} ns/op owned={:.1} ns/op speedup={:.2}x iters={} payload_bytes={}",
            ref_elapsed.as_secs_f64() * 1_000_000_000.0 / ITERS as f64,
            owned_elapsed.as_secs_f64() * 1_000_000_000.0 / ITERS as f64,
            owned_elapsed.as_secs_f64() / ref_elapsed.as_secs_f64(),
            ITERS,
            link_payload.len() - 35
        );
    }

    #[test]
    fn test_session_datagram_empty_payload() {
        let dg = SessionDatagram::new(make_node_addr(1), make_node_addr(2), Vec::new());

        let encoded = dg.encode();
        assert_eq!(encoded.len(), SESSION_DATAGRAM_HEADER_SIZE);

        let decoded = SessionDatagram::decode(&encoded[1..]).unwrap();
        assert!(decoded.payload.is_empty());
    }

    #[test]
    fn test_session_datagram_decode_too_short() {
        assert!(SessionDatagram::decode(&[]).is_err());
        assert!(SessionDatagram::decode(&[0x00; 20]).is_err());
    }

    #[test]
    fn test_session_datagram_ttl_roundtrip() {
        for hop in [0u8, 1, 64, 128, 255] {
            let dg = SessionDatagram::new(make_node_addr(1), make_node_addr(2), vec![0x42])
                .with_ttl(hop);

            let encoded = dg.encode();
            let decoded = SessionDatagram::decode(&encoded[1..]).unwrap();
            assert_eq!(decoded.ttl, hop);
        }
    }
}