path: root/src/msg.rs

//! Message body parsing and writing.
//!
//! Each protocol message has a fixed header (parsed by
//! `wire::MsgHeader`) followed by a variable body. This
//! module provides parse/write for all body types.
//!
//! All multi-byte fields are big-endian (network order).

use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

use crate::error::Error;
use crate::id::{ID_LEN, NodeId};
use crate::peers::PeerInfo;
use crate::wire::{DOMAIN_INET, HEADER_SIZE};

// ── Helpers ─────────────────────────────────────────

// Callers MUST validate `off + 2 <= buf.len()` before calling.
fn read_u16(buf: &[u8], off: usize) -> u16 {
    u16::from_be_bytes([buf[off], buf[off + 1]])
}

fn read_u32(buf: &[u8], off: usize) -> u32 {
    u32::from_be_bytes([buf[off], buf[off + 1], buf[off + 2], buf[off + 3]])
}

fn write_u16(buf: &mut [u8], off: usize, v: u16) {
    buf[off..off + 2].copy_from_slice(&v.to_be_bytes());
}

fn write_u32(buf: &mut [u8], off: usize, v: u32) {
    buf[off..off + 4].copy_from_slice(&v.to_be_bytes());
}

// ── Ping (DHT + DTUN) ──────────────────────────────

/// Body: just a nonce (4 bytes after header).
/// Used by: DhtPing, DhtPingReply, DtunPing,
/// DtunPingReply, DtunRequestReply.
pub fn parse_ping(buf: &[u8]) -> Result<u32, Error> {
    if buf.len() < HEADER_SIZE + 4 {
        return Err(Error::BufferTooSmall);
    }
    Ok(read_u32(buf, HEADER_SIZE))
}

pub fn write_ping(buf: &mut [u8], nonce: u32) {
    write_u32(buf, HEADER_SIZE, nonce);
}

/// Total size of a ping message.
pub const PING_MSG_SIZE: usize = HEADER_SIZE + 4;

// ── NAT Echo ────────────────────────────────────────

/// Parse NatEcho: nonce only.
pub fn parse_nat_echo(buf: &[u8]) -> Result<u32, Error> {
    parse_ping(buf) // same layout
}

/// NatEchoReply body: nonce(4) + domain(2) + port(2) + addr(16).
pub const NAT_ECHO_REPLY_BODY: usize = 4 + 2 + 2 + 16;

#[derive(Debug, Clone)]
pub struct NatEchoReply {
    pub nonce: u32,
    pub domain: u16,
    pub port: u16,
    pub addr: [u8; 16],
}

pub fn parse_nat_echo_reply(buf: &[u8]) -> Result<NatEchoReply, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + NAT_ECHO_REPLY_BODY {
        return Err(Error::BufferTooSmall);
    }
    Ok(NatEchoReply {
        nonce: read_u32(buf, off),
        domain: read_u16(buf, off + 4),
        port: read_u16(buf, off + 6),
        addr: {
            let mut a = [0u8; 16];
            a.copy_from_slice(&buf[off + 8..off + 24]);
            a
        },
    })
}

pub fn write_nat_echo_reply(buf: &mut [u8], reply: &NatEchoReply) {
    let off = HEADER_SIZE;
    write_u32(buf, off, reply.nonce);
    write_u16(buf, off + 4, reply.domain);
    write_u16(buf, off + 6, reply.port);
    buf[off + 8..off + 24].copy_from_slice(&reply.addr);
}

/// NatEchoRedirect body: nonce(4) + port(2) + padding(2).
pub fn parse_nat_echo_redirect(buf: &[u8]) -> Result<(u32, u16), Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + 8 {
        return Err(Error::BufferTooSmall);
    }
    Ok((read_u32(buf, off), read_u16(buf, off + 4)))
}

/// Write NatEchoRedirect body.
pub fn write_nat_echo_redirect(buf: &mut [u8], nonce: u32, port: u16) {
    let off = HEADER_SIZE;
    write_u32(buf, off, nonce);
    write_u16(buf, off + 4, port);
    write_u16(buf, off + 6, 0); // padding
}

/// Size of a NatEchoRedirect message.
pub const NAT_ECHO_REDIRECT_SIZE: usize = HEADER_SIZE + 8;

// ── FindNode (DHT + DTUN) ───────────────────────────

/// FindNode body: nonce(4) + id(20) + domain(2) + state_or_pad(2).
pub const FIND_NODE_BODY: usize = 4 + ID_LEN + 2 + 2;

#[derive(Debug, Clone)]
pub struct FindNodeMsg {
    pub nonce: u32,
    pub target: NodeId,
    pub domain: u16,
    pub state: u16,
}

pub fn parse_find_node(buf: &[u8]) -> Result<FindNodeMsg, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + FIND_NODE_BODY {
        return Err(Error::BufferTooSmall);
    }
    Ok(FindNodeMsg {
        nonce: read_u32(buf, off),
        target: NodeId::read_from(&buf[off + 4..off + 4 + ID_LEN]),
        domain: read_u16(buf, off + 4 + ID_LEN),
        state: read_u16(buf, off + 6 + ID_LEN),
    })
}

pub fn write_find_node(buf: &mut [u8], msg: &FindNodeMsg) {
    let off = HEADER_SIZE;
    write_u32(buf, off, msg.nonce);
    msg.target.write_to(&mut buf[off + 4..off + 4 + ID_LEN]);
    write_u16(buf, off + 4 + ID_LEN, msg.domain);
    write_u16(buf, off + 6 + ID_LEN, msg.state);
}

pub const FIND_NODE_MSG_SIZE: usize = HEADER_SIZE + FIND_NODE_BODY;

// ── FindNodeReply (DHT + DTUN) ──────────────────────

/// FindNodeReply fixed part: nonce(4) + id(20) +
/// domain(2) + num(1) + padding(1).
pub const FIND_NODE_REPLY_FIXED: usize = 4 + ID_LEN + 4;

#[derive(Debug, Clone)]
pub struct FindNodeReplyMsg {
    pub nonce: u32,
    pub id: NodeId,
    pub domain: u16,
    pub nodes: Vec<PeerInfo>,
}

/// Size of an IPv4 node entry: port(2) + reserved(2) +
/// addr(4) + id(ID_LEN).
pub const INET_NODE_SIZE: usize = 8 + ID_LEN;

/// Size of an IPv6 node entry: port(2) + reserved(2) +
/// addr(16) + id(ID_LEN).
pub const INET6_NODE_SIZE: usize = 20 + ID_LEN;

pub fn parse_find_node_reply(buf: &[u8]) -> Result<FindNodeReplyMsg, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + FIND_NODE_REPLY_FIXED {
        return Err(Error::BufferTooSmall);
    }

    let nonce = read_u32(buf, off);
    let id = NodeId::read_from(&buf[off + 4..off + 4 + ID_LEN]);
    let domain = read_u16(buf, off + 4 + ID_LEN);
    let num = buf[off + 4 + ID_LEN + 2] as usize;

    let nodes_off = off + FIND_NODE_REPLY_FIXED;
    let nodes = if domain == DOMAIN_INET {
        read_nodes_inet(&buf[nodes_off..], num)
    } else {
        read_nodes_inet6(&buf[nodes_off..], num)
    };

    Ok(FindNodeReplyMsg {
        nonce,
        id,
        domain,
        nodes,
    })
}

pub fn write_find_node_reply(buf: &mut [u8], msg: &FindNodeReplyMsg) -> usize {
    let off = HEADER_SIZE;
    write_u32(buf, off, msg.nonce);
    msg.id.write_to(&mut buf[off + 4..off + 4 + ID_LEN]);
    write_u16(buf, off + 4 + ID_LEN, msg.domain);
    let num = msg.nodes.len().min(MAX_NODES_PER_REPLY);
    buf[off + 4 + ID_LEN + 2] = num as u8;
    buf[off + 4 + ID_LEN + 3] = 0; // padding

    let nodes_off = off + FIND_NODE_REPLY_FIXED;
    let nodes_len = if msg.domain == DOMAIN_INET {
        write_nodes_inet(&mut buf[nodes_off..], &msg.nodes)
    } else {
        write_nodes_inet6(&mut buf[nodes_off..], &msg.nodes)
    };

    HEADER_SIZE + FIND_NODE_REPLY_FIXED + nodes_len
}

// ── Store (DHT) ─────────────────────────────────────

/// Store fixed part: id(20) + from(20) + keylen(2) +
/// valuelen(2) + ttl(2) + flags(1) + reserved(1) = 48.
pub const STORE_FIXED: usize = ID_LEN * 2 + 8;

#[derive(Debug, Clone)]
pub struct StoreMsg {
    pub id: NodeId,
    pub from: NodeId,
    pub key: Vec<u8>,
    pub value: Vec<u8>,
    pub ttl: u16,
    pub is_unique: bool,
}

pub fn parse_store(buf: &[u8]) -> Result<StoreMsg, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + STORE_FIXED {
        return Err(Error::BufferTooSmall);
    }

    let id = NodeId::read_from(&buf[off..off + ID_LEN]);
    let from = NodeId::read_from(&buf[off + ID_LEN..off + ID_LEN * 2]);
    let keylen = read_u16(buf, off + ID_LEN * 2) as usize;
    let valuelen = read_u16(buf, off + ID_LEN * 2 + 2) as usize;
    let ttl = read_u16(buf, off + ID_LEN * 2 + 4);
    let flags = buf[off + ID_LEN * 2 + 6];

    let data_off = off + STORE_FIXED;
    let total = data_off
        .checked_add(keylen)
        .and_then(|v| v.checked_add(valuelen))
        .ok_or(Error::InvalidMessage)?;

    if buf.len() < total {
        return Err(Error::BufferTooSmall);
    }

    let key = buf[data_off..data_off + keylen].to_vec();
    let value = buf[data_off + keylen..data_off + keylen + valuelen].to_vec();

    Ok(StoreMsg {
        id,
        from,
        key,
        value,
        ttl,
        is_unique: flags & crate::wire::DHT_FLAG_UNIQUE != 0,
    })
}

pub fn write_store(buf: &mut [u8], msg: &StoreMsg) -> Result<usize, Error> {
    let off = HEADER_SIZE;
    let total = off + STORE_FIXED + msg.key.len() + msg.value.len();
    if buf.len() < total {
        return Err(Error::BufferTooSmall);
    }

    msg.id.write_to(&mut buf[off..off + ID_LEN]);
    msg.from.write_to(&mut buf[off + ID_LEN..off + ID_LEN * 2]);
    let keylen =
        u16::try_from(msg.key.len()).map_err(|_| Error::BufferTooSmall)?;
    let valuelen =
        u16::try_from(msg.value.len()).map_err(|_| Error::BufferTooSmall)?;
    write_u16(buf, off + ID_LEN * 2, keylen);
    write_u16(buf, off + ID_LEN * 2 + 2, valuelen);
    write_u16(buf, off + ID_LEN * 2 + 4, msg.ttl);
    buf[off + ID_LEN * 2 + 6] = if msg.is_unique {
        crate::wire::DHT_FLAG_UNIQUE
    } else {
        0
    };
    buf[off + ID_LEN * 2 + 7] = 0; // reserved

    let data_off = off + STORE_FIXED;
    buf[data_off..data_off + msg.key.len()].copy_from_slice(&msg.key);
    buf[data_off + msg.key.len()..data_off + msg.key.len() + msg.value.len()]
        .copy_from_slice(&msg.value);

    Ok(total)
}

// ── FindValue (DHT) ─────────────────────────────────

/// FindValue fixed: nonce(4) + id(20) + domain(2) +
/// keylen(2) + flag(1) + padding(3) = 32.
pub const FIND_VALUE_FIXED: usize = 4 + ID_LEN + 8;

#[derive(Debug, Clone)]
pub struct FindValueMsg {
    pub nonce: u32,
    pub target: NodeId,
    pub domain: u16,
    pub key: Vec<u8>,
    pub use_rdp: bool,
}

pub fn parse_find_value(buf: &[u8]) -> Result<FindValueMsg, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + FIND_VALUE_FIXED {
        return Err(Error::BufferTooSmall);
    }

    let nonce = read_u32(buf, off);
    let target = NodeId::read_from(&buf[off + 4..off + 4 + ID_LEN]);
    let domain = read_u16(buf, off + 4 + ID_LEN);
    let keylen = read_u16(buf, off + 6 + ID_LEN) as usize;
    let flag = buf[off + 8 + ID_LEN];

    let key_off = off + FIND_VALUE_FIXED;
    if buf.len() < key_off + keylen {
        return Err(Error::BufferTooSmall);
    }

    Ok(FindValueMsg {
        nonce,
        target,
        domain,
        key: buf[key_off..key_off + keylen].to_vec(),
        use_rdp: flag == 1,
    })
}

pub fn write_find_value(
    buf: &mut [u8],
    msg: &FindValueMsg,
) -> Result<usize, Error> {
    let off = HEADER_SIZE;
    let total = off + FIND_VALUE_FIXED + msg.key.len();
    if buf.len() < total {
        return Err(Error::BufferTooSmall);
    }

    write_u32(buf, off, msg.nonce);
    msg.target.write_to(&mut buf[off + 4..off + 4 + ID_LEN]);
    write_u16(buf, off + 4 + ID_LEN, msg.domain);
    let keylen =
        u16::try_from(msg.key.len()).map_err(|_| Error::BufferTooSmall)?;
    write_u16(buf, off + 6 + ID_LEN, keylen);
    buf[off + 8 + ID_LEN] = if msg.use_rdp { 1 } else { 0 };
    buf[off + 9 + ID_LEN] = 0;
    buf[off + 10 + ID_LEN] = 0;
    buf[off + 11 + ID_LEN] = 0;

    let key_off = off + FIND_VALUE_FIXED;
    buf[key_off..key_off + msg.key.len()].copy_from_slice(&msg.key);

    Ok(total)
}

// ── FindValueReply (DHT) ────────────────────────────

/// Fixed: nonce(4) + id(20) + index(2) + total(2) +
/// flag(1) + padding(3) = 32.
pub const FIND_VALUE_REPLY_FIXED: usize = 4 + ID_LEN + 8;

#[derive(Debug, Clone)]
pub enum FindValueReplyData {
    /// flag=0xa0: node list
    Nodes { domain: u16, nodes: Vec<PeerInfo> },

    /// flag=0xa1: a value chunk
    Value {
        index: u16,
        total: u16,
        data: Vec<u8>,
    },

    /// flag=0xa2: no data
    Nul,
}

#[derive(Debug, Clone)]
pub struct FindValueReplyMsg {
    pub nonce: u32,
    pub id: NodeId,
    pub data: FindValueReplyData,
}

pub fn parse_find_value_reply(buf: &[u8]) -> Result<FindValueReplyMsg, Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + FIND_VALUE_REPLY_FIXED {
        return Err(Error::BufferTooSmall);
    }

    let nonce = read_u32(buf, off);
    let id = NodeId::read_from(&buf[off + 4..off + 4 + ID_LEN]);
    let index = read_u16(buf, off + 4 + ID_LEN);
    let total = read_u16(buf, off + 6 + ID_LEN);
    let flag = buf[off + 8 + ID_LEN];

    let data_off = off + FIND_VALUE_REPLY_FIXED;

    let data = match flag {
        crate::wire::DATA_ARE_NODES => {
            if buf.len() < data_off + 4 {
                return Err(Error::BufferTooSmall);
            }
            let domain = read_u16(buf, data_off);
            let num = buf[data_off + 2] as usize;
            let nodes_off = data_off + 4;
            let nodes = if domain == DOMAIN_INET {
                read_nodes_inet(&buf[nodes_off..], num)
            } else {
                read_nodes_inet6(&buf[nodes_off..], num)
            };
            FindValueReplyData::Nodes { domain, nodes }
        }
        crate::wire::DATA_ARE_VALUES => {
            let payload = buf[data_off..].to_vec();
            FindValueReplyData::Value {
                index,
                total,
                data: payload,
            }
        }
        crate::wire::DATA_ARE_NUL => FindValueReplyData::Nul,
        _ => return Err(Error::InvalidMessage),
    };

    Ok(FindValueReplyMsg { nonce, id, data })
}

// ── DtunRegister ────────────────────────────────────

pub fn parse_dtun_register(buf: &[u8]) -> Result<u32, Error> {
    if buf.len() < HEADER_SIZE + 4 {
        return Err(Error::BufferTooSmall);
    }
    Ok(read_u32(buf, HEADER_SIZE)) // session
}

// ── DtunRequest ─────────────────────────────────────

pub fn parse_dtun_request(buf: &[u8]) -> Result<(u32, NodeId), Error> {
    let off = HEADER_SIZE;
    if buf.len() < off + 4 + ID_LEN {
        return Err(Error::BufferTooSmall);
    }
    let nonce = read_u32(buf, off);
    let target = NodeId::read_from(&buf[off + 4..off + 4 + ID_LEN]);
    Ok((nonce, target))
}

// ── Node list serialization (IPv4 / IPv6) ───────────

/// Maximum nodes per reply (prevents OOM from malicious num).
const MAX_NODES_PER_REPLY: usize = 20;

/// Read `num` IPv4 node entries from `buf`.
pub fn read_nodes_inet(buf: &[u8], num: usize) -> Vec<PeerInfo> {
    let num = num.min(MAX_NODES_PER_REPLY);
    let mut nodes = Vec::with_capacity(num);
    for i in 0..num {
        let off = i * INET_NODE_SIZE;
        if off + INET_NODE_SIZE > buf.len() {
            break;
        }
        let port = read_u16(buf, off);
        let ip = Ipv4Addr::new(
            buf[off + 4],
            buf[off + 5],
            buf[off + 6],
            buf[off + 7],
        );
        let id = NodeId::read_from(&buf[off + 8..off + 8 + ID_LEN]);
        let addr = SocketAddr::V4(SocketAddrV4::new(ip, port));
        nodes.push(PeerInfo::new(id, addr));
    }
    nodes
}

/// Read `num` IPv6 node entries from `buf`.
pub fn read_nodes_inet6(buf: &[u8], num: usize) -> Vec<PeerInfo> {
    let num = num.min(MAX_NODES_PER_REPLY);
    let mut nodes = Vec::with_capacity(num);
    for i in 0..num {
        let off = i * INET6_NODE_SIZE;
        if off + INET6_NODE_SIZE > buf.len() {
            break;
        }
        let port = read_u16(buf, off);
        let mut octets = [0u8; 16];
        octets.copy_from_slice(&buf[off + 4..off + 20]);
        let ip = Ipv6Addr::from(octets);
        let id = NodeId::read_from(&buf[off + 20..off + 20 + ID_LEN]);
        let addr = SocketAddr::V6(SocketAddrV6::new(ip, port, 0, 0));
        nodes.push(PeerInfo::new(id, addr));
    }
    nodes
}

/// Write IPv4 node entries. Returns bytes written.
pub fn write_nodes_inet(buf: &mut [u8], nodes: &[PeerInfo]) -> usize {
    let mut written = 0;
    for node in nodes {
        if written + INET_NODE_SIZE > buf.len() {
            break;
        }
        let off = written;
        write_u16(buf, off, node.addr.port());
        write_u16(buf, off + 2, 0); // reserved

        if let SocketAddr::V4(v4) = node.addr {
            let octets = v4.ip().octets();
            buf[off + 4..off + 8].copy_from_slice(&octets);
        } else {
            buf[off + 4..off + 8].fill(0);
        }

        node.id.write_to(&mut buf[off + 8..off + 8 + ID_LEN]);
        written += INET_NODE_SIZE;
    }
    written
}

/// Write IPv6 node entries. Returns bytes written.
pub fn write_nodes_inet6(buf: &mut [u8], nodes: &[PeerInfo]) -> usize {
    let mut written = 0;
    for node in nodes {
        if written + INET6_NODE_SIZE > buf.len() {
            break;
        }
        let off = written;
        write_u16(buf, off, node.addr.port());
        write_u16(buf, off + 2, 0); // reserved

        if let SocketAddr::V6(v6) = node.addr {
            let octets = v6.ip().octets();
            buf[off + 4..off + 20].copy_from_slice(&octets);
        } else {
            buf[off + 4..off + 20].fill(0);
        }

        node.id.write_to(&mut buf[off + 20..off + 20 + ID_LEN]);
        written += INET6_NODE_SIZE;
    }
    written
}

/// Create a PeerInfo from a message header and source
/// address.
pub fn peer_from_header(src_id: NodeId, from: SocketAddr) -> PeerInfo {
    PeerInfo::new(src_id, from)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::wire::{MsgHeader, MsgType};

    fn make_buf(msg_type: MsgType, body_len: usize) -> Vec<u8> {
        let total = HEADER_SIZE + body_len;
        let mut buf = vec![0u8; total];
        let hdr = MsgHeader::new(
            msg_type,
            total as u16,
            NodeId::from_bytes([0xAA; 32]),
            NodeId::from_bytes([0xBB; 32]),
        );
        hdr.write(&mut buf).unwrap();
        buf
    }

    // ── Ping ────────────────────────────────────────

    #[test]
    fn ping_roundtrip() {
        let mut buf = make_buf(MsgType::DhtPing, 4);
        write_ping(&mut buf, 0xDEADBEEF);
        let nonce = parse_ping(&buf).unwrap();
        assert_eq!(nonce, 0xDEADBEEF);
    }

    // ── NatEchoReply ────────────────────────────────

    #[test]
    fn nat_echo_reply_roundtrip() {
        let mut buf = make_buf(MsgType::NatEchoReply, NAT_ECHO_REPLY_BODY);
        let reply = NatEchoReply {
            nonce: 42,
            domain: DOMAIN_INET,
            port: 3000,
            addr: {
                let mut a = [0u8; 16];
                a[0..4].copy_from_slice(&[192, 168, 1, 1]);
                a
            },
        };
        write_nat_echo_reply(&mut buf, &reply);
        let parsed = parse_nat_echo_reply(&buf).unwrap();
        assert_eq!(parsed.nonce, 42);
        assert_eq!(parsed.domain, DOMAIN_INET);
        assert_eq!(parsed.port, 3000);
        assert_eq!(parsed.addr[0..4], [192, 168, 1, 1]);
    }

    // ── FindNode ────────────────────────────────────

    #[test]
    fn find_node_roundtrip() {
        let mut buf = make_buf(MsgType::DhtFindNode, FIND_NODE_BODY);
        let msg = FindNodeMsg {
            nonce: 99,
            target: NodeId::from_bytes([0x42; 32]),
            domain: DOMAIN_INET,
            state: 1,
        };
        write_find_node(&mut buf, &msg);
        let parsed = parse_find_node(&buf).unwrap();
        assert_eq!(parsed.nonce, 99);
        assert_eq!(parsed.target, msg.target);
        assert_eq!(parsed.domain, DOMAIN_INET);
    }

    // ── FindNodeReply ───────────────────────────────

    #[test]
    fn find_node_reply_roundtrip() {
        let nodes = vec![
            PeerInfo::new(
                NodeId::from_bytes([0x01; 32]),
                "127.0.0.1:3000".parse().unwrap(),
            ),
            PeerInfo::new(
                NodeId::from_bytes([0x02; 32]),
                "127.0.0.1:3001".parse().unwrap(),
            ),
        ];
        let msg = FindNodeReplyMsg {
            nonce: 55,
            id: NodeId::from_bytes([0x42; 32]),
            domain: DOMAIN_INET,
            nodes: nodes.clone(),
        };

        let mut buf = vec![0u8; 1024];
        let hdr = MsgHeader::new(
            MsgType::DhtFindNodeReply,
            0, // will fix
            NodeId::from_bytes([0xAA; 32]),
            NodeId::from_bytes([0xBB; 32]),
        );
        hdr.write(&mut buf).unwrap();
        let _total = write_find_node_reply(&mut buf, &msg);

        let parsed = parse_find_node_reply(&buf).unwrap();
        assert_eq!(parsed.nonce, 55);
        assert_eq!(parsed.nodes.len(), 2);
        assert_eq!(parsed.nodes[0].id, nodes[0].id);
        assert_eq!(parsed.nodes[0].addr.port(), 3000);
        assert_eq!(parsed.nodes[1].id, nodes[1].id);
    }

    // ── Store ───────────────────────────────────────

    #[test]
    fn store_roundtrip() {
        let msg = StoreMsg {
            id: NodeId::from_bytes([0x10; 32]),
            from: NodeId::from_bytes([0x20; 32]),
            key: b"mykey".to_vec(),
            value: b"myvalue".to_vec(),
            ttl: 300,
            is_unique: true,
        };
        let total = HEADER_SIZE + STORE_FIXED + msg.key.len() + msg.value.len();
        let mut buf = vec![0u8; total];
        let hdr = MsgHeader::new(
            MsgType::DhtStore,
            total as u16,
            NodeId::from_bytes([0xAA; 32]),
            NodeId::from_bytes([0xBB; 32]),
        );
        hdr.write(&mut buf).unwrap();
        write_store(&mut buf, &msg).unwrap();

        let parsed = parse_store(&buf).unwrap();
        assert_eq!(parsed.id, msg.id);
        assert_eq!(parsed.from, msg.from);
        assert_eq!(parsed.key, b"mykey");
        assert_eq!(parsed.value, b"myvalue");
        assert_eq!(parsed.ttl, 300);
        assert!(parsed.is_unique);
    }

    #[test]
    fn store_not_unique() {
        let msg = StoreMsg {
            id: NodeId::from_bytes([0x10; 32]),
            from: NodeId::from_bytes([0x20; 32]),
            key: b"k".to_vec(),
            value: b"v".to_vec(),
            ttl: 60,
            is_unique: false,
        };
        let total = HEADER_SIZE + STORE_FIXED + msg.key.len() + msg.value.len();
        let mut buf = vec![0u8; total];
        let hdr = MsgHeader::new(
            MsgType::DhtStore,
            total as u16,
            NodeId::from_bytes([0xAA; 32]),
            NodeId::from_bytes([0xBB; 32]),
        );
        hdr.write(&mut buf).unwrap();
        write_store(&mut buf, &msg).unwrap();

        let parsed = parse_store(&buf).unwrap();
        assert!(!parsed.is_unique);
    }

    // ── FindValue ───────────────────────────────────

    #[test]
    fn find_value_roundtrip() {
        let msg = FindValueMsg {
            nonce: 77,
            target: NodeId::from_bytes([0x33; 32]),
            domain: DOMAIN_INET,
            key: b"lookup-key".to_vec(),
            use_rdp: false,
        };
        let total = HEADER_SIZE + FIND_VALUE_FIXED + msg.key.len();
        let mut buf = vec![0u8; total];
        let hdr = MsgHeader::new(
            MsgType::DhtFindValue,
            total as u16,
            NodeId::from_bytes([0xAA; 32]),
            NodeId::from_bytes([0xBB; 32]),
        );
        hdr.write(&mut buf).unwrap();
        write_find_value(&mut buf, &msg).unwrap();

        let parsed = parse_find_value(&buf).unwrap();
        assert_eq!(parsed.nonce, 77);
        assert_eq!(parsed.target, msg.target);
        assert_eq!(parsed.key, b"lookup-key");
        assert!(!parsed.use_rdp);
    }

    // ── DtunRegister ────────────────────────────────

    #[test]
    fn dtun_register_roundtrip() {
        let mut buf = make_buf(MsgType::DtunRegister, 4);
        write_u32(&mut buf, HEADER_SIZE, 12345);
        let session = parse_dtun_register(&buf).unwrap();
        assert_eq!(session, 12345);
    }

    // ── DtunRequest ─────────────────────────────────

    #[test]
    fn dtun_request_roundtrip() {
        let mut buf = make_buf(MsgType::DtunRequest, 4 + ID_LEN);
        let nonce = 88u32;
        let target = NodeId::from_bytes([0x55; 32]);
        write_u32(&mut buf, HEADER_SIZE, nonce);
        target.write_to(&mut buf[HEADER_SIZE + 4..HEADER_SIZE + 4 + ID_LEN]);

        let (n, t) = parse_dtun_request(&buf).unwrap();
        assert_eq!(n, 88);
        assert_eq!(t, target);
    }

    // ── Node list ───────────────────────────────────

    #[test]
    fn inet_nodes_roundtrip() {
        let nodes = vec![
            PeerInfo::new(
                NodeId::from_bytes([0x01; 32]),
                "10.0.0.1:8000".parse().unwrap(),
            ),
            PeerInfo::new(
                NodeId::from_bytes([0x02; 32]),
                "10.0.0.2:9000".parse().unwrap(),
            ),
        ];
        let mut buf = vec![0u8; INET_NODE_SIZE * 2];
        let written = write_nodes_inet(&mut buf, &nodes);
        assert_eq!(written, INET_NODE_SIZE * 2);

        let parsed = read_nodes_inet(&buf, 2);
        assert_eq!(parsed.len(), 2);
        assert_eq!(parsed[0].id, nodes[0].id);
        assert_eq!(parsed[0].addr.port(), 8000);
        assert_eq!(parsed[1].addr.port(), 9000);
    }

    // ── Truncated inputs ────────────────────────────

    #[test]
    fn parse_store_truncated() {
        let buf = make_buf(MsgType::DhtStore, 2); // too small
        assert!(matches!(parse_store(&buf), Err(Error::BufferTooSmall)));
    }

    #[test]
    fn parse_find_node_truncated() {
        let buf = make_buf(MsgType::DhtFindNode, 2);
        assert!(matches!(parse_find_node(&buf), Err(Error::BufferTooSmall)));
    }

    #[test]
    fn parse_find_value_truncated() {
        let buf = make_buf(MsgType::DhtFindValue, 2);
        assert!(matches!(parse_find_value(&buf), Err(Error::BufferTooSmall)));
    }
}