feat(crypto,proto): explicit-nonce AeadKey + datagram record codec

Contract for the custom UDP transport (v2): - aura-crypto: AeadKey — ChaCha20-Poly1305 with an EXPLICIT per-message nonce (caller passes the counter), for datagram transports where packets may be lost or reordered. AeadSession::into_parts() hands off (AeadKey, counter). Same nonce scheme as AeadSession, so they interoperate on one key with disjoint counter ranges. +4 tests. - aura-proto: DatagramSender/DatagramReceiver (record = seq(8) || AEAD(frame, aad=seq), sliding replay window) and Session::into_datagram_parts(); reuse for a UDP data path. +1 test. Existing 16 crypto / 13 proto tests still green. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-25 18:57:56 +03:00
parent 46513354c0
commit fa9f18ec17
4 changed files with 277 additions and 10 deletions
@@ -98,6 +98,14 @@ impl AeadSession {
        result.map_err(|_| CryptoError::AeadDecrypt)
    }

+    /// Consume the session, returning a reusable explicit-nonce [`AeadKey`] plus the current
+    /// counter value. Datagram transports use this to continue from the post-handshake counter
+    /// while carrying the nonce on the wire (see [`AeadKey`]).
+    #[must_use]
+    pub fn into_parts(self) -> (AeadKey, u64) {
+        (AeadKey::new(self.key), self.counter)
+    }
+
    /// Current counter value (next nonce to be used). Test-only accessor.
    #[cfg(test)]
    #[must_use]
@@ -114,6 +122,79 @@ impl Drop for AeadSession {

 impl zeroize::ZeroizeOnDrop for AeadSession {}

+/// A 256-bit ChaCha20-Poly1305 key used with EXPLICIT per-message nonces.
+///
+/// Unlike [`AeadSession`] (which derives the nonce from an internal, lock-step counter), `AeadKey`
+/// takes the nonce counter as an argument on every call. Datagram transports need exactly this:
+/// packets may be lost or reordered, so the per-record counter is carried on the wire and supplied
+/// by the caller rather than tracked implicitly. The nonce scheme is identical to [`AeadSession`]
+/// (`LE(counter) || 0x0000_0000`), so the two interoperate on the same key as long as their
+/// counter ranges do not overlap.
+pub struct AeadKey {
+    key: [u8; 32],
+}
+
+impl AeadKey {
+    /// Create a key holder from 256 bits of key material.
+    #[must_use]
+    pub fn new(key: [u8; 32]) -> Self {
+        Self { key }
+    }
+
+    /// Build the cipher instance for the current key.
+    fn cipher(&self) -> ChaCha20Poly1305 {
+        ChaCha20Poly1305::new(Key::from_slice(&self.key))
+    }
+
+    /// Encrypt `plaintext` with associated data `aad` under the nonce derived from `counter`,
+    /// returning `ciphertext || tag`. The caller owns nonce uniqueness: never reuse a `counter`
+    /// value with the same key.
+    ///
+    /// # Panics
+    /// Panics only if the underlying AEAD reports an error, which for ChaCha20-Poly1305 encryption
+    /// happens solely when the plaintext exceeds the cipher's maximum supported length.
+    #[must_use]
+    pub fn seal(&self, counter: u64, plaintext: &[u8], aad: &[u8]) -> Vec<u8> {
+        let nonce = AeadSession::nonce_for(counter);
+        self.cipher()
+            .encrypt(
+                Nonce::from_slice(&nonce),
+                Payload {
+                    msg: plaintext,
+                    aad,
+                },
+            )
+            .expect("ChaCha20-Poly1305 encryption never fails for in-range plaintext")
+    }
+
+    /// Decrypt `ciphertext` (`ciphertext || tag`) with associated data `aad` under the nonce
+    /// derived from `counter`.
+    ///
+    /// # Errors
+    /// Returns [`CryptoError::AeadDecrypt`] if authentication fails (tampered ciphertext, wrong
+    /// AAD, wrong key, or wrong counter).
+    pub fn open(&self, counter: u64, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
+        let nonce = AeadSession::nonce_for(counter);
+        self.cipher()
+            .decrypt(
+                Nonce::from_slice(&nonce),
+                Payload {
+                    msg: ciphertext,
+                    aad,
+                },
+            )
+            .map_err(|_| CryptoError::AeadDecrypt)
+    }
+}
+
+impl Drop for AeadKey {
+    fn drop(&mut self) {
+        self.key.zeroize();
+    }
+}
+
+impl zeroize::ZeroizeOnDrop for AeadKey {}
+
 #[cfg(test)]
 mod tests {
    use super::*;
@@ -155,4 +236,50 @@ mod tests {
        }
        assert_eq!(seen.len(), 10_000);
    }
+
+    #[test]
+    fn aead_key_explicit_nonce_roundtrip() {
+        let k = AeadKey::new([7u8; 32]);
+        let ct = k.seal(42, b"hello datagram", b"aad");
+        let pt = k.open(42, &ct, b"aad").expect("open at same counter");
+        assert_eq!(pt, b"hello datagram");
+    }
+
+    #[test]
+    fn aead_key_wrong_counter_or_aad_fails() {
+        let k = AeadKey::new([9u8; 32]);
+        let ct = k.seal(5, b"msg", b"aad");
+        assert!(k.open(6, &ct, b"aad").is_err(), "wrong counter must fail");
+        assert!(k.open(5, &ct, b"other").is_err(), "wrong aad must fail");
+        let mut tampered = ct.clone();
+        tampered[0] ^= 1;
+        assert!(k.open(5, &tampered, b"aad").is_err(), "tamper must fail");
+    }
+
+    #[test]
+    fn aead_key_matches_session_nonce_scheme() {
+        // An AeadKey at counter c must reproduce what an AeadSession produces at counter c.
+        let key = [3u8; 32];
+        let mut s = AeadSession::new(key);
+        let from_session = s.seal(b"x", b"a"); // counter 0, then advances
+        let from_key = AeadKey::new(key).seal(0, b"x", b"a");
+        assert_eq!(from_session, from_key);
+    }
+
+    #[test]
+    fn into_parts_preserves_key_and_counter() {
+        let mut s = AeadSession::new([5u8; 32]);
+        let _ = s.seal(b"a", b"");
+        let _ = s.seal(b"b", b""); // counter now 2
+        let (key, counter) = s.into_parts();
+        assert_eq!(counter, 2);
+        // The recovered key, used at the next counter, matches a fresh session advanced to 2.
+        let expect = {
+            let mut s2 = AeadSession::new([5u8; 32]);
+            let _ = s2.seal(b"a", b"");
+            let _ = s2.seal(b"b", b"");
+            s2.seal(b"c", b"d")
+        };
+        assert_eq!(key.seal(counter, b"c", b"d"), expect);
+    }
 }
@@ -16,7 +16,7 @@ pub mod aead;
 pub mod kdf;
 pub mod kem;

-pub use aead::AeadSession;
+pub use aead::{AeadKey, AeadSession};
 pub use kdf::{derive_session_keys, SessionKeys};
 pub use kem::{HybridCiphertext, HybridPrivateKey, HybridPublicKey, HybridSharedSecret};

@@ -49,7 +49,9 @@ pub mod session;
 pub use conn::PacketConnection;
 pub use frame::{Frame, MsgType};
 pub use handshake::{client_handshake, server_handshake};
-pub use session::{Session, SessionReceiver, SessionSender};
+pub use session::{
+    DatagramReceiver, DatagramSender, Session, SessionReceiver, SessionSender,
+};

 use thiserror::Error;

@@ -3,8 +3,8 @@
 //! A [`Session`] owns the transport reader + writer and the two directional [`AeadSession`]s
 //! produced by the handshake. It exposes [`Session::send_frame`] / [`Session::recv_frame`], which
 //! serialize a [`Frame`], AEAD-seal/open it, and ship it inside a [`MsgType::Data`] record framed
-//! by the 5-byte protocol header. For full-duplex use (e.g. a VPN data path) call
-//! [`Session::split`] to get independent [`SessionSender`] / [`SessionReceiver`] halves.
+//! by the 5-byte protocol header. For full-duplex stream use call [`Session::split`]; for a
+//! connectionless (UDP) data path call [`Session::into_datagram_parts`].
 //!
 //! ## Record format and replay protection
 //!
@@ -23,8 +23,14 @@
 //!
 //! The `seq` is also folded into the AEAD AAD (alongside the frame header), cryptographically
 //! binding the record to its claimed position.
+//!
+//! ## Datagram mode
+//!
+//! [`Session::into_datagram_parts`] yields [`DatagramSender`] / [`DatagramReceiver`] for UDP-style
+//! transports: they use `aura_crypto::AeadKey` (explicit per-record nonce = the carried `seq`), so
+//! datagrams can be lost or reordered. The datagram record is `seq(8) || AEAD(frame, aad = seq)`.

-use aura_crypto::AeadSession;
+use aura_crypto::{AeadKey, AeadSession};

 use crate::frame::{self, Frame, MsgType, RawFrame, HEADER_LEN};
 use crate::ProtoError;
@@ -236,10 +242,9 @@ where
 /// An established, encrypted Aura session over a transport reader `R` and writer `W`.
 ///
 /// Created by [`crate::client_handshake`] / [`crate::server_handshake`]. Use
-/// [`Session::send_frame`] / [`Session::recv_frame`] for half-duplex convenience, or
-/// [`Session::split`] to obtain independent [`SessionSender`] / [`SessionReceiver`] halves that can
-/// be moved into separate tasks for full-duplex operation (e.g. a VPN data path with concurrent
-/// read and write).
+/// [`Session::send_frame`] / [`Session::recv_frame`] for half-duplex convenience, [`Session::split`]
+/// for full-duplex stream halves, or [`Session::into_datagram_parts`] for a connectionless (UDP)
+/// data path.
 pub struct Session<R, W> {
    sender: SessionSender<W>,
    receiver: SessionReceiver<R>,
@@ -303,7 +308,7 @@ where
        self.receiver.recv_frame().await
    }

-    /// Split into independent send and receive halves for full-duplex operation.
+    /// Split into independent send and receive halves for full-duplex stream operation.
    ///
    /// The two halves own disjoint state (writer + outbound AEAD vs. reader + inbound AEAD +
    /// replay window), so they can be moved into separate tasks and driven concurrently. Capture
@@ -313,6 +318,31 @@ where
        (self.sender, self.receiver)
    }

+    /// Split into datagram send/receive halves for a connectionless (UDP) data path.
+    ///
+    /// Unlike [`Session::split`] (stream halves bound to the reader/writer), this discards the
+    /// transport stream and returns explicit-nonce AEAD codecs: each datagram carries its own
+    /// sequence number as the nonce, so packets can be lost or reordered. The codecs continue from
+    /// the post-handshake AEAD counters, so they never reuse a nonce already used by the encrypted
+    /// handshake messages. Returns `(sender, receiver, peer_id)`.
+    #[must_use]
+    pub fn into_datagram_parts(self) -> (DatagramSender, DatagramReceiver, Option<String>) {
+        let peer = self.peer_id;
+        let (send_key, send_ctr) = self.sender.send_aead.into_parts();
+        let (recv_key, recv_ctr) = self.receiver.recv_aead.into_parts();
+        (
+            DatagramSender {
+                key: send_key,
+                seq: send_ctr,
+            },
+            DatagramReceiver {
+                key: recv_key,
+                replay: ReplayWindow::new(recv_ctr),
+            },
+            peer,
+        )
+    }
+
    /// Consume the session, returning its transport halves (for clean shutdown / reuse).
    #[must_use]
    pub fn into_inner(self) -> (R, W) {
@@ -322,6 +352,73 @@ where
    }
 }

+/// Datagram (connectionless) send half: an explicit-nonce AEAD plus the next sequence number.
+///
+/// Produced by [`Session::into_datagram_parts`]. Each [`DatagramSender::seal`] returns a complete
+/// datagram payload `seq(8, big-endian) || ChaCha20Poly1305_seal(frame, aad = seq)`. The transport
+/// (e.g. the UDP backend) sends one such payload per UDP datagram.
+pub struct DatagramSender {
+    key: AeadKey,
+    seq: u64,
+}
+
+impl DatagramSender {
+    /// Serialize and seal one [`Frame`] into a datagram payload, advancing the sequence number.
+    #[must_use]
+    pub fn seal(&mut self, frame: &Frame) -> Vec<u8> {
+        let seq = self.seq;
+        let frame_bytes = frame.encode();
+        let seq_be = seq.to_be_bytes();
+        let ciphertext = self.key.seal(seq, &frame_bytes, &seq_be);
+        self.seq = self.seq.checked_add(1).expect("datagram sequence overflow");
+        let mut out = Vec::with_capacity(SEQ_LEN + ciphertext.len());
+        out.extend_from_slice(&seq_be);
+        out.extend_from_slice(&ciphertext);
+        out
+    }
+
+    /// The sequence number that the next [`DatagramSender::seal`] will stamp.
+    #[must_use]
+    pub fn next_seq(&self) -> u64 {
+        self.seq
+    }
+}
+
+/// Datagram (connectionless) receive half: an explicit-nonce AEAD plus a replay window.
+///
+/// Produced by [`Session::into_datagram_parts`]. [`DatagramReceiver::open`] parses the carried
+/// sequence number, rejects replays/too-old records via the sliding window, then authenticates and
+/// decodes the [`Frame`].
+pub struct DatagramReceiver {
+    key: AeadKey,
+    replay: ReplayWindow,
+}
+
+impl DatagramReceiver {
+    /// Authenticate, replay-check, and decode one datagram payload produced by a peer
+    /// [`DatagramSender`].
+    ///
+    /// # Errors
+    /// * [`ProtoError::Replay`] — duplicate or too-old sequence number.
+    /// * [`ProtoError::Crypto`] — AEAD authentication failed.
+    /// * [`ProtoError::MalformedFrame`] — datagram too short or undecodable frame.
+    pub fn open(&mut self, datagram: &[u8]) -> Result<Frame, ProtoError> {
+        if datagram.len() < SEQ_LEN {
+            return Err(ProtoError::MalformedFrame("datagram shorter than seq prefix"));
+        }
+        let mut seq_be = [0u8; SEQ_LEN];
+        seq_be.copy_from_slice(&datagram[..SEQ_LEN]);
+        let seq = u64::from_be_bytes(seq_be);
+        let ciphertext = &datagram[SEQ_LEN..];
+
+        // Replay check FIRST — a duplicate/old record must not be processed.
+        self.replay.check_and_set(seq)?;
+
+        let plaintext = self.key.open(seq, ciphertext, &seq_be)?;
+        Frame::decode(&plaintext)
+    }
+}
+
 #[cfg(test)]
 mod tests {
    use super::*;
@@ -367,4 +464,45 @@ mod tests {
        // Just inside the window edge => still acceptable.
        assert!(w.check_and_set(200 - REPLAY_WINDOW + 1).is_ok());
    }
+
+    #[test]
+    fn datagram_roundtrip_reorder_and_replay() {
+        let key = [11u8; 32];
+        let mut tx = DatagramSender {
+            key: AeadKey::new(key),
+            seq: 2,
+        };
+        let mut rx = DatagramReceiver {
+            key: AeadKey::new(key),
+            replay: ReplayWindow::new(2),
+        };
+
+        let d0 = tx.seal(&Frame::Data {
+            stream_id: 0,
+            payload: bytes::Bytes::from_static(b"pkt-a"),
+        });
+        let d1 = tx.seal(&Frame::Data {
+            stream_id: 0,
+            payload: bytes::Bytes::from_static(b"pkt-b"),
+        });
+
+        // Out-of-order delivery within the window is accepted.
+        match rx.open(&d1).expect("open d1") {
+            Frame::Data { payload, .. } => assert_eq!(&payload[..], b"pkt-b"),
+            _ => panic!("expected Data frame"),
+        }
+        match rx.open(&d0).expect("open d0") {
+            Frame::Data { payload, .. } => assert_eq!(&payload[..], b"pkt-a"),
+            _ => panic!("expected Data frame"),
+        }
+
+        // A replay of an already-accepted datagram is rejected.
+        assert!(matches!(rx.open(&d1), Err(ProtoError::Replay(_))));
+
+        // Tampered ciphertext is rejected.
+        let mut bad = tx.seal(&Frame::Ping { seq: 7 });
+        let last = bad.len() - 1;
+        bad[last] ^= 1;
+        assert!(rx.open(&bad).is_err());
+    }
 }