Skip to main content

yip_crypto/
lib.rs

1//! Noise-IKpsk2 post-quantum hybrid handshake and AEAD session crypto for the yip data plane, built
2//! on the `snow` Noise Protocol Framework. Establishing a [`Session`] requires
3//! completing an IKpsk2 [`Handshake`]; the session then seals/opens inner frames
4//! with explicit per-frame nonces and a sliding anti-replay window.
5//!
6//! `snow` drives the Noise handshake only. Once the handshake completes, the two
7//! secret transport keys are extracted via snow's `dangerously_get_raw_split()`
8//! (the same HKDF-derived bytes snow's own `split()` uses internally) and handed
9//! to `ring`'s asm ChaCha20-Poly1305 for the data-plane hot path, keyed with the
10//! Noise nonce convention (4 zero bytes ++ 8-byte little-endian counter). This is
11//! byte-identical to snow's own transport state but faster.
12#![forbid(unsafe_code)]
13
14use ml_kem::{
15    array::Array,
16    kem::{Decapsulate, Encapsulate, Kem, KeyExport},
17    ml_kem_512::{self, MlKem512},
18};
19use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
20use zeroize::Zeroizing;
21
22/// The Noise parameter set: IKpsk2 pattern, X25519, ChaCha20-Poly1305, BLAKE2s.
23pub(crate) const NOISE_PARAMS: &str = "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s";
24
25/// An X25519 static keypair (32-byte private and public halves).
26#[derive(Debug, Clone)]
27pub struct Keypair {
28    /// X25519 private key.
29    pub private: [u8; 32],
30    /// X25519 public key.
31    pub public: [u8; 32],
32}
33
34/// Generate a fresh X25519 static keypair.
35pub fn generate_keypair() -> Keypair {
36    let kp = snow::Builder::new(NOISE_PARAMS.parse().expect("valid params"))
37        .generate_keypair()
38        .expect("keypair generation");
39    let mut private = [0u8; 32];
40    let mut public = [0u8; 32];
41    private.copy_from_slice(&kp.private);
42    public.copy_from_slice(&kp.public);
43    Keypair { private, public }
44}
45
46/// ML-KEM-512 encapsulation key length in bytes (FIPS 203).
47pub const PQ_ENCAPS_KEY_LEN: usize = 800;
48
49/// ML-KEM-512 ciphertext length in bytes (FIPS 203).
50pub const PQ_CIPHERTEXT_LEN: usize = 768;
51
52/// ML-KEM-512 shared secret length in bytes (FIPS 203).
53pub const PQ_SHARED_SECRET_LEN: usize = 32;
54
55/// An ML-KEM-512 public encapsulation key.
56#[derive(Clone, Debug, PartialEq, Eq)]
57pub struct PqEncapsKey(pub [u8; PQ_ENCAPS_KEY_LEN]);
58
59/// An ML-KEM-512 private decapsulation key wrapped in zeroizing memory.
60pub struct PqDecapsKey(Box<ml_kem_512::DecapsulationKey>);
61
62impl std::fmt::Debug for PqDecapsKey {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        f.write_str("PqDecapsKey([REDACTED])")
65    }
66}
67
68/// An ML-KEM-512 encapsulated ciphertext.
69#[derive(Clone, Debug, PartialEq, Eq)]
70pub struct PqCiphertext(pub [u8; PQ_CIPHERTEXT_LEN]);
71
72/// Generate a fresh ML-KEM-512 keypair.
73pub fn generate_pq_keypair() -> (PqEncapsKey, PqDecapsKey) {
74    let (dk, ek) = MlKem512::generate_keypair();
75    let ek_bytes = ek.to_bytes();
76    let mut ek_arr = [0u8; PQ_ENCAPS_KEY_LEN];
77    ek_arr.copy_from_slice(ek_bytes.as_slice());
78    (PqEncapsKey(ek_arr), PqDecapsKey(Box::new(dk)))
79}
80
81/// Encapsulate a shared secret toward `ek`, producing `(ciphertext, shared_secret)`.
82///
83/// Explicitly validates FIPS 203 Section 7.2 ByteDecode_12 (all 12-bit coefficients < 3329)
84/// across the first 768 bytes before invoking key encapsulation.
85pub fn pq_encapsulate(
86    ek: &PqEncapsKey,
87) -> Result<(PqCiphertext, Zeroizing<[u8; PQ_SHARED_SECRET_LEN]>), CryptoError> {
88    // FIPS 203 Section 7.2: ByteDecode_12 modulus check (all 12-bit coeffs < 3329)
89    for chunk in ek.0[..768].as_chunks::<3>().0 {
90        let d0 = u16::from(chunk[0]) | ((u16::from(chunk[1]) & 0x0F) << 8);
91        let d1 = (u16::from(chunk[1]) >> 4) | (u16::from(chunk[2]) << 4);
92        if d0 >= 3329 || d1 >= 3329 {
93            return Err(CryptoError::Handshake);
94        }
95    }
96
97    let ek_arr = Array::from(ek.0);
98    let key = ml_kem_512::EncapsulationKey::new(&ek_arr).map_err(|_| CryptoError::Handshake)?;
99    let (ct, ss) = key.encapsulate();
100    let ss = Zeroizing::new(ss);
101    let mut ct_arr = [0u8; PQ_CIPHERTEXT_LEN];
102    ct_arr.copy_from_slice(ct.as_slice());
103    let mut ss_arr = Zeroizing::new([0u8; PQ_SHARED_SECRET_LEN]);
104    ss_arr.copy_from_slice(ss.as_slice());
105    Ok((PqCiphertext(ct_arr), ss_arr))
106}
107
108/// Decapsulate `ct` using `dk`, recovering the 32-byte shared secret wrapped in zeroizing memory.
109pub fn pq_decapsulate(
110    dk: &PqDecapsKey,
111    ct: &PqCiphertext,
112) -> Result<Zeroizing<[u8; PQ_SHARED_SECRET_LEN]>, CryptoError> {
113    let ct_arr = Array::from(ct.0);
114    let ss = Zeroizing::new(dk.0.decapsulate(&ct_arr));
115    let mut ss_arr = Zeroizing::new([0u8; PQ_SHARED_SECRET_LEN]);
116    ss_arr.copy_from_slice(ss.as_slice());
117    Ok(ss_arr)
118}
119
120/// Standard profile replay window parameters (1 KB bitmap).
121pub const REPLAY_WINDOW_BITS_STANDARD: u64 = 8_192;
122pub const REPLAY_WORDS_STANDARD: usize = 128;
123
124/// HighThroughput profile replay window parameters (16 KB bitmap).
125pub const REPLAY_WINDOW_BITS_HIGH_THROUGHPUT: u64 = 131_072;
126pub const REPLAY_WORDS_HIGH_THROUGHPUT: usize = 2048;
127
128/// Number of past counters the replay window tracks behind the latest in HighThroughput mode.
129pub const REPLAY_WINDOW_BITS: u64 = REPLAY_WINDOW_BITS_HIGH_THROUGHPUT;
130pub const REPLAY_WORDS: usize = REPLAY_WORDS_HIGH_THROUGHPUT;
131
132/// Sizing profile for the sliding replay window bitmap.
133#[derive(Debug, Clone, Copy, PartialEq, Eq)]
134pub enum ReplayProfile {
135    /// 1 KB bitmap tracking 8,192 bits behind the latest seen counter.
136    /// Fits tightly into L1/L2 cache for standard latency-critical connections.
137    Standard,
138    /// 16 KB bitmap tracking 131,072 bits behind the latest seen counter.
139    /// Designed for multi-gigabit multi-core pipelines absorbing high burst jitter.
140    HighThroughput,
141}
142
143impl ReplayProfile {
144    /// Number of past counters tracked behind the latest counter.
145    #[inline]
146    pub const fn bits(self) -> u64 {
147        match self {
148            Self::Standard => REPLAY_WINDOW_BITS_STANDARD,
149            Self::HighThroughput => REPLAY_WINDOW_BITS_HIGH_THROUGHPUT,
150        }
151    }
152
153    /// Number of 64-bit words in the circular ring bitmap.
154    #[inline]
155    pub const fn words(self) -> usize {
156        match self {
157            Self::Standard => REPLAY_WORDS_STANDARD,
158            Self::HighThroughput => REPLAY_WORDS_HIGH_THROUGHPUT,
159        }
160    }
161}
162
163/// A wide sliding replay window over a monotonic `u64` counter using a circular word ring.
164#[derive(Clone)]
165pub struct ReplayWindow {
166    profile: ReplayProfile,
167    latest: u64,
168    bitmap: Box<[u64]>,
169    started: bool,
170}
171
172impl ReplayWindow {
173    /// Create a replay window with the default `HighThroughput` profile (16 KB / 131,072 bits).
174    pub fn new() -> Self {
175        Self::new_with_profile(ReplayProfile::HighThroughput)
176    }
177
178    /// Create a replay window configured with a specific sizing profile.
179    pub fn new_with_profile(profile: ReplayProfile) -> Self {
180        let words = profile.words();
181        Self {
182            profile,
183            latest: 0,
184            bitmap: vec![0u64; words].into_boxed_slice(),
185            started: false,
186        }
187    }
188
189    /// Current sizing profile of the replay window.
190    pub fn profile(&self) -> ReplayProfile {
191        self.profile
192    }
193
194    /// Dynamically promote window capacity to `HighThroughput` (16 KB / 131,072 bits).
195    /// Preserves all previously seen counters without dropping replay protection.
196    pub fn promote_to_high_throughput(&mut self) {
197        if self.profile == ReplayProfile::HighThroughput {
198            return;
199        }
200        let old_words = self.bitmap.len();
201        let new_words = ReplayProfile::HighThroughput.words();
202        let mut new_bitmap = vec![0u64; new_words].into_boxed_slice();
203
204        if self.started {
205            let latest_word = self.latest / 64;
206            let start_word = latest_word.saturating_sub((old_words - 1) as u64);
207            for w in start_word..=latest_word {
208                let old_idx = (w as usize) & (old_words - 1);
209                let new_idx = (w as usize) & (new_words - 1);
210                new_bitmap[new_idx] = self.bitmap[old_idx];
211            }
212        }
213
214        self.profile = ReplayProfile::HighThroughput;
215        self.bitmap = new_bitmap;
216    }
217
218    #[inline]
219    fn word_idx(&self, counter: u64) -> usize {
220        ((counter / 64) as usize) & (self.bitmap.len() - 1)
221    }
222
223    #[inline]
224    fn bit_mask(counter: u64) -> u64 {
225        1u64 << (counter % 64)
226    }
227
228    /// Would `counter` be accepted right now? Read-only — does not mutate state.
229    pub fn check(&self, counter: u64) -> bool {
230        if !self.started {
231            return true;
232        }
233        if counter > self.latest {
234            true
235        } else {
236            let diff = self.latest - counter;
237            if diff >= self.profile.bits() {
238                return false; // too old
239            }
240            let idx = self.word_idx(counter);
241            let mask = Self::bit_mask(counter);
242            (self.bitmap[idx] & mask) == 0
243        }
244    }
245
246    /// Record `counter` as seen, advancing the window. Must be preceded by check().
247    pub fn commit(&mut self, counter: u64) {
248        if !self.started {
249            self.started = true;
250            self.latest = counter;
251            let idx = self.word_idx(counter);
252            self.bitmap[idx] = Self::bit_mask(counter);
253            return;
254        }
255
256        if counter > self.latest {
257            let diff = counter - self.latest;
258            let window_bits = self.profile.bits();
259            let total_words = self.bitmap.len();
260            if diff >= window_bits {
261                // Large leap: clear entire bitmap
262                self.bitmap.fill(0);
263            } else {
264                // Clear any words in the circular ring that were overtaken
265                let old_word = self.latest / 64;
266                let new_word = counter / 64;
267                if new_word > old_word {
268                    let words_to_clear = ((new_word - old_word) as usize).min(total_words);
269                    for w in 1..=words_to_clear {
270                        let idx = ((old_word + w as u64) as usize) & (total_words - 1);
271                        self.bitmap[idx] = 0;
272                    }
273                }
274            }
275            self.latest = counter;
276            let idx = self.word_idx(counter);
277            self.bitmap[idx] |= Self::bit_mask(counter);
278        } else {
279            let diff = self.latest - counter;
280            if diff < self.profile.bits() {
281                let idx = self.word_idx(counter);
282                self.bitmap[idx] |= Self::bit_mask(counter);
283            }
284        }
285    }
286
287    /// Convenience helper to check and commit a counter in one step.
288    /// Used by test harnesses and benchmark suites.
289    pub fn check_and_set(&mut self, counter: u64) -> bool {
290        if self.check(counter) {
291            self.commit(counter);
292            true
293        } else {
294            false
295        }
296    }
297}
298
299impl Default for ReplayWindow {
300    fn default() -> Self {
301        Self::new()
302    }
303}
304
305/// Errors from the crypto layer.
306
307#[derive(Debug, thiserror::Error, PartialEq, Eq)]
308pub enum CryptoError {
309    /// AEAD tag did not verify / decryption failed.
310    #[error("decryption failed")]
311    Decrypt,
312    /// Nonce/counter outside the anti-replay window (replayed or too old).
313    #[error("replayed message")]
314    Replay,
315    /// Handshake step failed (bad message, wrong state, or key error).
316    #[error("handshake failed")]
317    Handshake,
318}
319
320/// Internal role-specific state for post-quantum hybrid key exchange.
321enum HandshakePqRole {
322    Initiator {
323        dk: Option<PqDecapsKey>,
324        ek: PqEncapsKey,
325    },
326    Responder {
327        pending_ct: Option<PqCiphertext>,
328    },
329}
330
331/// An in-progress Noise-IKpsk2 post-quantum hybrid handshake.
332pub struct Handshake {
333    inner: std::sync::Mutex<snow::HandshakeState>,
334    role: HandshakePqRole,
335}
336
337impl Handshake {
338    /// Begin as the initiator, which must already know the responder's static public key.
339    pub fn initiator(
340        local_private: &[u8; 32],
341        peer_public: &[u8; 32],
342    ) -> Result<Handshake, CryptoError> {
343        let inner = snow::Builder::new(NOISE_PARAMS.parse().map_err(|_| CryptoError::Handshake)?)
344            .local_private_key(local_private)
345            .map_err(|_| CryptoError::Handshake)?
346            .remote_public_key(peer_public)
347            .map_err(|_| CryptoError::Handshake)?
348            .build_initiator()
349            .map_err(|_| CryptoError::Handshake)?;
350
351        let (ek, dk) = generate_pq_keypair();
352        Ok(Handshake {
353            inner: std::sync::Mutex::new(inner),
354            role: HandshakePqRole::Initiator { dk: Some(dk), ek },
355        })
356    }
357
358    /// Begin as the responder; learns the initiator's static key during the handshake.
359    pub fn responder(local_private: &[u8; 32]) -> Result<Handshake, CryptoError> {
360        let inner = snow::Builder::new(NOISE_PARAMS.parse().map_err(|_| CryptoError::Handshake)?)
361            .local_private_key(local_private)
362            .map_err(|_| CryptoError::Handshake)?
363            .build_responder()
364            .map_err(|_| CryptoError::Handshake)?;
365
366        Ok(Handshake {
367            inner: std::sync::Mutex::new(inner),
368            role: HandshakePqRole::Responder { pending_ct: None },
369        })
370    }
371
372    /// Produce the next handshake message to send to the peer.
373    ///
374    /// - For Initiator (Msg 1): Encrypts `[ek_A (800B) || payload]` inside Noise Msg 1.
375    /// - For Responder (Msg 2): Encrypts `payload` inside Noise Msg 2, framed as `[ct (768B) || Noise Msg 2]`.
376    pub fn write_message(&mut self, payload: &[u8]) -> Result<Vec<u8>, CryptoError> {
377        let mut buf = [0u8; 4096];
378        match &mut self.role {
379            HandshakePqRole::Initiator { ek, .. } => {
380                let mut full_payload = Vec::with_capacity(PQ_ENCAPS_KEY_LEN + payload.len());
381                full_payload.extend_from_slice(&ek.0);
382                full_payload.extend_from_slice(payload);
383                let n = self
384                    .inner
385                    .get_mut()
386                    .expect("handshake lock")
387                    .write_message(&full_payload, &mut buf)
388                    .map_err(|_| CryptoError::Handshake)?;
389                Ok(buf[..n].to_vec())
390            }
391            HandshakePqRole::Responder { pending_ct } => {
392                let ct = pending_ct.take().ok_or(CryptoError::Handshake)?;
393                let n = self
394                    .inner
395                    .get_mut()
396                    .expect("handshake lock")
397                    .write_message(payload, &mut buf)
398                    .map_err(|_| CryptoError::Handshake)?;
399                let mut out = Vec::with_capacity(PQ_CIPHERTEXT_LEN + n);
400                out.extend_from_slice(&ct.0);
401                out.extend_from_slice(&buf[..n]);
402                Ok(out)
403            }
404        }
405    }
406
407    /// Consume a handshake message received from the peer, returning the decrypted app payload.
408    ///
409    /// - For Responder (Msg 1): Decrypts Noise Msg 1, extracts `ek_A`, encapsulates `(ct, ss)`,
410    ///   calls `set_psk(2, &ss)`, and returns `payload`.
411    /// - For Initiator (Msg 2): Slices `ct` from Msg 2 prefix, decapsulates `ss`,
412    ///   calls `set_psk(2, &ss)`, decrypts Noise Msg 2, scrubs `dk`, and returns `payload`.
413    pub fn read_message(&mut self, msg: &[u8]) -> Result<Vec<u8>, CryptoError> {
414        let mut buf = [0u8; 4096];
415        match &mut self.role {
416            HandshakePqRole::Responder { pending_ct } => {
417                let n = self
418                    .inner
419                    .get_mut()
420                    .expect("handshake lock")
421                    .read_message(msg, &mut buf)
422                    .map_err(|_| CryptoError::Handshake)?;
423                if n < PQ_ENCAPS_KEY_LEN {
424                    return Err(CryptoError::Handshake);
425                }
426                let ek_bytes: [u8; PQ_ENCAPS_KEY_LEN] = buf[..PQ_ENCAPS_KEY_LEN]
427                    .try_into()
428                    .map_err(|_| CryptoError::Handshake)?;
429                let ek = PqEncapsKey(ek_bytes);
430                let (ct, ss) = pq_encapsulate(&ek)?;
431                self.inner
432                    .get_mut()
433                    .expect("handshake lock")
434                    .set_psk(2, &ss[..])
435                    .map_err(|_| CryptoError::Handshake)?;
436                *pending_ct = Some(ct);
437                Ok(buf[PQ_ENCAPS_KEY_LEN..n].to_vec())
438            }
439            HandshakePqRole::Initiator { dk, .. } => {
440                let ct_slice = msg.get(..PQ_CIPHERTEXT_LEN).ok_or(CryptoError::Handshake)?;
441                let ct_bytes: [u8; PQ_CIPHERTEXT_LEN] =
442                    ct_slice.try_into().map_err(|_| CryptoError::Handshake)?;
443                let ct = PqCiphertext(ct_bytes);
444                let decaps_key = dk.take().ok_or(CryptoError::Handshake)?;
445                let ss = pq_decapsulate(&decaps_key, &ct)?;
446                drop(decaps_key); // Explicit scrub of dk_A
447
448                let inner = self.inner.get_mut().expect("handshake lock");
449                inner
450                    .set_psk(2, &ss[..])
451                    .map_err(|_| CryptoError::Handshake)?;
452                let n = inner
453                    .read_message(&msg[PQ_CIPHERTEXT_LEN..], &mut buf)
454                    .map_err(|_| CryptoError::Handshake)?;
455                Ok(buf[..n].to_vec())
456            }
457        }
458    }
459
460    /// Whether the handshake has completed and a session can be derived.
461    pub fn is_finished(&self) -> bool {
462        self.inner
463            .lock()
464            .expect("handshake lock")
465            .is_handshake_finished()
466    }
467
468    /// The peer's authenticated static public key, if learned yet.
469    pub fn remote_static(&self) -> Option<[u8; 32]> {
470        self.inner
471            .lock()
472            .expect("handshake lock")
473            .get_remote_static()
474            .map(|k| {
475                let mut out = [0u8; 32];
476                out.copy_from_slice(k);
477                out
478            })
479    }
480
481    /// The Noise channel-binding hash (snow's handshake hash), identical on both
482    /// peers after the handshake completes. Use it to derive subkeys (e.g. the
483    /// wire codec keys) bound to this session.
484    pub fn channel_binding(&self) -> [u8; 32] {
485        let mut out = [0u8; 32];
486        out.copy_from_slice(
487            self.inner
488                .lock()
489                .expect("handshake lock")
490                .get_handshake_hash(),
491        );
492        out
493    }
494
495    /// Extract the two raw 32-byte secret Noise transport keys (send key, recv key)
496    /// derived from the handshake.
497    pub fn raw_split_keys(&self) -> ([u8; 32], [u8; 32]) {
498        let mut inner = self.inner.lock().expect("handshake lock");
499        let is_initiator = inner.is_initiator();
500        let (k0, k1) = inner.dangerously_get_raw_split();
501        if is_initiator {
502            (k0, k1)
503        } else {
504            (k1, k0)
505        }
506    }
507
508    /// Convert a completed handshake into an AEAD [`Session`].
509    ///
510    /// Extracts the two secret Noise transport keys via snow's raw split and
511    /// builds `ring` AEAD keys from them directly; snow's own transport state is
512    /// not used for the data plane. Per snow's split convention, element 0 of
513    /// the pair is the initiator's send key (= responder's receive key) and
514    /// element 1 is the responder's send key (= initiator's receive key), so
515    /// the mapping below is role-dependent.
516    pub fn into_session(self) -> Result<Session, CryptoError> {
517        let (k_send, k_recv) = self.raw_split_keys();
518        Session::from_raw_keys(&k_send, &k_recv, 0, 1)
519    }
520
521    #[cfg(test)]
522    pub(crate) fn to_session_for_test(&self) -> Result<Session, CryptoError> {
523        let (k_send, k_recv) = self.raw_split_keys();
524        Session::from_raw_keys(&k_send, &k_recv, 0, 1)
525    }
526
527    #[cfg(test)]
528    pub(crate) fn into_inner_snow_for_test(self) -> snow::HandshakeState {
529        self.inner.into_inner().expect("handshake lock")
530    }
531}
532
533/// Noise ChaChaPoly nonce: 4 zero bytes ++ 8-byte little-endian counter.
534fn noise_nonce(counter: u64) -> Nonce {
535    let mut n = [0u8; 12];
536    n[4..].copy_from_slice(&counter.to_le_bytes());
537    Nonce::assume_unique_for_key(n)
538}
539
540/// A sealed frame: the AEAD ciphertext plus the explicit nonce it was sealed
541/// under. The caller carries `counter` on the wire so the peer can `open`.
542#[derive(Debug, Clone)]
543pub struct Sealed {
544    /// The explicit AEAD nonce assigned to this frame.
545    pub counter: u64,
546    /// The AEAD ciphertext (plaintext length + 16-byte tag).
547    pub ciphertext: Vec<u8>,
548}
549
550/// An established AEAD session. Seals outgoing frames under a monotonic counter
551/// and opens incoming frames out of order, rejecting replays.
552///
553/// Uses `ring`'s ChaCha20-Poly1305 keyed by the Noise Split() transport keys
554/// (see [`Handshake::into_session`]) rather than snow's own transport state.
555pub struct Session {
556    send_key: LessSafeKey,
557    recv_key: LessSafeKey,
558    send_counter: u64,
559    stride: u64,
560    replay: ReplayWindow,
561}
562
563impl Session {
564    /// Construct a session directly from raw 32-byte send and receive keys,
565    /// with an explicit start counter and stride increment (e.g. for worker sharding).
566    pub fn from_raw_keys(
567        k_send: &[u8; 32],
568        k_recv: &[u8; 32],
569        start_counter: u64,
570        stride: u64,
571    ) -> Result<Self, CryptoError> {
572        let send =
573            UnboundKey::new(&CHACHA20_POLY1305, k_send).map_err(|_| CryptoError::Handshake)?;
574        let recv =
575            UnboundKey::new(&CHACHA20_POLY1305, k_recv).map_err(|_| CryptoError::Handshake)?;
576        Ok(Session {
577            send_key: LessSafeKey::new(send),
578            recv_key: LessSafeKey::new(recv),
579            send_counter: start_counter,
580            stride: if stride == 0 { 1 } else { stride },
581            replay: ReplayWindow::new(),
582        })
583    }
584
585    /// Reconfigure the send counter and stride increment for this session.
586    pub fn set_stride(&mut self, start_counter: u64, stride: u64) {
587        let stride = if stride == 0 { 1 } else { stride };
588        if self.send_counter == 0 {
589            self.send_counter = start_counter;
590        } else {
591            let rem = self.send_counter % stride;
592            let target_rem = start_counter % stride;
593            let diff = if rem <= target_rem {
594                target_rem - rem
595            } else {
596                stride - (rem - target_rem)
597            };
598            self.send_counter = self.send_counter.saturating_add(diff);
599        }
600        self.stride = stride;
601    }
602
603    /// The current nonce stride increment.
604    pub fn stride(&self) -> u64 {
605        self.stride
606    }
607
608    /// The current send counter.
609    pub fn send_counter(&self) -> u64 {
610        self.send_counter
611    }
612
613    /// Seal one inner frame, assigning it the next send counter.
614    pub fn seal(&mut self, plaintext: &[u8]) -> Result<Sealed, CryptoError> {
615        let counter = self.send_counter;
616        let mut buf = plaintext.to_vec();
617        self.send_key
618            .seal_in_place_append_tag(noise_nonce(counter), Aad::empty(), &mut buf)
619            .map_err(|_| CryptoError::Decrypt)?;
620        self.send_counter = self
621            .send_counter
622            .checked_add(self.stride)
623            .ok_or(CryptoError::Decrypt)?;
624        Ok(Sealed {
625            counter,
626            ciphertext: buf,
627        })
628    }
629
630    /// Open one inner frame received under explicit `counter`, enforcing replay protection.
631    ///
632    /// The replay window is checked read-only *before* AEAD, then committed only
633    /// *after* the frame authenticates (matching WireGuard, which advances its
634    /// window post-decrypt). This ordering prevents a forged frame carrying an
635    /// arbitrary counter from advancing the window and starving legitimate
636    /// packets — an off-path DoS the mark-before-auth ordering was open to.
637    pub fn open(&mut self, counter: u64, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
638        if !self.replay.check(counter) {
639            return Err(CryptoError::Replay);
640        }
641        let mut buf = ciphertext.to_vec();
642        let plain = self
643            .recv_key
644            .open_in_place(noise_nonce(counter), Aad::empty(), &mut buf)
645            .map_err(|_| CryptoError::Decrypt)?;
646        self.replay.commit(counter);
647        Ok(plain.to_vec())
648    }
649
650    /// Check whether `counter` is valid within the anti-replay window without committing it.
651    /// Used for pre-decoding packet validation and DoS rejection.
652    #[inline]
653    pub fn check_replay(&self, counter: u64) -> bool {
654        self.replay.check(counter)
655    }
656
657    /// Seal into a caller-owned reusable buffer (no per-call allocation).
658    pub fn seal_into(&mut self, plaintext: &[u8], out: &mut Vec<u8>) -> Result<u64, CryptoError> {
659        let counter = self.send_counter;
660        out.clear();
661        out.extend_from_slice(plaintext);
662        self.send_key
663            .seal_in_place_append_tag(noise_nonce(counter), Aad::empty(), out)
664            .map_err(|_| CryptoError::Decrypt)?;
665        self.send_counter = self
666            .send_counter
667            .checked_add(self.stride)
668            .ok_or(CryptoError::Decrypt)?;
669        Ok(counter)
670    }
671
672    /// Seal one inner frame under an explicit counter, without mutating internal counter.
673    /// This allows multi-threaded workers sharing a single peer connection to seal
674    /// concurrently using nonces dispensed from a ChunkedNonceDispenser.
675    pub fn seal_with_counter(&self, counter: u64, plaintext: &[u8]) -> Result<Sealed, CryptoError> {
676        let mut buf = plaintext.to_vec();
677        self.send_key
678            .seal_in_place_append_tag(noise_nonce(counter), Aad::empty(), &mut buf)
679            .map_err(|_| CryptoError::Decrypt)?;
680        Ok(Sealed {
681            counter,
682            ciphertext: buf,
683        })
684    }
685
686    /// Seal into a caller-owned reusable buffer under an explicit counter.
687    pub fn seal_into_with_counter(
688        &self,
689        counter: u64,
690        plaintext: &[u8],
691        out: &mut Vec<u8>,
692    ) -> Result<(), CryptoError> {
693        out.clear();
694        out.extend_from_slice(plaintext);
695        self.send_key
696            .seal_in_place_append_tag(noise_nonce(counter), Aad::empty(), out)
697            .map_err(|_| CryptoError::Decrypt)?;
698        Ok(())
699    }
700
701    /// Open one inner frame received under explicit counter against a caller-provided replay window.
702    pub fn open_with_window(
703        &self,
704        counter: u64,
705        ciphertext: &[u8],
706        replay: &mut ReplayWindow,
707    ) -> Result<Vec<u8>, CryptoError> {
708        if !replay.check(counter) {
709            return Err(CryptoError::Replay);
710        }
711        let mut buf = ciphertext.to_vec();
712        let plain = self
713            .recv_key
714            .open_in_place(noise_nonce(counter), Aad::empty(), &mut buf)
715            .map_err(|_| CryptoError::Decrypt)?;
716        replay.commit(counter);
717        Ok(plain.to_vec())
718    }
719
720    /// Open into a caller-owned reusable buffer under explicit counter against a caller-provided replay window.
721    pub fn open_into_with_window(
722        &self,
723        counter: u64,
724        ciphertext: &[u8],
725        replay: &mut ReplayWindow,
726        out: &mut Vec<u8>,
727    ) -> Result<(), CryptoError> {
728        if !replay.check(counter) {
729            return Err(CryptoError::Replay);
730        }
731        out.clear();
732        out.extend_from_slice(ciphertext);
733        let n = {
734            let plain = self
735                .recv_key
736                .open_in_place(noise_nonce(counter), Aad::empty(), out)
737                .map_err(|_| CryptoError::Decrypt)?;
738            plain.len()
739        };
740        replay.commit(counter);
741        out.truncate(n);
742        Ok(())
743    }
744
745    /// Open into a caller-owned reusable buffer (no per-call allocation).
746    pub fn open_into(
747        &mut self,
748        counter: u64,
749        ciphertext: &[u8],
750        out: &mut Vec<u8>,
751    ) -> Result<(), CryptoError> {
752        if !self.replay.check(counter) {
753            return Err(CryptoError::Replay);
754        }
755        out.clear();
756        out.extend_from_slice(ciphertext);
757        let n = {
758            let plain = self
759                .recv_key
760                .open_in_place(noise_nonce(counter), Aad::empty(), out)
761                .map_err(|_| CryptoError::Decrypt)?;
762            plain.len()
763        };
764        self.replay.commit(counter);
765        out.truncate(n);
766        Ok(())
767    }
768
769    /// Seal plaintext in-place inside `buf[..plaintext_len]`, writing the 16-byte Poly1305
770    /// authentication tag immediately following the ciphertext into `buf[plaintext_len..plaintext_len + 16]`.
771    ///
772    /// The buffer must have capacity of at least `plaintext_len + 16` bytes.
773    /// Returns the explicit counter assigned to this frame.
774    pub fn seal_in_place(
775        &mut self,
776        buf: &mut [u8],
777        plaintext_len: usize,
778    ) -> Result<u64, CryptoError> {
779        let counter = self.send_counter;
780        let total_len = plaintext_len.checked_add(16).ok_or(CryptoError::Decrypt)?;
781        if buf.len() < total_len {
782            return Err(CryptoError::Decrypt);
783        }
784        let (in_out, tag_out) = buf[..total_len].split_at_mut(plaintext_len);
785        let tag = self
786            .send_key
787            .seal_in_place_separate_tag(noise_nonce(counter), Aad::empty(), in_out)
788            .map_err(|_| CryptoError::Decrypt)?;
789        tag_out.copy_from_slice(tag.as_ref());
790        self.send_counter = self
791            .send_counter
792            .checked_add(self.stride)
793            .ok_or(CryptoError::Decrypt)?;
794        Ok(counter)
795    }
796
797    /// Open and authenticate ciphertext in-place inside `buf[..sealed_len]`, enforcing anti-replay.
798    ///
799    /// Expects the 16-byte Poly1305 authentication tag at the end of the ciphertext:
800    /// `buf[sealed_len - 16..sealed_len]`.
801    /// Returns the decrypted plaintext length (`sealed_len - 16`).
802    pub fn open_in_place(
803        &mut self,
804        counter: u64,
805        buf: &mut [u8],
806        sealed_len: usize,
807    ) -> Result<usize, CryptoError> {
808        if !self.replay.check(counter) {
809            return Err(CryptoError::Replay);
810        }
811        if sealed_len < 16 || buf.len() < sealed_len {
812            return Err(CryptoError::Decrypt);
813        }
814        let plain = self
815            .recv_key
816            .open_in_place(noise_nonce(counter), Aad::empty(), &mut buf[..sealed_len])
817            .map_err(|_| CryptoError::Decrypt)?;
818        self.replay.commit(counter);
819        Ok(plain.len())
820    }
821}
822
823/// ChaCha20-Poly1305 AEAD cipher supporting zero-copy in-place encryption and decryption
824/// directly within pre-allocated network packet buffers (such as AF_XDP UMEM chunks).
825pub struct ChaCha20Poly1305Cipher {
826    key: LessSafeKey,
827}
828
829impl std::fmt::Debug for ChaCha20Poly1305Cipher {
830    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
831        f.debug_struct("ChaCha20Poly1305Cipher")
832            .finish_non_exhaustive()
833    }
834}
835
836impl ChaCha20Poly1305Cipher {
837    /// Creates a new cipher instance from a 32-byte symmetric key.
838    pub fn new(key: [u8; 32]) -> Self {
839        let unbound = UnboundKey::new(&CHACHA20_POLY1305, &key)
840            .expect("32-byte key is valid for ChaCha20-Poly1305");
841        Self {
842            key: LessSafeKey::new(unbound),
843        }
844    }
845
846    /// Seal plaintext in-place inside `buf[..plaintext_len]`, writing the 16-byte Poly1305
847    /// authentication tag immediately following the ciphertext into `buf[plaintext_len..plaintext_len + 16]`.
848    ///
849    /// The buffer must have capacity of at least `plaintext_len + 16` bytes.
850    /// Returns the total sealed ciphertext length (`plaintext_len + 16`).
851    pub fn seal_in_place(
852        &self,
853        counter: u64,
854        buf: &mut [u8],
855        plaintext_len: usize,
856    ) -> Result<usize, CryptoError> {
857        let total_len = plaintext_len.checked_add(16).ok_or(CryptoError::Decrypt)?;
858        if buf.len() < total_len {
859            return Err(CryptoError::Decrypt);
860        }
861        let (in_out, tag_out) = buf[..total_len].split_at_mut(plaintext_len);
862        let tag = self
863            .key
864            .seal_in_place_separate_tag(noise_nonce(counter), Aad::empty(), in_out)
865            .map_err(|_| CryptoError::Decrypt)?;
866        tag_out.copy_from_slice(tag.as_ref());
867        Ok(total_len)
868    }
869
870    /// Open and authenticate ciphertext in-place inside `buf[..sealed_len]`.
871    ///
872    /// Expects the 16-byte Poly1305 authentication tag at the end of the ciphertext:
873    /// `buf[sealed_len - 16..sealed_len]`.
874    /// Returns the decrypted plaintext length (`sealed_len - 16`).
875    pub fn open_in_place(
876        &self,
877        counter: u64,
878        buf: &mut [u8],
879        sealed_len: usize,
880    ) -> Result<usize, CryptoError> {
881        if sealed_len < 16 || buf.len() < sealed_len {
882            return Err(CryptoError::Decrypt);
883        }
884        let plain = self
885            .key
886            .open_in_place(noise_nonce(counter), Aad::empty(), &mut buf[..sealed_len])
887            .map_err(|_| CryptoError::Decrypt)?;
888        Ok(plain.len())
889    }
890
891    /// Open and authenticate ciphertext in-place, verifying against a sliding anti-replay window.
892    ///
893    /// The replay window is checked before AEAD decryption and committed only upon
894    /// successful authentication.
895    pub fn open_in_place_with_window(
896        &self,
897        counter: u64,
898        buf: &mut [u8],
899        sealed_len: usize,
900        replay: &mut ReplayWindow,
901    ) -> Result<usize, CryptoError> {
902        if !replay.check(counter) {
903            return Err(CryptoError::Replay);
904        }
905        let plain_len = self.open_in_place(counter, buf, sealed_len)?;
906        replay.commit(counter);
907        Ok(plain_len)
908    }
909}
910
911/// Test-only helper: drive a full initiator/responder handshake to completion
912/// and return the two established sessions. Mirrors `yip_bench::established_pair`.
913#[cfg(test)]
914pub(crate) fn test_session_pair() -> (Session, Session) {
915    let resp_kp = generate_keypair();
916    let init_kp = generate_keypair();
917    let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
918    let mut res = Handshake::responder(&resp_kp.private).unwrap();
919    let m1 = ini.write_message(&[]).unwrap();
920    let _ = res.read_message(&m1).unwrap();
921    let m2 = res.write_message(&[]).unwrap();
922    let _ = ini.read_message(&m2).unwrap();
923    (ini.into_session().unwrap(), res.into_session().unwrap())
924}
925
926#[cfg(test)]
927mod tests {
928    use super::*;
929
930    #[test]
931    fn generated_keypairs_are_distinct_32_byte_keys() {
932        let a = generate_keypair();
933        let b = generate_keypair();
934        assert_eq!(a.private.len(), 32);
935        assert_eq!(a.public.len(), 32);
936        assert_ne!(a.private, b.private, "two keypairs differ");
937        assert_ne!(a.public, [0u8; 32], "public key is not all-zero");
938    }
939
940    #[test]
941    fn session_seal_and_open_in_place_with_stride() {
942        let (k_send, k_recv) = ([0x11u8; 32], [0x22u8; 32]);
943        let mut s_tx = Session::from_raw_keys(&k_send, &k_recv, 2, 4).unwrap();
944        let mut s_rx = Session::from_raw_keys(&k_recv, &k_send, 0, 1).unwrap();
945
946        let mut buf = vec![0u8; 64];
947        buf[..10].copy_from_slice(b"0123456789");
948
949        let c1 = s_tx.seal_in_place(&mut buf, 10).unwrap();
950        assert_eq!(c1, 2);
951        let plain_len = s_rx.open_in_place(c1, &mut buf, 26).unwrap();
952        assert_eq!(plain_len, 10);
953        assert_eq!(&buf[..10], b"0123456789");
954
955        // Next counter with stride 4
956        buf[..10].copy_from_slice(b"abcdefghij");
957        let c2 = s_tx.seal_in_place(&mut buf, 10).unwrap();
958        assert_eq!(c2, 6);
959        let plain_len2 = s_rx.open_in_place(c2, &mut buf, 26).unwrap();
960        assert_eq!(plain_len2, 10);
961        assert_eq!(&buf[..10], b"abcdefghij");
962
963        // Replay of c2 should fail
964        assert!(s_rx.open_in_place(c2, &mut buf, 26).is_err());
965    }
966
967    #[test]
968    fn replay_window_accepts_fresh_rejects_replays_and_old() {
969        let mut w = ReplayWindow::new();
970        assert!(w.check_and_set(0), "first counter accepted");
971        assert!(!w.check_and_set(0), "exact replay rejected");
972        assert!(w.check_and_set(1), "next in order accepted");
973        assert!(w.check_and_set(5), "jump ahead accepted");
974        assert!(w.check_and_set(3), "in-window out-of-order accepted");
975        assert!(!w.check_and_set(3), "replay of out-of-order rejected");
976        assert!(w.check_and_set(100), "large advance accepted");
977        assert!(
978            !w.check_and_set(5),
979            "counter now far below window rejected as too old"
980        );
981    }
982
983    /// Advancing the window from counter A to B and then replaying A must be
984    /// rejected.  Kills the `shift = counter + latest` mutant on the advance
985    /// path (line 60) and the `bitmap << shift` → `bitmap >> shift` mutant on
986    /// the same path (line 64): both mutations misplace the old-counter bits so
987    /// the replay is no longer detected.
988    #[test]
989    fn replay_window_advance_then_replay_old_counter_rejected() {
990        let mut w = ReplayWindow::new();
991        // Establish counter 10 as the first-ever packet.
992        assert!(w.check_and_set(10), "counter 10 accepted as first");
993        // Advance to counter 15 (shift = 5 in real code; shift = 25 under the
994        // `counter + latest` mutant, and bits shift wrong under `>> shift`).
995        assert!(w.check_and_set(15), "advance to 15 accepted");
996        // Counter 10 is now diff=5 from latest=15; its bit must still be set.
997        assert!(
998            !w.check_and_set(10),
999            "replay of counter 10 after advance to 15 rejected"
1000        );
1001    }
1002
1003    /// After advancing the window to a new latest, an immediate replay of that
1004    /// latest counter must be rejected.  Kills the `bitmap | 1` → `bitmap & 1`
1005    /// and `bitmap | 1` → `bitmap ^ 1` mutants: both can leave bit-0 of the
1006    /// new bitmap unset, making the just-accepted counter replayable.
1007    #[test]
1008    fn replay_window_new_latest_immediately_replayable_rejected() {
1009        let mut w = ReplayWindow::new();
1010        // Build a window where counter 9 is also recorded (bit 1 will be set
1011        // at latest=10), so that when we advance to 11 the shifted bitmap has
1012        // its LSB = 1.  Under `^ 1` that would clear bit-0, leaving latest-11
1013        // unprotected.
1014        assert!(w.check_and_set(10), "first packet at 10");
1015        assert!(w.check_and_set(9), "counter 9 accepted in-order-ish");
1016        // Advance from 10 to 11: shift=1, old bitmap has bit-1 set (counter 9).
1017        // `bitmap << 1` yields a value with LSB = old-bit-1 = 1.
1018        // Under `^ 1` that XORs the LSB back to 0, so the replay check below
1019        // would wrongly accept.
1020        assert!(w.check_and_set(11), "advance to 11 accepted");
1021        assert!(!w.check_and_set(11), "replay of new latest 11 rejected");
1022    }
1023
1024    /// The `diff = self.latest - counter` on the in-window path (line 69) must
1025    /// use subtraction, not addition.  With `diff = latest + counter` the
1026    /// diff is 25 (not 5) for latest=15 counter=10, so bit-5 (which is set)
1027    /// is not checked and the replay slips through.
1028    #[test]
1029    fn replay_window_in_window_replay_rejected_after_advance() {
1030        let mut w = ReplayWindow::new();
1031        assert!(w.check_and_set(10), "first packet");
1032        assert!(w.check_and_set(15), "advance to 15");
1033        // Counter 10 is 5 below the new latest; real diff = 5, mutant diff = 25.
1034        // Both are < 64, but bit-5 is set while bit-25 is not.
1035        assert!(!w.check_and_set(10), "in-window replay at diff=5 rejected");
1036    }
1037
1038    /// A freshly-built responder `Handshake` must not report `is_finished`
1039    /// before any messages have been exchanged.  Kills the mutant that
1040    /// replaces the `is_finished` body with `true` (line 150).
1041    #[test]
1042    fn handshake_not_finished_before_message_exchange() {
1043        let kp = generate_keypair();
1044        let res = Handshake::responder(&kp.private).unwrap();
1045        assert!(
1046            !res.is_finished(),
1047            "responder reports not-finished before any messages"
1048        );
1049        let ini = Handshake::initiator(&kp.private, &kp.public).unwrap();
1050        assert!(
1051            !ini.is_finished(),
1052            "initiator reports not-finished before any messages"
1053        );
1054    }
1055
1056    #[test]
1057    fn handshake_payload_round_trips_and_sessions_match() {
1058        // msg1 carries an app payload (the initiator's cert, in 2c); msg2
1059        // carries a different one (the responder's cert). Noise-IK encrypts
1060        // both (msg1 under `es`, msg2 fully), so this also documents that
1061        // certs never appear in cleartext on the wire.
1062        let resp_kp = generate_keypair();
1063        let init_kp = generate_keypair();
1064        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1065        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1066
1067        let m1 = ini.write_message(b"cert-A").unwrap();
1068        let got_a = res.read_message(&m1).unwrap();
1069        assert_eq!(got_a, b"cert-A");
1070
1071        let m2 = res.write_message(b"cert-B").unwrap();
1072        let got_b = ini.read_message(&m2).unwrap();
1073        assert_eq!(got_b, b"cert-B");
1074
1075        // Both sides derive the same channel binding (proof the payloads
1076        // didn't perturb the handshake transcript) before consuming into a
1077        // transport-mode session, then prove the sessions actually talk.
1078        assert_eq!(ini.channel_binding(), res.channel_binding());
1079        let mut ini_session = ini.into_session().unwrap();
1080        let mut res_session = res.into_session().unwrap();
1081        let sealed = ini_session.seal(b"payload round-trip ok").unwrap();
1082        assert_eq!(
1083            res_session
1084                .open(sealed.counter, &sealed.ciphertext)
1085                .unwrap(),
1086            b"payload round-trip ok"
1087        );
1088    }
1089
1090    #[test]
1091    fn session_seals_and_opens_roundtrip() {
1092        let (mut a, mut b) = test_session_pair();
1093        let s = a.seal(b"inner packet").unwrap();
1094        assert_eq!(s.counter, 0, "first counter is 0");
1095        assert_eq!(b.open(s.counter, &s.ciphertext).unwrap(), b"inner packet");
1096    }
1097
1098    #[test]
1099    fn session_opens_out_of_order() {
1100        let (mut a, mut b) = test_session_pair();
1101        let s0 = a.seal(b"zero").unwrap();
1102        let s1 = a.seal(b"one").unwrap();
1103        assert_eq!(s1.counter, 1);
1104        // deliver 1 before 0
1105        assert_eq!(b.open(s1.counter, &s1.ciphertext).unwrap(), b"one");
1106        assert_eq!(b.open(s0.counter, &s0.ciphertext).unwrap(), b"zero");
1107    }
1108
1109    #[test]
1110    fn session_rejects_replay() {
1111        let (mut a, mut b) = test_session_pair();
1112        let s = a.seal(b"x").unwrap();
1113        assert!(b.open(s.counter, &s.ciphertext).is_ok());
1114        assert_eq!(b.open(s.counter, &s.ciphertext), Err(CryptoError::Replay));
1115    }
1116
1117    #[test]
1118    fn session_rejects_tampered_ciphertext() {
1119        let (mut a, mut b) = test_session_pair();
1120        let s = a.seal(b"y").unwrap();
1121        let mut bad = s.ciphertext.clone();
1122        bad[0] ^= 0x01;
1123        assert_eq!(b.open(s.counter, &bad), Err(CryptoError::Decrypt));
1124    }
1125
1126    /// A forged frame carrying a large counter but garbage ciphertext must fail
1127    /// AEAD *without* advancing the anti-replay window. Otherwise an off-path
1128    /// attacker who injects one such frame slides `latest` far forward, so every
1129    /// subsequent legitimate packet is rejected as "too old" — a session-killing
1130    /// DoS. The window slot must only be committed after AEAD verification.
1131    #[test]
1132    fn forged_frame_does_not_advance_replay_window() {
1133        let (mut a, mut b) = test_session_pair();
1134        let s0 = a.seal(b"zero").unwrap();
1135        let s1 = a.seal(b"one").unwrap();
1136
1137        // Establish the receive window with a legitimate frame.
1138        assert_eq!(b.open(s0.counter, &s0.ciphertext).unwrap(), b"zero");
1139
1140        // Off-path attacker injects a forged frame at a far-future counter.
1141        let garbage = vec![0u8; s0.ciphertext.len()];
1142        assert_eq!(
1143            b.open(1_000_000, &garbage),
1144            Err(CryptoError::Decrypt),
1145            "forged frame fails AEAD"
1146        );
1147
1148        // The forged frame must not have moved the window: the next legitimate
1149        // in-flight frame still opens.
1150        assert_eq!(
1151            b.open(s1.counter, &s1.ciphertext).unwrap(),
1152            b"one",
1153            "legit frame still opens after forged far-future injection"
1154        );
1155    }
1156
1157    /// Same invariant on the alloc-free `open_into` path.
1158    #[test]
1159    fn forged_frame_does_not_advance_replay_window_open_into() {
1160        let (mut a, mut b) = test_session_pair();
1161        let s0 = a.seal(b"zero").unwrap();
1162        let s1 = a.seal(b"one").unwrap();
1163        let mut out = Vec::new();
1164
1165        b.open_into(s0.counter, &s0.ciphertext, &mut out).unwrap();
1166        assert_eq!(out, b"zero");
1167
1168        let garbage = vec![0u8; s0.ciphertext.len()];
1169        assert_eq!(
1170            b.open_into(1_000_000, &garbage, &mut out),
1171            Err(CryptoError::Decrypt),
1172            "forged frame fails AEAD"
1173        );
1174
1175        b.open_into(s1.counter, &s1.ciphertext, &mut out).unwrap();
1176        assert_eq!(
1177            out, b"one",
1178            "legit frame still opens after forged injection"
1179        );
1180    }
1181
1182    #[test]
1183    fn channel_binding_matches_on_both_peers() {
1184        let resp_kp = generate_keypair();
1185        let init_kp = generate_keypair();
1186        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1187        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1188        let m1 = ini.write_message(&[]).unwrap();
1189        let _ = res.read_message(&m1).unwrap();
1190        let m2 = res.write_message(&[]).unwrap();
1191        let _ = ini.read_message(&m2).unwrap();
1192        assert!(ini.is_finished() && res.is_finished());
1193        assert_eq!(
1194            ini.channel_binding(),
1195            res.channel_binding(),
1196            "both peers derive the same binding"
1197        );
1198        assert_ne!(ini.channel_binding(), [0u8; 32]);
1199    }
1200
1201    /// `channel_binding` must be a genuine, transcript-dependent hash — not a
1202    /// constant. Kills the `replace channel_binding -> [u8; 32] with [1; 32]`
1203    /// mutant (line 207): two independently-keyed handshakes must derive
1204    /// DIFFERENT bindings, and neither may equal the all-ones constant.
1205    #[test]
1206    fn channel_binding_is_transcript_dependent_not_constant() {
1207        fn complete_binding() -> [u8; 32] {
1208            let resp_kp = generate_keypair();
1209            let init_kp = generate_keypair();
1210            let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1211            let mut res = Handshake::responder(&resp_kp.private).unwrap();
1212            let m1 = ini.write_message(&[]).unwrap();
1213            let _ = res.read_message(&m1).unwrap();
1214            let m2 = res.write_message(&[]).unwrap();
1215            let _ = ini.read_message(&m2).unwrap();
1216            assert!(ini.is_finished() && res.is_finished());
1217            let binding = ini.channel_binding();
1218            assert_eq!(binding, res.channel_binding());
1219            binding
1220        }
1221
1222        let binding_a = complete_binding();
1223        let binding_b = complete_binding();
1224        assert_ne!(
1225            binding_a, [1u8; 32],
1226            "channel_binding must not be the constant [1; 32]"
1227        );
1228        assert_ne!(
1229            binding_b, [1u8; 32],
1230            "channel_binding must not be the constant [1; 32]"
1231        );
1232        assert_ne!(
1233            binding_a, binding_b,
1234            "two independent handshakes (different keys/transcripts) must \
1235             derive different channel bindings"
1236        );
1237    }
1238
1239    #[test]
1240    fn ik_handshake_completes_and_authenticates_initiator() {
1241        let resp_kp = generate_keypair();
1242        let init_kp = generate_keypair();
1243
1244        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1245        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1246
1247        let msg1 = ini.write_message(&[]).unwrap();
1248        let _ = res.read_message(&msg1).unwrap();
1249        let msg2 = res.write_message(&[]).unwrap();
1250        let _ = ini.read_message(&msg2).unwrap();
1251
1252        assert!(ini.is_finished() && res.is_finished());
1253        // IK: the responder learns the initiator's static public key.
1254        assert_eq!(res.remote_static(), Some(init_kp.public));
1255    }
1256
1257    #[test]
1258    fn seal_is_byte_identical_across_a_reference_session() {
1259        // Two independently-built sessions from the same handshake produce the same
1260        // keystream for the same counter+plaintext; a receiver opens what a sender seals.
1261        let (mut a, mut b) = crate::test_session_pair();
1262        for ctr in 0u64..8 {
1263            let s = a.seal(&[0x5Au8; 64]).unwrap();
1264            assert_eq!(s.counter, ctr);
1265            assert_eq!(b.open(s.counter, &s.ciphertext).unwrap(), vec![0x5Au8; 64]);
1266        }
1267    }
1268
1269    #[test]
1270    fn open_rejects_tampered_ciphertext() {
1271        let (mut a, mut b) = crate::test_session_pair();
1272        let s = a.seal(b"secret").unwrap();
1273        let mut bad = s.ciphertext.clone();
1274        bad[0] ^= 1;
1275        assert_eq!(b.open(s.counter, &bad), Err(CryptoError::Decrypt));
1276    }
1277
1278    #[test]
1279    fn open_rejects_replay_and_opens_out_of_order() {
1280        let (mut a, mut b) = crate::test_session_pair();
1281        let s0 = a.seal(b"zero").unwrap();
1282        let s1 = a.seal(b"one").unwrap();
1283        assert_eq!(b.open(s1.counter, &s1.ciphertext).unwrap(), b"one"); // out of order
1284        assert_eq!(b.open(s0.counter, &s0.ciphertext).unwrap(), b"zero");
1285        assert_eq!(b.open(s1.counter, &s1.ciphertext), Err(CryptoError::Replay));
1286        // replay
1287    }
1288
1289    /// Durable KAT (spec §5): the production `Session`'s `ring` ChaCha20-Poly1305
1290    /// output must be byte-for-byte identical to snow's own transport AEAD, for
1291    /// both handshake directions and several counters. This is the regression
1292    /// guard for `Handshake::into_session`'s role-dependent key mapping and for
1293    /// `noise_nonce`'s counter encoding: a swapped send/recv mapping or a
1294    /// big-endian (or otherwise wrong) nonce would still round-trip internally
1295    /// (seal/open use the same buggy convention on both sides) but would no
1296    /// longer match the expected output.
1297    ///
1298    /// Verify that driving completed initiator and responder handshakes yields
1299    /// symmetric raw split keys, and confirm that `Session::seal` produces
1300    /// byte-identical ciphertext to `snow::StatelessTransportState::write_message`
1301    /// across both peers in both directions for several counters.
1302    #[test]
1303    fn session_seal_is_byte_identical_to_snow_write_message_both_directions() {
1304        let resp_kp = generate_keypair();
1305        let init_kp = generate_keypair();
1306
1307        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1308        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1309
1310        let m1 = ini.write_message(&[]).unwrap();
1311        let _ = res.read_message(&m1).unwrap();
1312        let m2 = res.write_message(&[]).unwrap();
1313        let _ = ini.read_message(&m2).unwrap();
1314
1315        assert!(ini.is_finished() && res.is_finished());
1316
1317        let (ini_send, ini_recv) = ini.raw_split_keys();
1318        let (res_send, res_recv) = res.raw_split_keys();
1319        assert_eq!(
1320            ini_send, res_recv,
1321            "initiator send key matches responder recv key"
1322        );
1323        assert_eq!(
1324            ini_recv, res_send,
1325            "initiator recv key matches responder send key"
1326        );
1327
1328        let mut ini_session = ini.to_session_for_test().unwrap();
1329        let mut res_session = res.to_session_for_test().unwrap();
1330        let snow_ini_ref = ini
1331            .into_inner_snow_for_test()
1332            .into_stateless_transport_mode()
1333            .unwrap();
1334        let snow_res_ref = res
1335            .into_inner_snow_for_test()
1336            .into_stateless_transport_mode()
1337            .unwrap();
1338
1339        let plaintext: Vec<u8> = (0u8..200).collect();
1340
1341        // initiator -> responder: production seal byte-identical to snow's
1342        // write_message (initiator role), and the responder session opens it.
1343        for ctr in 0u64..=4 {
1344            let sealed = ini_session.seal(&plaintext).unwrap();
1345            assert_eq!(sealed.counter, ctr);
1346            let mut snow_out = vec![0u8; plaintext.len() + 16];
1347            let n = snow_ini_ref
1348                .write_message(ctr, &plaintext, &mut snow_out)
1349                .unwrap();
1350            snow_out.truncate(n);
1351            assert_eq!(
1352                sealed.ciphertext, snow_out,
1353                "initiator->responder ciphertext byte-identical to snow at counter {ctr}"
1354            );
1355            assert_eq!(
1356                res_session
1357                    .open(sealed.counter, &sealed.ciphertext)
1358                    .unwrap(),
1359                plaintext,
1360                "responder opens what the initiator sealed at counter {ctr}"
1361            );
1362        }
1363
1364        // responder -> initiator: the symmetric case
1365        for ctr in 0u64..=4 {
1366            let sealed = res_session.seal(&plaintext).unwrap();
1367            assert_eq!(sealed.counter, ctr);
1368            let mut snow_out = vec![0u8; plaintext.len() + 16];
1369            let n = snow_res_ref
1370                .write_message(ctr, &plaintext, &mut snow_out)
1371                .unwrap();
1372            snow_out.truncate(n);
1373            assert_eq!(
1374                sealed.ciphertext, snow_out,
1375                "responder->initiator ciphertext byte-identical to snow at counter {ctr}"
1376            );
1377            assert_eq!(
1378                ini_session
1379                    .open(sealed.counter, &sealed.ciphertext)
1380                    .unwrap(),
1381                plaintext,
1382                "initiator opens what the responder sealed at counter {ctr}"
1383            );
1384        }
1385    }
1386
1387    #[test]
1388    fn seal_into_matches_seal_and_opens() {
1389        let (mut a, mut b) = crate::test_session_pair();
1390        let mut sbuf = Vec::new();
1391        let ctr = a.seal_into(b"reuse me", &mut sbuf).unwrap();
1392        let mut obuf = Vec::new();
1393        b.open_into(ctr, &sbuf, &mut obuf).unwrap();
1394        assert_eq!(obuf, b"reuse me");
1395    }
1396
1397    #[test]
1398    fn test_wide_replay_window_jitter_and_rejection() {
1399        let mut w = ReplayWindow::new();
1400        // Initially accepts 0
1401        assert!(w.check(0));
1402        w.commit(0);
1403
1404        // Advance latest to 100,000
1405        assert!(w.check(100_000));
1406        w.commit(100_000);
1407
1408        // Packet 50,000 (diff = 50,000 < 131,072) must be accepted
1409        assert!(w.check(50_000));
1410        w.commit(50_000);
1411
1412        // Duplicate 50,000 must be rejected
1413        assert!(!w.check(50_000));
1414
1415        // Packet 0 is now 100,000 behind, which is < 131,072, but it was already seen
1416        assert!(!w.check(0));
1417
1418        // Advance to 250,000
1419        assert!(w.check(250_000));
1420        w.commit(250_000);
1421
1422        // Packet 100,000 is 150,000 behind (>= 131,072), must be rejected as too old
1423        assert!(!w.check(100_000));
1424    }
1425
1426    #[test]
1427    fn test_seal_with_counter_and_open_with_window() {
1428        let (a, b) = crate::test_session_pair();
1429        let payload = b"concurrent multi-worker payload";
1430        let sealed = a.seal_with_counter(42, payload).unwrap();
1431        assert_eq!(sealed.counter, 42);
1432
1433        let mut replay = ReplayWindow::new();
1434        let decrypted = b
1435            .open_with_window(sealed.counter, &sealed.ciphertext, &mut replay)
1436            .unwrap();
1437        assert_eq!(decrypted, payload);
1438
1439        // Replay attempt must fail
1440        assert_eq!(
1441            b.open_with_window(sealed.counter, &sealed.ciphertext, &mut replay),
1442            Err(CryptoError::Replay)
1443        );
1444    }
1445
1446    #[test]
1447    fn test_chacha20_poly1305_cipher_in_place() {
1448        let key = [0x5au8; 32];
1449        let cipher = ChaCha20Poly1305Cipher::new(key);
1450        let mut buffer = [0u8; 128];
1451        let plaintext = b"zero-copy in-place packet buffer payload";
1452        buffer[..plaintext.len()].copy_from_slice(plaintext);
1453
1454        let counter = 100u64;
1455        let sealed_len = cipher
1456            .seal_in_place(counter, &mut buffer, plaintext.len())
1457            .expect("seal_in_place should succeed");
1458        assert_eq!(sealed_len, plaintext.len() + 16);
1459        assert_ne!(&buffer[..plaintext.len()], plaintext);
1460
1461        let mut replay = ReplayWindow::new();
1462        let plain_len = cipher
1463            .open_in_place_with_window(counter, &mut buffer, sealed_len, &mut replay)
1464            .expect("open_in_place_with_window should succeed");
1465        assert_eq!(plain_len, plaintext.len());
1466        assert_eq!(&buffer[..plain_len], plaintext);
1467
1468        // Replay rejection
1469        assert_eq!(
1470            cipher.open_in_place_with_window(counter, &mut buffer, sealed_len, &mut replay),
1471            Err(CryptoError::Replay)
1472        );
1473    }
1474
1475    #[test]
1476    fn test_chacha20_poly1305_cipher_debug_and_boundaries() {
1477        let key = [0x5au8; 32];
1478        let cipher = ChaCha20Poly1305Cipher::new(key);
1479        let debug_str = format!("{cipher:?}");
1480        assert!(debug_str.contains("ChaCha20Poly1305Cipher"));
1481
1482        let mut buf = [0u8; 32];
1483        // Exact buffer length (buf.len() == plaintext_len + 16): must succeed
1484        assert_eq!(cipher.seal_in_place(1, &mut buf[..16], 0).unwrap(), 16);
1485
1486        // Buffer smaller than total_len (buf.len() < plaintext_len + 16): must fail
1487        assert_eq!(
1488            cipher.seal_in_place(1, &mut buf[..15], 0),
1489            Err(CryptoError::Decrypt)
1490        );
1491
1492        // open_in_place: sealed_len < 16 must fail
1493        assert_eq!(
1494            cipher.open_in_place(1, &mut buf, 15),
1495            Err(CryptoError::Decrypt)
1496        );
1497        assert_eq!(
1498            cipher.open_in_place(1, &mut buf, 0),
1499            Err(CryptoError::Decrypt)
1500        );
1501
1502        // open_in_place: buf.len() < sealed_len must fail (even when sealed_len >= 16)
1503        assert_eq!(
1504            cipher.open_in_place(1, &mut buf[..10], 16),
1505            Err(CryptoError::Decrypt)
1506        );
1507
1508        // open_in_place: exact boundary sealed_len == 16 and buf.len() == 16
1509        let mut exact = [0u8; 16];
1510        cipher.seal_in_place(2, &mut exact, 0).unwrap();
1511        assert_eq!(cipher.open_in_place(2, &mut exact, 16).unwrap(), 0);
1512    }
1513
1514    #[test]
1515    fn test_seal_into_with_counter_and_open_into_with_window() {
1516        let (a, b) = crate::test_session_pair();
1517        let payload = b"buffer test payload";
1518        let mut sealed_buf = Vec::new();
1519        a.seal_into_with_counter(42, payload, &mut sealed_buf)
1520            .unwrap();
1521        assert_ne!(sealed_buf, payload);
1522        assert_eq!(sealed_buf.len(), payload.len() + 16);
1523
1524        let mut replay = ReplayWindow::new();
1525        let mut out = Vec::new();
1526        b.open_into_with_window(42, &sealed_buf, &mut replay, &mut out)
1527            .unwrap();
1528        assert_eq!(out, payload);
1529
1530        // Replay attempt must fail and not modify output buffer
1531        out.clear();
1532        let err = b.open_into_with_window(42, &sealed_buf, &mut replay, &mut out);
1533        assert_eq!(err, Err(CryptoError::Replay));
1534        assert!(out.is_empty());
1535
1536        // Forged tag must fail
1537        let mut bad_buf = sealed_buf.clone();
1538        bad_buf[0] ^= 1;
1539        let mut fresh_replay = ReplayWindow::new();
1540        let bad_err = b.open_into_with_window(43, &bad_buf, &mut fresh_replay, &mut out);
1541        assert_eq!(bad_err, Err(CryptoError::Decrypt));
1542    }
1543
1544    #[test]
1545    fn test_replay_window_word_idx_distinct_words() {
1546        let mut w = ReplayWindow::new();
1547        w.commit(0);
1548        // Counters across different word indices (64, 128) must not falsely collide with 0
1549        assert!(w.check(64));
1550        assert!(w.check(128));
1551        w.commit(64);
1552        assert!(!w.check(64));
1553        assert!(!w.check(0));
1554        assert!(w.check(128));
1555    }
1556
1557    #[test]
1558    fn test_promote_to_high_throughput_preserves_multi_word_state() {
1559        let mut w = ReplayWindow::new_with_profile(ReplayProfile::Standard);
1560        w.commit(100);
1561        w.commit(200);
1562        w.commit(1000);
1563        w.commit(1500);
1564        w.promote_to_high_throughput();
1565        assert_eq!(w.profile(), ReplayProfile::HighThroughput);
1566
1567        assert!(!w.check(100));
1568        assert!(!w.check(200));
1569        assert!(!w.check(1000));
1570        assert!(!w.check(1500));
1571
1572        assert!(w.check(101));
1573        assert!(w.check(201));
1574        assert!(w.check(1001));
1575        assert!(w.check(1499));
1576    }
1577
1578    #[test]
1579    fn test_promote_to_high_throughput_large_latest() {
1580        // Case 1: latest_word is small (< old_words = 128 words).
1581        // latest_word = 10, so saturating_sub(127) == 0.
1582        // Under mutant `old_words + 1` = 129, saturating_sub(129) == 0 (no diff).
1583        // But under mutant `old_words / 1` = 128, saturating_sub(128) == 0.
1584        //
1585        // Case 2: latest_word >= old_words (e.g. latest_word = 200).
1586        // start_word should be 200 - 127 = 73.
1587        // Under mutant `+`: (old_words + 1) = 129 -> 200 - 129 = 71.
1588        // Under mutant `/`: (old_words / 1) = 128 -> 200 - 128 = 72.
1589        // Counter at word 72 is (72 * 64). Under the correct code (start_word = 73),
1590        // word 72 is NOT copied (it's older than 127 words from latest).
1591        // BUT wait: in standard window of 128 words, word 72 is diff = 200 - 72 = 128 words = 8192 bits!
1592        // 8192 bits is beyond standard window bits (8192), so its slot in old_bitmap is actually word 200's slot!
1593        // Notice: 72 & 127 = 72. 200 & 127 = 72!
1594        // If word 72 is copied to new_bitmap[72], while word 200 is copied to new_bitmap[200],
1595        // then word 72 in new_bitmap would receive the contents of old_bitmap[72]!
1596        // In the correct code (start_word = 73), word 72 is NOT visited, so new_bitmap[72] is 0!
1597        // Under mutant `+` or `/`: start_word <= 72, so loop visits w = 72 and copies old_bitmap[72] into new_bitmap[72]!
1598        // Therefore, counter (72 * 64 + bit) would falsely appear as committed in new_bitmap!
1599        let mut w = ReplayWindow::new_with_profile(ReplayProfile::Standard);
1600        let latest = 200 * 64 + 10;
1601        w.commit(latest);
1602        w.promote_to_high_throughput();
1603
1604        // Counter at word 72 (e.g. 72 * 64 + 10) was never committed (it is 128 words older than latest).
1605        // In correct code, new_bitmap[72] == 0, so check(72 * 64 + 10) is true (not seen, acceptable).
1606        // Under mutant `-` -> `+` or `/`, w = 72 is included in the loop, copying old_bitmap[72] (which has bit 10 set from latest!)
1607        // into new_bitmap[72], so check(72 * 64 + 10) would falsely return false (already seen)!
1608        assert!(
1609            w.check(72 * 64 + 10),
1610            "counter 72 * 64 + 10 was never committed and must be acceptable"
1611        );
1612
1613        // Also test that actual valid tail counter (word 73: 73 * 64 + 10) was committed and preserved:
1614        let mut w2 = ReplayWindow::new_with_profile(ReplayProfile::Standard);
1615        w2.commit(73 * 64 + 10);
1616        w2.commit(latest);
1617        w2.promote_to_high_throughput();
1618        assert!(!w2.check(latest), "latest must still be marked seen");
1619        assert!(
1620            !w2.check(73 * 64 + 10),
1621            "tail counter in word 73 must be marked seen"
1622        );
1623    }
1624
1625    #[test]
1626    fn test_replay_window_circular_word_clearing() {
1627        let mut w = ReplayWindow::new_with_profile(ReplayProfile::Standard);
1628        w.commit(69);
1629        assert!(!w.check(69));
1630
1631        w.commit(8192 + 74);
1632        assert!(!w.check(8192 + 74));
1633
1634        assert!(
1635            w.check(8192 + 69),
1636            "counter 8192 + 69 must not be blocked by stale bit from counter 69"
1637        );
1638    }
1639
1640    #[test]
1641    fn test_replay_window_consecutive_word_clearing() {
1642        let mut w = ReplayWindow::new_with_profile(ReplayProfile::Standard);
1643        w.commit(130);
1644        w.commit(8192 + 70);
1645        w.commit(8192 + 135);
1646        assert!(
1647            w.check(8192 + 130),
1648            "counter 8192 + 130 must be accepted after word 2 is overtaken and cleared"
1649        );
1650
1651        // Kills: replace - with + in diff = counter - self.latest
1652        w.commit(3000);
1653        w.commit(3001);
1654        w.commit(3500); // 3500 - 3001 = 499 < 8192. Under +, 3500 + 3001 = 6501, wait!
1655                        // Under +, diff = 3500 + 3001 = 6501, which is < 8192! We need counter + latest >= 8192:
1656                        // Say latest = 5000, counter = 5100: diff = 5100 - 5000 = 100 < 8192.
1657                        // Under +, diff = 5100 + 5000 = 10100 >= 8192 (large leap clears bitmap!).
1658        w.commit(5000);
1659        w.commit(5001);
1660        w.commit(5100);
1661        assert!(!w.check(5001), "counter 5001 must still be marked seen");
1662    }
1663
1664    #[test]
1665    fn test_session_set_stride_and_send_counter() {
1666        let (mut a, _) = crate::test_session_pair();
1667        assert_eq!(a.send_counter(), 0);
1668        assert_ne!(a.send_counter(), 1);
1669
1670        // Initial configuration on fresh session (send_counter == 0):
1671        a.set_stride(5, 4);
1672        assert_eq!(a.stride(), 4);
1673        assert_eq!(a.send_counter(), 5);
1674
1675        // Stride 0 falls back to 1:
1676        // send_counter was 5. With stride 1, target_rem = 10 % 1 = 0, rem = 5 % 1 = 0.
1677        // send_counter remains 5.
1678        a.set_stride(10, 0);
1679        assert_eq!(a.stride(), 1);
1680        assert_eq!(a.send_counter(), 5);
1681
1682        // Reconfiguration with send_counter > 0 and rem <= target_rem:
1683        // currently send_counter is 5.
1684        // stride = 4, start_counter = 3.
1685        // rem = 5 % 4 = 1.
1686        // target_rem = 3 % 4 = 3.
1687        // rem (1) <= target_rem (3) -> diff = 3 - 1 = 2.
1688        // send_counter becomes 5 + 2 = 7. 7 % 4 == 3.
1689        a.set_stride(3, 4);
1690        assert_eq!(a.stride(), 4);
1691        assert_eq!(a.send_counter(), 7);
1692        assert_eq!(a.send_counter() % 4, 3);
1693
1694        // Reconfiguration with send_counter > 0 and rem > target_rem where (rem - target_rem) > 1:
1695        // currently send_counter is 7.
1696        // Let's set stride = 8, start_counter = 1.
1697        // rem = 7 % 8 = 7.
1698        // target_rem = 1 % 8 = 1.
1699        // rem (7) > target_rem (1) -> (rem - target_rem) = 6.
1700        // Correct diff: 8 - 6 = 2. send_counter becomes 7 + 2 = 9. 9 % 8 = 1.
1701        // Mutant diff: 8 / 6 = 1. send_counter becomes 7 + 1 = 8. 8 % 8 = 0 != 1.
1702        a.set_stride(1, 8);
1703        assert_eq!(a.stride(), 8);
1704        assert_eq!(a.send_counter(), 9);
1705        assert_eq!(a.send_counter() % 8, 1);
1706    }
1707
1708    #[test]
1709    fn test_session_seal_in_place_and_open_in_place_boundaries() {
1710        let (mut a, mut b) = crate::test_session_pair();
1711        let mut buf = [0u8; 32];
1712        // Exact length buf.len() == plaintext_len + 16: succeeds
1713        let ctr = a.seal_in_place(&mut buf[..16], 0).unwrap();
1714        assert_eq!(ctr, 0);
1715
1716        // Buffer smaller than total_len: fails
1717        assert_eq!(
1718            a.seal_in_place(&mut buf[..15], 0),
1719            Err(CryptoError::Decrypt)
1720        );
1721
1722        // open_in_place: sealed_len < 16 fails
1723        assert_eq!(b.open_in_place(0, &mut buf, 15), Err(CryptoError::Decrypt));
1724        assert_eq!(b.open_in_place(0, &mut buf, 0), Err(CryptoError::Decrypt));
1725
1726        // open_in_place: buf.len() < sealed_len fails
1727        assert_eq!(
1728            b.open_in_place(0, &mut buf[..10], 16),
1729            Err(CryptoError::Decrypt)
1730        );
1731
1732        // open_in_place: exact boundary sealed_len == 16 and buf.len() == 16: succeeds
1733        let mut exact = [0u8; 16];
1734        let ctr2 = a.seal_in_place(&mut exact, 0).unwrap();
1735        assert_eq!(b.open_in_place(ctr2, &mut exact, 16).unwrap(), 0);
1736    }
1737
1738    #[test]
1739    fn pq_keypair_encapsulate_decapsulate_roundtrip() {
1740        let (ek, dk) = generate_pq_keypair();
1741        assert_eq!(ek.0.len(), PQ_ENCAPS_KEY_LEN);
1742
1743        let (ct, ss_sender) = pq_encapsulate(&ek).expect("encapsulation succeeds");
1744        assert_eq!(ct.0.len(), PQ_CIPHERTEXT_LEN);
1745        assert_eq!(ss_sender.len(), PQ_SHARED_SECRET_LEN);
1746
1747        let ss_recipient = pq_decapsulate(&dk, &ct).expect("decapsulation succeeds");
1748        assert_eq!(
1749            *ss_sender, *ss_recipient,
1750            "both sides agree on 32-byte shared secret"
1751        );
1752    }
1753
1754    #[test]
1755    fn pq_encapsulate_rejects_fips203_invalid_modulus() {
1756        let (mut ek, _dk) = generate_pq_keypair();
1757        // FIPS 203 Section 7.2 ByteDecode_12 requires all 12-bit coefficients < q (3329).
1758        // Inject 0x0FFF (4095 >= 3329) into the first coefficient pair.
1759        ek.0[0] = 0xFF;
1760        ek.0[1] = 0x0F;
1761        let res = pq_encapsulate(&ek);
1762        assert!(
1763            res.is_err(),
1764            "encapsulation must reject coefficient >= 3329"
1765        );
1766    }
1767
1768    #[test]
1769    fn pq_decapsulate_implicit_rejection_on_tampered_ciphertext() {
1770        let (ek, dk) = generate_pq_keypair();
1771        let (mut ct, ss_sender) = pq_encapsulate(&ek).expect("encapsulation succeeds");
1772
1773        // Tamper with the ciphertext: flip one bit in ct
1774        ct.0[0] ^= 0x01;
1775        let ss_recipient =
1776            pq_decapsulate(&dk, &ct).expect("decapsulation succeeds via implicit rejection");
1777        assert_ne!(
1778            *ss_sender, *ss_recipient,
1779            "FIPS 203 implicit rejection must yield a distinct pseudorandom secret on tampered ciphertext"
1780        );
1781    }
1782
1783    #[test]
1784    fn pq_decaps_key_debug_redacted() {
1785        let (_ek, dk) = generate_pq_keypair();
1786        assert_eq!(format!("{:?}", dk), "PqDecapsKey([REDACTED])");
1787    }
1788
1789    #[test]
1790    fn pq_encapsulate_rejects_fips203_invalid_modulus_second_coeff() {
1791        let (mut ek, _dk) = generate_pq_keypair();
1792        // FIPS 203 Section 7.2 ByteDecode_12 requires all 12-bit coefficients < q (3329).
1793        // Inject 0x0FFF (4095 >= 3329) into the second coefficient (d1) of the first chunk.
1794        // d1 = (chunk[1] >> 4) | (chunk[2] << 4)
1795        // Setting chunk[1] = 0xF0 and chunk[2] = 0xFF gives d1 = 0xFFF = 4095.
1796        // Ensure d0 = chunk[0] | ((chunk[1] & 0x0F) << 8) remains < 3329: chunk[0]=0, chunk[1]&0x0F=0 => d0 = 0.
1797        ek.0[0] = 0x00;
1798        ek.0[1] = 0xF0;
1799        ek.0[2] = 0xFF;
1800        let res = pq_encapsulate(&ek);
1801        assert!(
1802            res.is_err(),
1803            "encapsulation must reject second coefficient >= 3329"
1804        );
1805    }
1806
1807    #[test]
1808    fn pq_hybrid_handshake_roundtrip_with_payload() {
1809        let init_kp = generate_keypair();
1810        let resp_kp = generate_keypair();
1811
1812        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1813        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1814
1815        let cert_a = b"initiator-test-certificate-128b";
1816        let cert_b = b"responder-test-certificate-128b";
1817
1818        // Msg 1: Initiator -> Responder
1819        let msg1 = ini.write_message(cert_a).unwrap();
1820        // Msg 1 wire size = 32(e) + 48(s) + 800(ek) + cert_a.len() + 16(tag)
1821        assert_eq!(msg1.len(), 32 + 48 + PQ_ENCAPS_KEY_LEN + cert_a.len() + 16);
1822
1823        let p1 = res.read_message(&msg1).unwrap();
1824        assert_eq!(p1, cert_a);
1825        assert_eq!(res.remote_static(), Some(init_kp.public));
1826
1827        // Msg 2: Responder -> Initiator
1828        let msg2 = res.write_message(cert_b).unwrap();
1829        // Msg 2 wire size = 768(ct) + 32(e) + cert_b.len() + 16(tag)
1830        assert_eq!(msg2.len(), PQ_CIPHERTEXT_LEN + 32 + cert_b.len() + 16);
1831
1832        let p2 = ini.read_message(&msg2).unwrap();
1833        assert_eq!(p2, cert_b);
1834        assert_eq!(ini.remote_static(), Some(resp_kp.public));
1835
1836        assert!(ini.is_finished());
1837        assert!(res.is_finished());
1838        assert_eq!(ini.channel_binding(), res.channel_binding());
1839
1840        let (ini_send, ini_recv) = ini.raw_split_keys();
1841        let (res_send, res_recv) = res.raw_split_keys();
1842        assert_eq!(ini_send, res_recv);
1843        assert_eq!(ini_recv, res_send);
1844
1845        let mut ini_session = ini.into_session().unwrap();
1846        let mut res_session = res.into_session().unwrap();
1847        let sealed = ini_session.seal(b"quantum safe transport payload").unwrap();
1848        let opened = res_session
1849            .open(sealed.counter, &sealed.ciphertext)
1850            .unwrap();
1851        assert_eq!(opened, b"quantum safe transport payload");
1852    }
1853
1854    #[test]
1855    fn pq_hybrid_handshake_rejects_tampered_ciphertext() {
1856        let init_kp = generate_keypair();
1857        let resp_kp = generate_keypair();
1858
1859        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1860        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1861
1862        let msg1 = ini.write_message(b"certA").unwrap();
1863        let _ = res.read_message(&msg1).unwrap();
1864        let mut msg2 = res.write_message(b"certB").unwrap();
1865
1866        // Tamper with the cleartext ct prefix of msg2
1867        msg2[10] ^= 0x01;
1868
1869        // Initiator decapsulates to a mismatched ss, causing snow Poly1305 MAC to fail
1870        let err = ini.read_message(&msg2);
1871        assert!(err.is_err(), "tampered ct must fail handshake verification");
1872    }
1873
1874    #[test]
1875    fn pq_handshake_read_message_rejects_truncated_message() {
1876        let init_kp = generate_keypair();
1877        let resp_kp = generate_keypair();
1878
1879        let mut ini = Handshake::initiator(&init_kp.private, &resp_kp.public).unwrap();
1880        let mut res = Handshake::responder(&resp_kp.private).unwrap();
1881
1882        // Initiator sends valid msg1 to responder
1883        let msg1 = ini.write_message(b"certA").unwrap();
1884        // Responder rejects truncated msg1
1885        assert!(res.read_message(&msg1[..10]).is_err());
1886
1887        let _ = res.read_message(&msg1).unwrap();
1888        let msg2 = res.write_message(b"certB").unwrap();
1889
1890        // Initiator rejects msg2 shorter than PQ_CIPHERTEXT_LEN
1891        let truncated = vec![0u8; PQ_CIPHERTEXT_LEN - 1];
1892        assert!(ini.read_message(&truncated).is_err());
1893        assert!(ini.read_message(&msg2[..PQ_CIPHERTEXT_LEN]).is_err());
1894    }
1895
1896    #[test]
1897    fn session_check_replay_is_non_committing() {
1898        let (a, mut b) = crate::test_session_pair();
1899        assert!(b.check_replay(1));
1900        assert!(b.check_replay(1));
1901
1902        let sealed = a.seal_with_counter(1, b"hello").unwrap();
1903        let opened = b.open(sealed.counter, &sealed.ciphertext).unwrap();
1904        assert_eq!(opened, b"hello");
1905
1906        assert!(!b.check_replay(1));
1907    }
1908}