stochastic-routing-extended 1.0.2

//! Active session state.

use crate::crypto::{AeadPipeline, KeyDerivation};
use crate::seed::SeedRng;

/// An established SRX session holding all runtime state.
pub struct Session {
    /// Session identifier.
    pub id: u64,
    /// Current seed RNG (shared with peer).
    pub rng: SeedRng,
    /// Current data encryption key.
    pub data_key: [u8; 32],
    /// Packet counter for this session.
    pub packet_counter: u64,
    /// Whether the session is still active.
    pub active: bool,
    /// Key derivation epoch (incremented on each re-key).
    pub key_index: u64,
}

impl Session {
    /// Build session state from `K_master` and handshake parameters (matches peer’s `derive_initial_seed`).
    pub fn from_master_secret(
        id: u64,
        master_key: &[u8; 32],
        timestamp: u64,
        session_nonce: &[u8],
    ) -> crate::error::Result<Self> {
        let seed =
            KeyDerivation::derive_initial_seed(master_key.as_slice(), timestamp, session_nonce)?;
        let data_key = KeyDerivation::derive_data_key(&seed, 0)?;
        Ok(Self::new(id, seed, data_key))
    }

    /// Create a new session from handshake results.
    pub fn new(id: u64, seed: [u8; 32], data_key: [u8; 32]) -> Self {
        Self {
            id,
            rng: SeedRng::new(seed),
            data_key,
            packet_counter: 0,
            active: true,
            key_index: 0,
        }
    }

    /// Increment the packet counter and return the new value.
    pub fn next_packet_counter(&mut self) -> u64 {
        self.packet_counter += 1;
        self.packet_counter
    }

    /// Rotate the data key using KDF.
    ///
    /// Derives `data_key[key_index+1] = HKDF(seed, key_index+1)` and increments
    /// `key_index`. Returns the new key.
    pub fn rekey(&mut self) -> crate::error::Result<[u8; 32]> {
        self.key_index += 1;
        self.data_key = KeyDerivation::derive_data_key(&self.rng.seed_bytes(), self.key_index)?;
        Ok(self.data_key)
    }

    /// Mark the session as closed.
    pub fn close(&mut self) {
        self.active = false;
    }

    /// Encrypt application payload using the parallel AEAD pool and a KDF-derived nonce.
    ///
    /// Advances [`Session::packet_counter`] and derives `nonce = HKDF(seed ‖ counter)`.
    pub fn encrypt_with_pipeline(
        &mut self,
        pipeline: &AeadPipeline,
        plaintext: &[u8],
    ) -> crate::error::Result<Vec<u8>> {
        let counter = self.next_packet_counter();
        let nonce = KeyDerivation::derive_nonce(&self.rng.seed_bytes(), counter)?;
        pipeline.encrypt(nonce, plaintext.to_vec())
    }

    /// Decrypt using an explicit nonce (e.g. derived on receive from frame metadata or inner payload).
    pub fn decrypt_with_pipeline(
        &self,
        pipeline: &AeadPipeline,
        nonce: [u8; 12],
        ciphertext: Vec<u8>,
    ) -> crate::error::Result<Vec<u8>> {
        pipeline.decrypt(nonce, ciphertext)
    }
}

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

    #[test]
    fn from_master_secret_derivation() {
        let master = [0x77u8; 32];
        let s = Session::from_master_secret(1, &master, 1_700_000_000, b"nonce").unwrap();
        assert_eq!(s.id, 1);
        assert!(s.active);
    }

    #[test]
    fn encrypt_decrypt_via_pipeline_matches_kdf_nonce() {
        let seed = [0xABu8; 32];
        let key = [0xCDu8; 32];
        let pipe = AeadPipeline::new(Variant::ChaCha20Poly1305, &key, 3).unwrap();
        let mut session = Session::new(9, seed, key);

        let ct = session.encrypt_with_pipeline(&pipe, b"payload").unwrap();
        let nonce =
            KeyDerivation::derive_nonce(&session.rng.seed_bytes(), session.packet_counter).unwrap();
        let pt = session.decrypt_with_pipeline(&pipe, nonce, ct).unwrap();
        assert_eq!(pt.as_slice(), b"payload");
    }
}