stochastic-routing-extended 1.0.2

//! End-to-end SRX pipeline: wires together padding, encryption, framing,
//! mimicry, jitter, cover traffic, and transport dispatch into a single send/recv path.
//!
//! ## Send path
//!
//! ```text
//! VPN payload
//!   -> pad  (PaddingStrategy)
//!   -> encrypt  (AEAD via Session + AeadPipeline)
//!   -> frame  (FrameCodec: counter XOR mask as frame_id, routing mask, split ct/mac)
//!   -> encode  (wire bytes)
//!   -> mimicry wrap  (ProtocolMimicry)
//!   -> jitter delay  (JitterModel)
//!   -> dispatch  (TransportManager via RoutingMask)
//! ```
//!
//! ## Recv path
//!
//! ```text
//! wire bytes from transport
//!   -> mimicry unwrap
//!   -> frame decode
//!   -> reassemble ct||mac
//!   -> recover counter (frame_id XOR mask)
//!   -> decrypt  (AEAD)
//!   -> unpad
//!   -> VPN payload
//! ```

pub mod cover;
pub mod ratelimit;

use std::sync::{Arc, Mutex};

use bytes::Bytes;

use crate::config::{AeadCipher as AeadVariant, MimicryMode};
use crate::crypto::{AeadPipeline, KeyDerivation, ReplayState, ReplayWindow};
use crate::error::{FrameError, Result, SrxError};
use crate::frame::codec::{Frame, FrameCodec, WIRE_MAC_LEN};
use crate::masking::jitter::JitterModel;
use crate::masking::mimicry::ProtocolMimicry;
use crate::masking::padding::PaddingStrategy;
use crate::metrics::PipelineMetrics;
use crate::routing::RoutingMask;
use crate::routing::aggregation::{AggregationStrategy, BandwidthAggregator};
use crate::session::{ReKeyScheduler, Session};
use crate::signaling::inband::{InBandSignaling, Signal};
use crate::transport::{TransportKind, TransportManager};

use self::ratelimit::RateLimiter;

/// Full SRX send/recv pipeline tying together all protocol layers.
pub struct SrxPipeline {
    /// Active session (owns seed RNG, data key, packet counter).
    pub session: Session,
    /// Current parallel AEAD worker pool.
    pipeline: Arc<AeadPipeline>,
    /// Previous AEAD pipeline (kept for in-flight packets after the re-key).
    prev_pipeline: Option<Arc<AeadPipeline>>,
    /// Packet counter at which the current pipeline epoch started.
    rekey_boundary: u64,
    /// Re-key scheduler (deterministic pseudo-random intervals).
    rekey_scheduler: Option<ReKeyScheduler>,
    /// AEAD variant (needed to rebuild the pipeline on re-key).
    aead_variant: AeadVariant,
    /// AEAD worker count (needed to rebuild the pipeline on the re-key).
    aead_worker_count: usize,
    /// Random-length padding.
    padding: PaddingStrategy,
    /// Protocol mimicry wrapper.
    mimicry: ProtocolMimicry,
    /// Optional jitter model (inter-packet delay).
    jitter: Option<JitterModel>,
    /// Frame encoder/decoder.
    codec: FrameCodec,
    /// Transport dispatcher.
    transport_mgr: TransportManager,
    /// Transports available for routing mask selection.
    available_transports: Vec<TransportKind>,
    /// XOR mask for obfuscating packet counter in frame_id.
    frame_id_mask: u64,
    /// In-band signaling encoder/decoder.
    signaling: InBandSignaling,
    /// Optional bandwidth aggregator for multi-transport send.
    bandwidth_aggregator: Option<BandwidthAggregator>,
    /// Pipeline metrics (bytes, packets, rekeys, etc.).
    metrics: PipelineMetrics,
    /// Optional rate limiter for traffic shaping.
    rate_limiter: Option<RateLimiter>,
    /// Anti-replay sliding window for received packets.
    replay_window: Mutex<ReplayWindow>,
}

impl SrxPipeline {
    /// Build a pipeline from its components.
    ///
    /// Derives the frame-ID XOR mask from the session seed so that both peers
    /// produce the same mask without extra messages.
    pub fn new(
        session: Session,
        pipeline: Arc<AeadPipeline>,
        padding: PaddingStrategy,
        mimicry_mode: MimicryMode,
        jitter: Option<JitterModel>,
        transport_mgr: TransportManager,
    ) -> Self {
        let available_transports = transport_mgr.active_kinds();
        let frame_id_mask = KeyDerivation::frame_id_mask(&session.rng.seed_bytes()).unwrap_or(0);
        Self {
            session,
            pipeline,
            prev_pipeline: None,
            rekey_boundary: 0,
            rekey_scheduler: None,
            aead_variant: AeadVariant::ChaCha20Poly1305,
            aead_worker_count: 2,
            padding,
            mimicry: ProtocolMimicry::new(mimicry_mode),
            jitter,
            codec: FrameCodec::new(),
            transport_mgr,
            available_transports,
            frame_id_mask,
            signaling: InBandSignaling::new(),
            bandwidth_aggregator: None,
            metrics: PipelineMetrics::new(),
            rate_limiter: None,
            replay_window: Mutex::new(ReplayWindow::new()),
        }
    }

    /// Build a pipeline from [`SrxConfig`], a [`Session`] (from handshake), and a
    /// pre-built [`TransportManager`].
    ///
    /// Reads masking/jitter/padding settings from the config.
    pub fn from_config(
        config: &crate::config::SrxConfig,
        session: Session,
        pipeline: Arc<AeadPipeline>,
        transport_mgr: TransportManager,
    ) -> Self {
        let seed = session.rng.seed_bytes();
        let jitter = if config.masking.jitter_enabled {
            Some(JitterModel::new(
                crate::seed::SeedRng::new(seed),
                std::time::Duration::from_millis(5),
                std::time::Duration::from_millis(50),
            ))
        } else {
            None
        };
        let max_padding = if config.masking.padding_enabled {
            128
        } else {
            0
        };
        let (min_rekey, max_rekey) = config.crypto.rekey_interval;
        let rekey_scheduler = Some(ReKeyScheduler::new(
            crate::seed::SeedRng::new(seed),
            min_rekey,
            max_rekey,
        ));
        let mut pipe = Self::new(
            session,
            pipeline,
            PaddingStrategy::new(crate::seed::SeedRng::new(seed), max_padding),
            config.masking.mimicry_mode,
            jitter,
            transport_mgr,
        );
        pipe.rekey_scheduler = rekey_scheduler;
        pipe.aead_variant = config.crypto.aead;
        pipe.aead_worker_count = config.crypto.aead_worker_count;
        pipe
    }

    /// Refresh the cached list of available transports from the manager.
    pub fn refresh_transports(&mut self) {
        self.available_transports = self.transport_mgr.active_kinds();
    }

    /// Enable bandwidth aggregation across multiple transports.
    ///
    /// When enabled, large payloads are split into fragments and distributed
    /// across healthy transports proportional to their health scores.
    pub fn enable_aggregation(&mut self, strategy: AggregationStrategy, fragment_size: usize) {
        self.bandwidth_aggregator = Some(BandwidthAggregator::new(strategy, fragment_size));
    }

    /// Disable bandwidth aggregation.
    pub fn disable_aggregation(&mut self) {
        self.bandwidth_aggregator = None;
    }

    /// Check if aggregation is enabled.
    pub fn is_aggregation_enabled(&self) -> bool {
        self.bandwidth_aggregator.is_some()
    }

    /// Access the underlying transport manager (e.g. to add/remove transports).
    pub fn transport_mgr(&self) -> &TransportManager {
        &self.transport_mgr
    }

    /// Mutable access to the transport manager.
    pub fn transport_mgr_mut(&mut self) -> &mut TransportManager {
        &mut self.transport_mgr
    }

    /// Access pipeline metrics.
    pub fn metrics(&self) -> &PipelineMetrics {
        &self.metrics
    }

    /// Export anti-replay state for persistence across restarts.
    pub fn replay_state(&self) -> ReplayState {
        let rw = self.replay_window.lock().unwrap();
        rw.snapshot()
    }

    /// Restore anti-replay state (e.g. after process restart).
    pub fn set_replay_state(&self, state: &ReplayState) -> Result<()> {
        let mut rw = self.replay_window.lock().unwrap();
        if rw.restore(state) {
            Ok(())
        } else {
            Err(SrxError::Frame(FrameError::Corrupted(
                "invalid replay state snapshot".into(),
            )))
        }
    }

    /// Enable rate limiting (token bucket) for outgoing traffic.
    pub fn enable_rate_limit(&mut self, rate_bytes_per_sec: u64, burst_bytes: u64) {
        self.rate_limiter = Some(RateLimiter::new(rate_bytes_per_sec, burst_bytes));
    }

    /// Disable rate limiting.
    pub fn disable_rate_limit(&mut self) {
        self.rate_limiter = None;
    }

    /// Perform a re-key: derive new data key, build a new AEAD pipeline, keep
    /// the old one for in-flight packets.
    fn perform_rekey(&mut self) -> Result<()> {
        let new_key = self.session.rekey()?;
        let new_pipeline = Arc::new(
            AeadPipeline::new(self.aead_variant, &new_key, self.aead_worker_count).map_err(
                |e| {
                    SrxError::Session(crate::error::SessionError::ReKeyFailed(format!(
                        "failed to build AEAD pipeline: {e}"
                    )))
                },
            )?,
        );
        self.prev_pipeline = Some(std::mem::replace(&mut self.pipeline, new_pipeline));
        self.rekey_boundary = self.session.packet_counter;
        self.metrics.record_rekey();
        tracing::info!(
            key_index = self.session.key_index,
            boundary = self.rekey_boundary,
            "re-key performed"
        );
        Ok(())
    }

    // ------------------------------------------------------------------
    // Send path
    // ------------------------------------------------------------------

    /// Full send path: pad -> encrypt -> frame -> encode -> mimicry -> jitter -> dispatch.
    ///
    /// Returns the [`TransportKind`] that carried the frame.
    pub async fn send(&mut self, payload: &[u8]) -> Result<TransportKind> {
        // Rate limiting (wait for a token bucket if enabled).
        if let Some(ref mut rl) = self.rate_limiter {
            rl.consume(payload.len()).await;
        }

        let envelope = self.prepare_outgoing(payload)?;

        // Jitter delay (optional).
        if let Some(ref mut jitter) = self.jitter {
            let delay = jitter.next_delay();
            tokio::time::sleep(delay).await;
        }

        // Dispatch via routing mask: prefer healthy transports from the mask.
        let counter = self.session.packet_counter;
        let mask =
            RoutingMask::generate(&mut self.session.rng, &self.available_transports, counter);

        // First pass: mask-selected AND healthy.
        for kind in &mask.transports {
            if self.transport_mgr.get(*kind).is_some() && !self.transport_mgr.is_blocked(*kind) {
                self.transport_mgr
                    .send(*kind, Bytes::from(envelope.clone()))
                    .await?;
                return Ok(*kind);
            }
        }

        // Second pass: any healthy transport.
        let healthy = self.transport_mgr.healthy_kinds();
        if let Some(&kind) = healthy.first() {
            self.transport_mgr
                .send(kind, Bytes::from(envelope.clone()))
                .await?;
            return Ok(kind);
        }

        // Third pass: fallback controller (advances past blocked transports).
        if let Some(kind) = self.transport_mgr.try_fallback() {
            self.transport_mgr.send(kind, Bytes::from(envelope)).await?;
            return Ok(kind);
        }

        // Last resort: first registered transport (even if blocked).
        if let Some(&kind) = self.available_transports.first() {
            self.transport_mgr.send(kind, Bytes::from(envelope)).await?;
            return Ok(kind);
        }

        Err(SrxError::Transport(
            crate::error::TransportError::AllTransportsExhausted,
        ))
    }

    /// Send with bandwidth aggregation across multiple transports.
    ///
    /// Splits the payload into fragments and distributes them across
    /// healthy transports. Returns a map of transport -> list of envelopes
    /// that were sent via each transport.
    ///
    /// Requires aggregation to be enabled via [`Self::enable_aggregation`].
    /// If aggregation is not enabled or no healthy transports are available,
    /// falls back to regular [`Self::send`].
    pub async fn send_aggregated(
        &mut self,
        payload: &[u8],
    ) -> Result<Vec<(TransportKind, Vec<u8>)>> {
        // If no aggregator, fall back to regular send.
        let Some(ref mut aggregator) = self.bandwidth_aggregator else {
            let kind = self.send(payload).await?;
            return Ok(vec![(kind, payload.to_vec())]);
        };

        // Get transport health scores.
        let scores = self.transport_mgr.health_scores();
        let healthy_scores: Vec<(TransportKind, f64)> =
            scores.into_iter().filter(|(_, s)| *s > 0.0).collect();

        if healthy_scores.is_empty() {
            return Err(SrxError::Transport(
                crate::error::TransportError::AllTransportsExhausted,
            ));
        }

        // Split the payload into fragments.
        let assignments = aggregator.distribute(payload, &healthy_scores);
        if assignments.is_empty() {
            return Err(SrxError::Transport(
                crate::error::TransportError::AllTransportsExhausted,
            ));
        }

        // Send each fragment through its assigned transport.
        let mut sent: Vec<(TransportKind, Vec<u8>)> = Vec::with_capacity(assignments.len());
        for assignment in assignments {
            // Jitter delay per fragment (optional).
            if let Some(ref mut jitter) = self.jitter {
                let delay = jitter.next_delay();
                tokio::time::sleep(delay).await;
            }

            // Wrap the fragment through the pipeline (pad -> encrypt -> frame -> mimicry).
            let envelope = self.prepare_outgoing(&assignment.data)?;

            // Send via assigned transport.
            self.transport_mgr
                .send(assignment.transport, Bytes::from(envelope.clone()))
                .await?;

            sent.push((assignment.transport, envelope));
        }

        Ok(sent)
    }

    /// Prepare outgoing bytes without dispatching (useful for tests or custom transport).
    ///
    /// Returns the mimicry-wrapped wire bytes ready to send.
    pub fn prepare_outgoing(&mut self, payload: &[u8]) -> Result<Vec<u8>> {
        // 1. Pad.
        let padded = self.padding.pad(payload);

        // 2. Encrypt (advances session.packet_counter, derives nonce internally).
        let ct = self
            .session
            .encrypt_with_pipeline(&self.pipeline, &padded)?;

        // 3. Split ciphertext||tag into Frame fields.
        if ct.len() < WIRE_MAC_LEN {
            return Err(SrxError::Frame(FrameError::Corrupted(
                "ciphertext shorter than MAC length".into(),
            )));
        }
        let tag_start = ct.len() - WIRE_MAC_LEN;
        let payload_bytes = Bytes::from(ct[..tag_start].to_vec());
        let mac: [u8; WIRE_MAC_LEN] = ct[tag_start..].try_into().expect("WIRE_MAC_LEN slice");

        // 4. Build frame (frame_id = counter XOR mask; looks random on the wire).
        let counter = self.session.packet_counter;
        let mask =
            RoutingMask::generate(&mut self.session.rng, &self.available_transports, counter);

        let frame = Frame {
            frame_id: counter ^ self.frame_id_mask,
            routing_mask: mask.value,
            payload: payload_bytes,
            mac,
            fragment: None,
        };

        // 5. Encode to wire bytes.
        let wire = self.codec.encode(&frame)?;

        // 6. Mimicry wrap.
        let result = self.mimicry.wrap(&wire);

        // 7. Track sent metrics.
        self.metrics.total_sent.record(result.len() as u64);

        // 8. Check re-key: if the counter crossed the threshold, rotate the key for the next packet.
        let needs_rekey = self
            .rekey_scheduler
            .as_ref()
            .is_some_and(|s| s.should_rekey(counter));
        if needs_rekey {
            self.perform_rekey()?;
            if let Some(ref mut scheduler) = self.rekey_scheduler {
                scheduler.schedule_next();
            }
        }

        Ok(result)
    }

    // ------------------------------------------------------------------
    // Recv path
    // ------------------------------------------------------------------

    /// Receive from a specific transport and process the full recv pipeline.
    pub async fn recv_from(&mut self, kind: TransportKind) -> Result<Vec<u8>> {
        let envelope = self.transport_mgr.recv(kind).await?;
        self.process_incoming(&envelope)
    }

    /// Process a received envelope through the recv pipeline (unwrap -> decode -> decrypt -> unpad).
    ///
    /// This is the pure transformation without transport I/O, useful when the caller
    /// already has the raw bytes (e.g. from a worker queue).
    pub fn process_incoming(&self, envelope: &[u8]) -> Result<Vec<u8>> {
        // 1. Mimicry unwrap.
        let wire = self.mimicry.unwrap(envelope).ok_or_else(|| {
            SrxError::Frame(FrameError::Corrupted("mimicry unwrap failed".into()))
        })?;

        // 2. Decode frame.
        let frame = self.codec.decode(&wire)?;

        // 3. Reconstruct ciphertext||tag.
        let mut ct = frame.payload.to_vec();
        ct.extend_from_slice(&frame.mac);

        // 4. Recover sender's packet counter (undo XOR mask) and derive nonce.
        let counter = frame.frame_id ^ self.frame_id_mask;

        // 5. Anti-replay check (before decryption — reject early).
        {
            let rw = self.replay_window.lock().unwrap();
            if !rw.check(counter) {
                return Err(SrxError::Frame(FrameError::Corrupted(
                    "replay detected: duplicate or too-old packet counter".into(),
                )));
            }
        }

        let nonce = KeyDerivation::derive_nonce(&self.session.rng.seed_bytes(), counter)?;

        // 6. Decrypt: pick the right AEAD pipeline epoch.
        //    If the counter is from before the current re-key boundary, and we
        //    still have the previous pipeline, use it; otherwise use current.
        let pipe = match self.prev_pipeline {
            Some(ref prev) if self.rekey_boundary > 0 && counter < self.rekey_boundary => prev,
            _ => &self.pipeline,
        };
        let padded = self.session.decrypt_with_pipeline(pipe, nonce, ct)?;

        // 7. Mark counter as seen AFTER successful decryption (prevents poisoning).
        {
            let mut rw = self.replay_window.lock().unwrap();
            rw.accept(counter);
        }

        // 8. Unpad.
        let result = PaddingStrategy::unpad(&padded)
            .ok_or_else(|| SrxError::Frame(FrameError::Corrupted("unpad failed".into())))?;

        // Track metrics.
        self.metrics.total_received.record(envelope.len() as u64);

        Ok(result)
    }

    /// Process incoming with receiver-side re-key.
    ///
    /// Like [`process_incoming`], but also checks the re-key scheduler and
    /// rotates the reception key when the sender's counter crosses a threshold.
    /// Use this when the receiver needs to stay in sync with the sender's
    /// deterministic re-key schedule.
    pub fn process_incoming_with_rekey(&mut self, envelope: &[u8]) -> Result<Vec<u8>> {
        // 1. Mimicry unwrap.
        let wire = self.mimicry.unwrap(envelope).ok_or_else(|| {
            SrxError::Frame(FrameError::Corrupted("mimicry unwrap failed".into()))
        })?;

        // 2. Decode frame.
        let frame = self.codec.decode(&wire)?;

        // 3. Reconstruct ciphertext||tag.
        let mut ct = frame.payload.to_vec();
        ct.extend_from_slice(&frame.mac);

        // 4. Recover counter.
        let counter = frame.frame_id ^ self.frame_id_mask;

        // 5. Anti-replay check.
        {
            let rw = self.replay_window.lock().unwrap();
            if !rw.check(counter) {
                return Err(SrxError::Frame(FrameError::Corrupted(
                    "replay detected: duplicate or too-old packet counter".into(),
                )));
            }
        }

        let nonce = KeyDerivation::derive_nonce(&self.session.rng.seed_bytes(), counter)?;

        // 6. Decrypt with the correct epoch (before re-key, matching sender order).
        let pipe = match self.prev_pipeline {
            Some(ref prev) if self.rekey_boundary > 0 && counter < self.rekey_boundary => prev,
            _ => &self.pipeline,
        };
        let padded = self.session.decrypt_with_pipeline(pipe, nonce, ct)?;

        // 7. Mark counter as seen after successful decryption.
        {
            let mut rw = self.replay_window.lock().unwrap();
            rw.accept(counter);
        }

        // 8. Re-key AFTER decrypt (sender encrypted then re-keyed, so we must
        //    decrypt with the old key first, then rotate).
        let needs_rekey = self
            .rekey_scheduler
            .as_ref()
            .is_some_and(|s| s.should_rekey(counter));
        if needs_rekey {
            self.perform_rekey()?;
            if let Some(ref mut scheduler) = self.rekey_scheduler {
                scheduler.schedule_next();
            }
        }

        // 9. Unpad.
        PaddingStrategy::unpad(&padded)
            .ok_or_else(|| SrxError::Frame(FrameError::Corrupted("unpad failed".into())))
    }

    // ------------------------------------------------------------------
    // In-band signaling
    // ------------------------------------------------------------------

    /// Send a control signal through the full pipeline (encrypted, framed,
    /// mimicry-wrapped — indistinguishable from data on the wire).
    pub fn prepare_signal(&mut self, signal: &Signal) -> Result<Vec<u8>> {
        self.metrics.record_signal();
        let encoded = self.signaling.encode(signal);
        self.prepare_outgoing(&encoded)
    }

    /// Try to decode a decrypted payload as an in-band signal.
    ///
    /// Returns `Some(Signal)` if the payload contains a valid signal frame,
    /// `None` if it is regular application data.
    pub fn try_decode_signal(&self, payload: &[u8]) -> Option<Signal> {
        self.signaling.decode(payload)
    }

    /// Process a received envelope and, if it contains a signal, return it
    /// separately from application data.
    ///
    /// Returns `Ok(Payload::Signal(s))` for control signals or
    /// `Ok(Payload::Data(bytes))` for application payloads.
    pub fn process_incoming_dispatched(&self, envelope: &[u8]) -> Result<Payload> {
        let data = self.process_incoming(envelope)?;
        match self.signaling.decode(&data) {
            Some(sig) => Ok(Payload::Signal(sig)),
            None => Ok(Payload::Data(data)),
        }
    }
}

/// Discriminated result from [`SrxPipeline::process_incoming_dispatched`].
#[derive(Debug)]
pub enum Payload {
    /// Regular application data.
    Data(Vec<u8>),
    /// In-band control signal.
    Signal(Signal),
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::AeadCipher as Variant;
    use crate::seed::SeedRng;

    /// Build a minimal pipeline pair (no real transport) for unit testing the
    /// prepare_outgoing / process_incoming roundtrip.
    fn test_pipeline_pair() -> (SrxPipeline, SrxPipeline) {
        let seed = [0xABu8; 32];
        let key = [0xCDu8; 32];
        let pipe = Arc::new(AeadPipeline::new(Variant::ChaCha20Poly1305, &key, 2).unwrap());

        let make = |id| {
            SrxPipeline::new(
                Session::new(id, seed, key),
                pipe.clone(),
                PaddingStrategy::new(SeedRng::new(seed), 64),
                MimicryMode::Https,
                None,
                TransportManager::new(),
            )
        };

        (make(1), make(2))
    }

    #[test]
    fn prepare_process_roundtrip() {
        let (mut sender, receiver) = test_pipeline_pair();
        let plaintext = b"hello from pipeline";

        let envelope = sender.prepare_outgoing(plaintext).unwrap();
        let recovered = receiver.process_incoming(&envelope).unwrap();

        assert_eq!(recovered, plaintext);
    }

    #[test]
    fn multiple_packets_roundtrip() {
        let (mut sender, receiver) = test_pipeline_pair();

        for i in 0u8..5 {
            let msg = format!("packet-{i}");
            let envelope = sender.prepare_outgoing(msg.as_bytes()).unwrap();
            let recovered = receiver.process_incoming(&envelope).unwrap();
            assert_eq!(recovered, msg.as_bytes());
        }
    }

    #[test]
    fn mimicry_none_roundtrip() {
        let seed = [0x11u8; 32];
        let key = [0x22u8; 32];
        let pipe = Arc::new(AeadPipeline::new(Variant::ChaCha20Poly1305, &key, 2).unwrap());

        let mut sender = SrxPipeline::new(
            Session::new(1, seed, key),
            pipe.clone(),
            PaddingStrategy::new(SeedRng::new(seed), 0),
            MimicryMode::None,
            None,
            TransportManager::new(),
        );
        let receiver = SrxPipeline::new(
            Session::new(2, seed, key),
            pipe,
            PaddingStrategy::new(SeedRng::new(seed), 0),
            MimicryMode::None,
            None,
            TransportManager::new(),
        );

        let envelope = sender.prepare_outgoing(b"no-mimicry").unwrap();
        let recovered = receiver.process_incoming(&envelope).unwrap();
        assert_eq!(recovered, b"no-mimicry");
    }

    #[test]
    fn tampered_envelope_fails() {
        let (mut sender, receiver) = test_pipeline_pair();
        let mut envelope = sender.prepare_outgoing(b"secret").unwrap();
        // Flip a byte inside the encrypted region.
        if envelope.len() > 20 {
            envelope[20] ^= 0xFF;
        }
        assert!(receiver.process_incoming(&envelope).is_err());
    }

    /// Build a pipeline pair with re-key enabled every 3 packets.
    fn test_pipeline_pair_with_rekey() -> (SrxPipeline, SrxPipeline) {
        let seed = [0xABu8; 32];
        let key = [0xCDu8; 32];
        let pipe = Arc::new(AeadPipeline::new(Variant::ChaCha20Poly1305, &key, 2).unwrap());

        let make = |id| {
            let mut p = SrxPipeline::new(
                Session::new(id, seed, key),
                pipe.clone(),
                PaddingStrategy::new(SeedRng::new(seed), 64),
                MimicryMode::Https,
                None,
                TransportManager::new(),
            );
            // Re-key after every 3 packets (deterministic schedule, same seed).
            p.rekey_scheduler = Some(ReKeyScheduler::new(SeedRng::new(seed), 3, 3));
            p.aead_variant = Variant::ChaCha20Poly1305;
            p.aead_worker_count = 2;
            p
        };

        (make(1), make(2))
    }

    #[test]
    fn rekey_triggers_after_threshold() {
        let (mut sender, _) = test_pipeline_pair_with_rekey();

        // Send 4 packets — re-key should trigger after packet 3.
        for i in 0u8..4 {
            sender
                .prepare_outgoing(format!("msg-{i}").as_bytes())
                .unwrap();
        }
        assert!(
            sender.session.key_index > 0,
            "sender should have re-keyed after 3 packets"
        );
    }

    #[test]
    fn rekey_roundtrip_with_receiver_rekey() {
        let (mut sender, mut receiver) = test_pipeline_pair_with_rekey();

        // Send 8 packets through sender, collect envelopes.
        let mut envelopes = Vec::new();
        for i in 0u8..8 {
            let msg = format!("rekey-{i}");
            envelopes.push((
                msg.clone(),
                sender.prepare_outgoing(msg.as_bytes()).unwrap(),
            ));
        }

        // Receiver processes all using process_incoming_with_rekey.
        for (msg, envelope) in &envelopes {
            let recovered = receiver.process_incoming_with_rekey(envelope).unwrap();
            assert_eq!(recovered, msg.as_bytes(), "mismatch for {msg}");
        }

        // Both sides should have re-keyed the same number of times.
        assert_eq!(sender.session.key_index, receiver.session.key_index);
    }

    #[test]
    fn session_rekey_derives_new_key() {
        let seed = [0x44u8; 32];
        let key = [0x55u8; 32];
        let mut session = Session::new(1, seed, key);
        let old_key = session.data_key;
        let new_key = session.rekey().unwrap();
        assert_ne!(old_key, new_key);
        assert_eq!(session.key_index, 1);
        assert_eq!(session.data_key, new_key);
    }

    #[test]
    fn signal_roundtrip_through_pipeline() {
        let (mut sender, receiver) = test_pipeline_pair();

        let envelope = sender.prepare_signal(&Signal::ReKey).unwrap();
        let payload = receiver.process_incoming_dispatched(&envelope).unwrap();
        match payload {
            Payload::Signal(sig) => assert_eq!(sig, Signal::ReKey),
            Payload::Data(_) => panic!("expected signal, got data"),
        }
    }

    #[test]
    fn data_not_mistaken_for_signal() {
        let (mut sender, receiver) = test_pipeline_pair();

        let envelope = sender.prepare_outgoing(b"regular data").unwrap();
        let payload = receiver.process_incoming_dispatched(&envelope).unwrap();
        match payload {
            Payload::Data(d) => assert_eq!(d, b"regular data"),
            Payload::Signal(_) => panic!("expected data, got signal"),
        }
    }

    #[test]
    fn custom_signal_with_payload() {
        let (mut sender, receiver) = test_pipeline_pair();

        let sig = Signal::Custom(b"switch-to-quic".to_vec());
        let envelope = sender.prepare_signal(&sig).unwrap();
        let payload = receiver.process_incoming_dispatched(&envelope).unwrap();
        match payload {
            Payload::Signal(s) => assert_eq!(s, sig),
            Payload::Data(_) => panic!("expected signal"),
        }
    }
}