arcly-stream 0.8.3

//! SRT Key Material (KMREQ/KMRSP) message and AES-CTR media encryption
//! (feature `srt-encrypt`, draft-sharabayko-srt §6).
//!
//! A passphrase is stretched to a Key Encrypting Key (KEK) with PBKDF2-HMAC-SHA1
//! (2048 iterations, salt = the low 8 bytes of the 128-bit KM salt). The random
//! Stream Encrypting Key (SEK) is wrapped under the KEK with AES Key Wrap and
//! carried in the **Key Material** message; the listener unwraps it with the
//! same passphrase. Media payloads are then AES-CTR'd with a per-packet IV built
//! from the salt and the packet's sequence number.
//!
//! Scope: one SEK (the *even* key, `KK=1`); periodic key regeneration / odd-key
//! rotation is a follow-up. Interoperates with this crate's own caller/listener.

use super::crypto::{aes_key_unwrap, aes_key_wrap, pbkdf2_hmac_sha1, Aes};

const KM_VERSION: u8 = 1;
const KM_PT_KEYMAT: u8 = 2;
const KM_SIGN: u16 = 0x2029; // HAIVISION vendor signature
const CIPHER_AES_CTR: u8 = 2;
const SE_SRT: u8 = 2;
const KK_EVEN: u8 = 1;
const PBKDF2_ITERS: u32 = 2048;
const SALT_LEN: usize = 16;
/// Bytes of the salt fed to PBKDF2 (its low half), per libsrt.
const SALT_PBKDF_LEN: usize = 8;

/// The negotiated key material for an encrypted SRT session.
#[derive(Clone)]
pub(crate) struct KeyMaterial {
    sek: Vec<u8>,
    salt: [u8; SALT_LEN],
}

impl KeyMaterial {
    /// Generate fresh key material for `key_len` (16/24/32) under `passphrase`.
    pub(crate) fn generate(key_len: usize) -> KeyMaterial {
        let mut sek = vec![0u8; key_len];
        random_bytes(&mut sek);
        let mut salt = [0u8; SALT_LEN];
        random_bytes(&mut salt);
        KeyMaterial { sek, salt }
    }

    fn kek(&self, passphrase: &[u8]) -> Vec<u8> {
        pbkdf2_hmac_sha1(
            passphrase,
            &self.salt[SALT_LEN - SALT_PBKDF_LEN..],
            PBKDF2_ITERS,
            self.sek.len(),
        )
    }

    /// Serialize a Key Material message wrapping the SEK under `passphrase`
    /// (the body of a `KMREQ`/`KMRSP` handshake extension).
    pub(crate) fn to_message(&self, passphrase: &[u8]) -> Vec<u8> {
        let kek = self.kek(passphrase);
        let wrapped = aes_key_wrap(&kek, &self.sek);
        let mut m = Vec::with_capacity(16 + SALT_LEN + wrapped.len());
        m.push((KM_VERSION << 4) | KM_PT_KEYMAT); // 0 | Vers(3) | PT(4)
        m.extend_from_slice(&KM_SIGN.to_be_bytes());
        m.push(KK_EVEN); // Resv1(6) | KK(2)
        m.extend_from_slice(&0u32.to_be_bytes()); // KEKI
        m.push(CIPHER_AES_CTR);
        m.push(0); // Auth = none
        m.push(SE_SRT);
        m.push(0); // Resv2
        m.extend_from_slice(&0u16.to_be_bytes()); // Resv3
        m.push((SALT_LEN / 4) as u8); // SLen/4
        m.push((self.sek.len() / 4) as u8); // KLen/4
        m.extend_from_slice(&self.salt);
        m.extend_from_slice(&wrapped);
        m
    }

    /// Parse a Key Material message and unwrap its SEK with `passphrase`.
    /// Returns `None` on a malformed message or a passphrase mismatch (the
    /// AES-Key-Wrap integrity check fails).
    pub(crate) fn parse(passphrase: &[u8], m: &[u8]) -> Option<KeyMaterial> {
        if m.len() < 16 + SALT_LEN {
            return None;
        }
        if m[0] != (KM_VERSION << 4) | KM_PT_KEYMAT {
            return None;
        }
        if u16::from_be_bytes([m[1], m[2]]) != KM_SIGN {
            return None;
        }
        if m[8] != CIPHER_AES_CTR {
            return None;
        }
        let slen = (m[14] as usize) * 4;
        let klen = (m[15] as usize) * 4;
        if slen != SALT_LEN || !(klen == 16 || klen == 24 || klen == 32) {
            return None;
        }
        let salt: [u8; SALT_LEN] = m.get(16..16 + SALT_LEN)?.try_into().ok()?;
        let wrapped = m.get(16 + SALT_LEN..)?;
        if wrapped.len() != klen + 8 {
            return None;
        }
        // KEK depends only on the passphrase, salt, and key length.
        let probe = KeyMaterial {
            sek: vec![0u8; klen],
            salt,
        };
        let kek = probe.kek(passphrase);
        let sek = aes_key_unwrap(&kek, wrapped)?;
        Some(KeyMaterial { sek, salt })
    }

    /// AES-CTR (en/de)crypt `payload` in place for the packet with sequence
    /// number `seq`. The op is symmetric — the same call decrypts.
    pub(crate) fn transform(&self, seq: u32, payload: &mut [u8]) {
        // Per-packet IV: zeros with the packet index at bytes 10..14, XOR'd with
        // the MSB 112 bits (first 14 bytes) of the salt. The low 2 bytes are the
        // AES-CTR block counter (a single packet stays well under one carry).
        let mut iv = [0u8; 16];
        iv[10..14].copy_from_slice(&seq.to_be_bytes());
        for (b, s) in iv.iter_mut().zip(self.salt.iter()).take(14) {
            *b ^= *s;
        }
        Aes::new(&self.sek).ctr_xor(&iv, payload);
    }
}

/// Fill `out` with OS-seeded random bytes. Each `RandomState` pulls fresh
/// entropy from the platform RNG for its hasher keys, so the finished hash of a
/// per-byte-offset counter (mixed with the current time) inherits that entropy
/// — dependency-free, no `unsafe`.
fn random_bytes(out: &mut [u8]) {
    use std::collections::hash_map::RandomState;
    use std::hash::{BuildHasher, Hasher};
    use std::time::{SystemTime, UNIX_EPOCH};

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let mut i = 0;
    while i < out.len() {
        let mut h = RandomState::new().build_hasher();
        h.write_usize(i);
        h.write_u128(now);
        let v = h.finish().to_le_bytes();
        let n = (out.len() - i).min(8);
        out[i..i + n].copy_from_slice(&v[..n]);
        i += n;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn round_trips_key_material_with_correct_passphrase() {
        let pass = b"correcthorsebatterystaple";
        let km = KeyMaterial::generate(16);
        let msg = km.to_message(pass);
        let recovered = KeyMaterial::parse(pass, &msg).expect("unwrap with right passphrase");
        assert_eq!(recovered.sek, km.sek);
        assert_eq!(recovered.salt, km.salt);
    }

    #[test]
    fn rejects_wrong_passphrase() {
        let km = KeyMaterial::generate(32);
        let msg = km.to_message(b"the-real-secret");
        assert!(KeyMaterial::parse(b"a-different-secret", &msg).is_none());
    }

    #[test]
    fn ctr_round_trips_and_is_seq_bound() {
        let km = KeyMaterial::generate(16);
        let plain = b"MPEG-TS payload riding inside an SRT data packet".to_vec();

        let mut enc = plain.clone();
        km.transform(42, &mut enc);
        assert_ne!(enc, plain, "payload is actually encrypted");

        let mut dec = enc.clone();
        km.transform(42, &mut dec);
        assert_eq!(dec, plain, "same seq decrypts");

        // Decrypting under the wrong sequence number must NOT recover plaintext.
        let mut wrong = enc.clone();
        km.transform(43, &mut wrong);
        assert_ne!(wrong, plain);
    }

    #[test]
    fn generate_is_random() {
        let a = KeyMaterial::generate(16);
        let b = KeyMaterial::generate(16);
        assert_ne!(a.sek, b.sek, "SEKs differ across sessions");
        assert_ne!(a.salt, b.salt, "salts differ across sessions");
    }

    #[test]
    fn parses_at_all_key_lengths() {
        for &klen in &[16usize, 24, 32] {
            let pass = b"klen-test";
            let km = KeyMaterial::generate(klen);
            let msg = km.to_message(pass);
            let got = KeyMaterial::parse(pass, &msg).unwrap();
            assert_eq!(got.sek.len(), klen);
        }
    }
}