//! Datagram framing: a 21-byte cleartext header authenticated as AAD, followed
//! by a ChaCha20-Poly1305 sealed payload.
//!
//! ```text
//! off size field
//!  0   1   ver_type — high nibble version (1), low nibble message type
//!  1   8   token_id u64 BE — random, server-assigned at login
//!  9  12   nonce — random; also this message's id
//! ```
//!
//! The header is cleartext because the server must read `token_id` to select a
//! key before it can decrypt anything; it is authenticated as AAD so a
//! ciphertext cannot be retargeted to another token or another message type.
//!
//! `datagram = header || ChaCha20Poly1305(K_dir, nonce, payload, aad = header)`

use chacha20poly1305::aead::{AeadInPlace, KeyInit};
use chacha20poly1305::{ChaCha20Poly1305, Tag};

use crate::error::DecodeError;
use crate::kdf::Key;
use crate::msg::{Message, MsgType, Nonce};

pub const VERSION: u8 = 1;
pub const HEADER_LEN: usize = 21;
pub const TAG_LEN: usize = 16;

/// Smallest possible datagram: header, empty payload, tag. No message type
/// actually has an empty payload, but this is the floor a length check can use
/// before it knows the type.
pub const MIN_DATAGRAM: usize = HEADER_LEN + TAG_LEN;

/// Path-MTU-safe ceiling for both IPv4 and IPv6. The 40-point `LOC` cap keeps
/// the largest real datagram at 998 bytes.
pub const MAX_DATAGRAM: usize = 1200;

/// Total datagram size for a payload of `payload_len` bytes.
#[must_use]
pub const fn datagram_len(payload_len: usize) -> usize {
    HEADER_LEN + payload_len + TAG_LEN
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
    pub msg_type: MsgType,
    /// The u64 assigned at login. A *credential*, not an identity: it
    /// authorises writing into its owner's position stream and never appears in
    /// a stored point's key.
    pub token_id: u64,
    pub nonce: Nonce,
}

impl Header {
    #[must_use]
    pub const fn new(msg_type: MsgType, token_id: u64, nonce: Nonce) -> Self {
        Self {
            msg_type,
            token_id,
            nonce,
        }
    }

    #[must_use]
    pub fn to_bytes(self) -> [u8; HEADER_LEN] {
        let mut b = [0u8; HEADER_LEN];
        b[0] = (VERSION << 4) | (self.msg_type as u8);
        b[1..9].copy_from_slice(&self.token_id.to_be_bytes());
        b[9..21].copy_from_slice(&self.nonce);
        b
    }

    /// Read the header of a datagram without decrypting it.
    ///
    /// This is the server's first look at a packet: it yields the `token_id`
    /// needed to pick a key, and it is where length, version and type filtering
    /// happen — all before any crypto work is spent on the packet.
    pub fn peek(datagram: &[u8]) -> Result<Self, DecodeError> {
        if datagram.len() < MIN_DATAGRAM {
            return Err(DecodeError::TooShort(datagram.len()));
        }
        if datagram.len() > MAX_DATAGRAM {
            return Err(DecodeError::TooLong(datagram.len()));
        }
        let version = datagram[0] >> 4;
        if version != VERSION {
            return Err(DecodeError::BadVersion(version));
        }
        Ok(Self {
            msg_type: MsgType::try_from(datagram[0] & 0x0F)?,
            token_id: u64::from_be_bytes(datagram[1..9].try_into().expect("length checked")),
            nonce: datagram[9..HEADER_LEN].try_into().expect("length checked"),
        })
    }
}

/// Seal a message into a datagram.
///
/// The nonce is passed in rather than generated here: this crate does no I/O and
/// owns no RNG, which is what makes it deterministically testable and lets the
/// golden vectors be reproducible. Callers must supply 12 fresh random bytes per
/// message — at 12 bytes the collision probability after 16 million messages is
/// about 2⁻⁴⁹, so a counter buys nothing and costs persistence.
#[must_use]
pub fn seal(key: &Key, header: Header, payload: &[u8]) -> Vec<u8> {
    debug_assert!(
        datagram_len(payload.len()) <= MAX_DATAGRAM,
        "payload of {} bytes exceeds the datagram budget",
        payload.len()
    );
    let aad = header.to_bytes();
    let mut out = Vec::with_capacity(datagram_len(payload.len()));
    out.extend_from_slice(&aad);
    out.extend_from_slice(payload);

    let cipher = ChaCha20Poly1305::new(key.into());
    let tag = cipher
        .encrypt_in_place_detached((&header.nonce).into(), &aad, &mut out[HEADER_LEN..])
        .expect("in-place detached encryption of a bounded buffer cannot fail");
    out.extend_from_slice(&tag);
    out
}

/// Seal an already-typed message, deriving the header from it.
#[must_use]
pub fn seal_message(key: &Key, token_id: u64, nonce: Nonce, msg: &Message) -> Vec<u8> {
    seal(
        key,
        Header::new(msg.msg_type(), token_id, nonce),
        &msg.encode_payload(),
    )
}

/// Open a datagram, returning its header and decrypted payload.
///
/// A [`DecodeError::AuthFailed`] must never be answered — not with a `NACK`, not
/// with anything. The server cannot know who sent it, and replying would make
/// the open port both a forgery oracle and a reflector.
pub fn open(key: &Key, datagram: &[u8]) -> Result<(Header, Vec<u8>), DecodeError> {
    let header = Header::peek(datagram)?;
    let (aad, rest) = datagram.split_at(HEADER_LEN);
    let (ciphertext, tag) = rest.split_at(rest.len() - TAG_LEN);

    let mut payload = ciphertext.to_vec();
    ChaCha20Poly1305::new(key.into())
        .decrypt_in_place_detached(
            (&header.nonce).into(),
            aad,
            &mut payload,
            Tag::from_slice(tag),
        )
        .map_err(|_| DecodeError::AuthFailed)?;
    Ok((header, payload))
}

/// Open a datagram and parse its payload.
pub fn open_message(key: &Key, datagram: &[u8]) -> Result<(Header, Message), DecodeError> {
    let (header, payload) = open(key, datagram)?;
    let msg = Message::decode_payload(header.msg_type, &payload)?;
    Ok((header, msg))
}

/// Ceiling on any datagram the server sends in reply to one it received.
///
/// This is the anti-amplification control, and it replaces an earlier rule that
/// `len(response) <= len(request)` for every message type. That rule was
/// achievable only by padding requests with reserved bytes — paying real bytes on
/// every `HELLO` and `PING` to make replies "fit" — and it protected against a
/// threat that authentication already eliminates:
///
/// **Nearly every reply is sent only to a datagram that passed AEAD
/// verification.** A bad tag gets silence. So a reflection attacker must already
/// hold a live token key, and the leverage they gain is the ratio below — against
/// DNS at ~50× and NTP `monlist` at ~550×, which is the scale at which reflection
/// is worth doing at all.
///
/// [`Revoked`] is the one exception, and it is deliberately the smallest message
/// in the protocol. The server cannot verify a datagram naming a token it has no
/// record of, yet that is exactly the device it must tell to log in again. Two
/// rules keep it from being useful: it is sent only when the request was at least
/// as long (see [`may_answer_unverified`], so the ratio never exceeds 1.0), and
/// the server rate-limits it per *destination* address, which for a spoofed
/// packet is the victim. Whether it is sent at all is a config switch.
///
/// | request | bytes | reply | bytes | ratio |
/// |---|---|---|---|---|
/// | `LOC`, 1 point | 62 | `ACK` | 51 | 0.82 |
/// | `LOC`, 40 points | 998 | `ACK` | 51 | 0.05 |
/// | `HELLO` | 43 | `ACK` | 51 | 1.19 |
/// | `PING` | 43 | `PONG` | 43 | 1.00 |
/// | `CONFIG_GET` | 39 | `CONFIG` | 49 | 1.26 |
/// | any, ≥ 38 B | 38 | `REVOKED` | 38 | ≤ 1.00 |
///
/// An absolute ceiling is also a stronger statement than a relative one: however
/// the protocol grows, the open port cannot be made to emit more than this many
/// bytes for one received datagram. A multi-nonce `ACK` is the one reply that
/// scales, and it scales with the number of datagrams *already received* from that
/// token, so it cannot amplify either — but nothing in this build emits one, and
/// [`fits_reply_budget`] holds for every reply it does emit.
pub const MAX_REPLY: usize = 64;

/// Whether a reply respects [`MAX_REPLY`].
#[must_use]
pub const fn fits_reply_budget(response_len: usize) -> bool {
    response_len <= MAX_REPLY
}

/// Size of a sealed [`Revoked`] notice: the smallest datagram this protocol can
/// produce, since its payload is one byte.
pub const REVOKED_DATAGRAM_LEN: usize = HEADER_LEN + 1 + TAG_LEN;

/// Whether a request is long enough to earn an unverifiable [`Revoked`] reply.
///
/// The server cannot authenticate a datagram naming a token it does not know, so
/// answering one means answering an address the sender merely claimed. That is a
/// reflector. It is a tolerable one only while it can never be an *amplifier*, so
/// the reply must not exceed the request — which for a fixed 38-byte reply is
/// just this length test.
///
/// [`MIN_DATAGRAM`] is 37, one byte short, so a minimum-size datagram earns
/// nothing. Every real message is far larger: the smallest a device ever sends is
/// a 39-byte `CONFIG_GET`.
#[must_use]
pub const fn may_answer_unverified(request_len: usize) -> bool {
    request_len >= REVOKED_DATAGRAM_LEN
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::kdf;
    use crate::msg::*;
    use crate::point::Point;

    const TOKEN_KEY: Key = [0x5A; 32];
    const TOKEN_ID: u64 = 0x1122_3344_5566_7788;
    const NONCE: Nonce = [0xA0; NONCE_LEN];

    fn keys() -> (Key, Key) {
        kdf::derive_both(&TOKEN_KEY)
    }

    #[test]
    fn header_round_trips() {
        let h = Header::new(MsgType::Loc, TOKEN_ID, NONCE);
        assert_eq!(
            Header::peek(&[h.to_bytes().as_slice(), &[0; TAG_LEN]].concat()),
            Ok(h)
        );
    }

    #[test]
    fn seal_open_round_trips_every_type() {
        let (up, down) = keys();
        let messages = [
            Message::Loc(vec![Point::new(1_785_000_042, 525_200_080, 134_050_000)]),
            Message::Ack(Ack::single(NONCE)),
            Message::Nack(Nack {
                nonce: NONCE,
                reason: NackReason::UnknownToken,
                retry_after_s: 0,
            }),
            Message::Hello(Hello {
                app_version_code: 2,
                os_api_level: 34,
                flags: HelloFlags::NONE,
                config_version: 1,
            }),
            Message::Config(Config::default()),
            Message::ConfigGet(ConfigGet { have_version: 1 }),
            Message::Ping(Ping { echo: 1, seq: 2 }),
            Message::Pong(Pong { echo: 1, seq: 3 }),
        ];
        for msg in messages {
            let key = if msg.msg_type().is_uplink() {
                &up
            } else {
                &down
            };
            let dg = seal_message(key, TOKEN_ID, NONCE, &msg);
            assert_eq!(dg.len(), datagram_len(msg.payload_len()));
            assert!(dg.len() <= MAX_DATAGRAM);
            let (h, back) = open_message(key, &dg).expect("opens");
            assert_eq!(h.token_id, TOKEN_ID);
            assert_eq!(h.msg_type, msg.msg_type());
            assert_eq!(back, msg);
        }
    }

    #[test]
    fn the_wrong_direction_key_cannot_open_a_datagram() {
        let (up, down) = keys();
        let dg = seal_message(&up, TOKEN_ID, NONCE, &Message::Ping(Ping::default()));
        assert_eq!(open(&down, &dg), Err(DecodeError::AuthFailed));
    }

    #[test]
    fn every_header_byte_is_authenticated() {
        let (up, _) = keys();
        let dg = seal_message(
            &up,
            TOKEN_ID,
            NONCE,
            &Message::Ping(Ping { echo: 9, seq: 1 }),
        );
        for i in 0..HEADER_LEN {
            let mut bad = dg.clone();
            bad[i] ^= 0x01;
            // A flipped version or type nibble fails the header check; anything
            // else fails the tag. Either way it never yields a message.
            assert!(
                open(&up, &bad).is_err(),
                "byte {i} of the header was not authenticated"
            );
        }
    }

    #[test]
    fn a_flipped_ciphertext_or_tag_byte_fails() {
        let (up, _) = keys();
        let dg = seal_message(&up, TOKEN_ID, NONCE, &Message::Ping(Ping::default()));
        for i in HEADER_LEN..dg.len() {
            let mut bad = dg.clone();
            bad[i] ^= 0x80;
            assert_eq!(open(&up, &bad), Err(DecodeError::AuthFailed), "byte {i}");
        }
    }

    #[test]
    fn truncation_is_caught_before_any_crypto() {
        let (up, _) = keys();
        let dg = seal_message(&up, TOKEN_ID, NONCE, &Message::Ping(Ping::default()));
        for cut in 0..MIN_DATAGRAM {
            assert_eq!(Header::peek(&dg[..cut]), Err(DecodeError::TooShort(cut)));
        }
        // Long enough to look like a header, short enough to be corrupt.
        for cut in MIN_DATAGRAM..dg.len() {
            assert!(
                open(&up, &dg[..cut]).is_err(),
                "truncation to {cut} accepted"
            );
        }
    }

    #[test]
    fn oversized_and_misversioned_datagrams_are_rejected() {
        assert_eq!(
            Header::peek(&vec![0x11; MAX_DATAGRAM + 1]),
            Err(DecodeError::TooLong(MAX_DATAGRAM + 1))
        );
        let mut dg = vec![0u8; MIN_DATAGRAM];
        dg[0] = 0x21; // version 2
        assert_eq!(Header::peek(&dg), Err(DecodeError::BadVersion(2)));
        dg[0] = 0x1F; // version 1, type 0xF
        assert_eq!(Header::peek(&dg), Err(DecodeError::BadMsgType(0xF)));
    }

    /// The control that keeps the open UDP port useless as a reflector.
    #[test]
    fn every_reply_fits_the_budget() {
        // Every message the server can send downstream.
        let replies = [
            Message::Ack(Ack::single(NONCE)),
            Message::Nack(Nack {
                nonce: NONCE,
                reason: NackReason::Malformed,
                retry_after_s: 0,
            }),
            Message::Config(Config::default()),
            Message::Pong(Pong::default()),
            Message::Revoked(Revoked {
                reason: RevokeReason::Revoked,
            }),
        ];
        for reply in &replies {
            let len = datagram_len(reply.payload_len());
            assert!(
                fits_reply_budget(len),
                "{:?} is {len} B, over the {MAX_REPLY} B reply budget",
                reply.msg_type(),
            );
        }
    }

    /// The unverifiable reply can never be an amplifier.
    ///
    /// This is the whole justification for REVOKED being one byte of payload.
    /// A datagram short enough to be profitable to reflect is short enough to be
    /// refused an answer.
    #[test]
    fn an_unverified_notice_never_amplifies() {
        let revoked = Message::Revoked(Revoked {
            reason: RevokeReason::Unknown,
        });
        assert_eq!(datagram_len(revoked.payload_len()), REVOKED_DATAGRAM_LEN);
        assert!(
            !may_answer_unverified(MIN_DATAGRAM),
            "the smallest possible datagram must not earn a reply"
        );
        for request_len in MIN_DATAGRAM..=MAX_DATAGRAM {
            if may_answer_unverified(request_len) {
                assert!(
                    REVOKED_DATAGRAM_LEN <= request_len,
                    "a {request_len} B request drew a {REVOKED_DATAGRAM_LEN} B reply"
                );
            }
        }
    }

    /// The ratios the reply budget actually permits, pinned so a future field
    /// cannot quietly turn the port into a useful reflector. Nothing here is
    /// remotely in DNS (~50x) or NTP monlist (~550x) territory, and every one of
    /// them still requires the sender to hold a valid token key.
    #[test]
    fn amplification_ratios_stay_near_one() {
        let hello = Message::Hello(Hello {
            app_version_code: 0,
            os_api_level: 0,
            flags: HelloFlags::NONE,
            config_version: 0,
        });
        let nack = Message::Nack(Nack {
            nonce: NONCE,
            reason: NackReason::Malformed,
            retry_after_s: 0,
        });
        let exchanges = [
            (
                Message::Loc(vec![Point::new(0, 0, 0)]),
                Message::Ack(Ack::single(NONCE)),
            ),
            (
                Message::Loc(vec![Point::new(0, 0, 0); MAX_POINTS]),
                Message::Ack(Ack::single(NONCE)),
            ),
            (hello, Message::Ack(Ack::single(NONCE))),
            (
                Message::Ping(Ping::default()),
                Message::Pong(Pong::default()),
            ),
            (
                Message::ConfigGet(ConfigGet::default()),
                Message::Config(Config::default()),
            ),
            (Message::Loc(vec![Point::new(0, 0, 0)]), nack),
        ];

        for (req, resp) in &exchanges {
            let req_len = datagram_len(req.payload_len()) as f64;
            let resp_len = datagram_len(resp.payload_len()) as f64;
            let ratio = resp_len / req_len;
            assert!(
                ratio <= 1.5,
                "{:?} -> {:?} amplifies {ratio:.2}x, more leverage than this design allows",
                req.msg_type(),
                resp.msg_type(),
            );
        }
    }

    /// The sizes the ratio table in the module documentation is computed from.
    #[test]
    fn documented_message_sizes_hold() {
        let sizes = [
            (Message::Ack(Ack::single(NONCE)), 51),
            (
                Message::Nack(Nack {
                    nonce: NONCE,
                    reason: NackReason::Malformed,
                    retry_after_s: 0,
                }),
                51,
            ),
            (Message::Config(Config::default()), 49),
            (
                Message::Revoked(Revoked {
                    reason: RevokeReason::Revoked,
                }),
                REVOKED_DATAGRAM_LEN,
            ),
            (Message::ConfigGet(ConfigGet::default()), 39),
            (Message::Ping(Ping::default()), 43),
            (Message::Pong(Pong::default()), 43),
            (
                Message::Hello(Hello {
                    app_version_code: 0,
                    os_api_level: 0,
                    flags: HelloFlags::NONE,
                    config_version: 0,
                }),
                43,
            ),
        ];
        for (msg, want) in &sizes {
            assert_eq!(
                datagram_len(msg.payload_len()),
                *want,
                "{:?} changed size, which moves the amplification ratios",
                msg.msg_type()
            );
        }
    }

    #[test]
    fn documented_wire_sizes_hold() {
        let one = datagram_len(Message::Loc(vec![Point::new(0, 0, 0)]).payload_len());
        let twenty = datagram_len(Message::Loc(vec![Point::new(0, 0, 0); 20]).payload_len());
        let forty = datagram_len(Message::Loc(vec![Point::new(0, 0, 0); MAX_POINTS]).payload_len());
        assert_eq!((one, twenty, forty), (62, 518, 998));
    }
}