arcly-stream 0.8.3

//! SRT handshake parsing and a minimal listener-side responder.
//!
//! The handshake control payload follows the 16-byte packet header and carries
//! the version, encryption/extension fields, sequence number, MTU/window sizes,
//! the [request type][HandshakeType], the caller's socket id, and a SYN cookie
//! (draft-sharabayko-srt §3.2.1).

/// SRT handshake request type (the 32-bit `Handshake Type` field).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum HandshakeType {
    /// `INDUCTION` (1) — caller's first packet; listener replies with a cookie.
    Induction,
    /// `CONCLUSION` (0xFFFFFFFF) — caller echoes the cookie to finish setup.
    Conclusion,
    /// `WAVEAHAND` (0) — rendezvous probe.
    WaveAHand,
    /// `AGREEMENT` (0xFFFFFFFE) — rendezvous agreement.
    Agreement,
    /// Any other / error code.
    Other(u32),
}

impl HandshakeType {
    fn from_u32(v: u32) -> HandshakeType {
        match v {
            1 => HandshakeType::Induction,
            0xFFFF_FFFF => HandshakeType::Conclusion,
            0 => HandshakeType::WaveAHand,
            0xFFFF_FFFE => HandshakeType::Agreement,
            other => HandshakeType::Other(other),
        }
    }

    fn to_u32(self) -> u32 {
        match self {
            HandshakeType::Induction => 1,
            HandshakeType::Conclusion => 0xFFFF_FFFF,
            HandshakeType::WaveAHand => 0,
            HandshakeType::Agreement => 0xFFFF_FFFE,
            HandshakeType::Other(v) => v,
        }
    }
}

/// The fields of an SRT handshake this listener inspects.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SrtHandshake {
    /// Protocol version (4 for HSv4, 5 for HSv5).
    pub version: u32,
    /// Encryption field — non-zero means the caller wants an encrypted link,
    /// which this build rejects (see the module-level scope note).
    pub encryption: u16,
    /// Initial packet sequence number.
    pub initial_sequence: u32,
    /// The handshake request type.
    pub handshake_type: HandshakeType,
    /// The caller's SRT socket id.
    pub socket_id: u32,
    /// SYN cookie (0 on the caller's induction request).
    pub cookie: u32,
}

impl SrtHandshake {
    /// Offset of the handshake payload within a control datagram (past the
    /// 16-byte SRT packet header).
    const PAYLOAD: usize = 16;

    /// Parse the handshake from a full control datagram. Returns `None` if the
    /// datagram is too short to contain a handshake body.
    pub fn parse(datagram: &[u8]) -> Option<SrtHandshake> {
        let b = datagram.get(Self::PAYLOAD..)?;
        if b.len() < 32 {
            return None;
        }
        let w = |i: usize| u32::from_be_bytes([b[i], b[i + 1], b[i + 2], b[i + 3]]);
        Some(SrtHandshake {
            version: w(0),
            encryption: u16::from_be_bytes([b[4], b[5]]),
            initial_sequence: w(8),
            handshake_type: HandshakeType::from_u32(w(20)),
            socket_id: w(24),
            cookie: w(28),
        })
    }

    /// Whether the caller requested an encrypted link (unsupported here).
    pub fn wants_encryption(&self) -> bool {
        self.encryption != 0
    }
}

/// SRT `SRT_CMD_SID` (Stream ID) HSv5 extension type.
const EXT_SID: u16 = 5;

/// The Stream ID a caller advertised in its HSv5 `CONCLUSION` handshake, decoded
/// from the `SRT_CMD_SID` extension block, or `None` if absent.
///
/// `streamid` is libsrt's standard way for an SRT caller to select which stream
/// it publishes (a plain name, or the `#!::r=…,m=…` access-control form). It is
/// carried as 32-bit little-endian words with the characters byte-swapped within
/// each word; this reverses that encoding and trims the trailing NUL padding back
/// to the original string. Walking the post-body extension blocks mirrors the
/// `srt-encrypt` KMREQ locator, so it interoperates with libsrt callers that send
/// a Stream ID rather than only this crate's own caller.
pub fn stream_id(datagram: &[u8]) -> Option<String> {
    let mut p = SrtHandshake::PAYLOAD + 32; // past header + 32-byte body
    while p + 4 <= datagram.len() {
        let ext_type = u16::from_be_bytes([datagram[p], datagram[p + 1]]);
        let words = u16::from_be_bytes([datagram[p + 2], datagram[p + 3]]) as usize;
        let end = p + 4 + words * 4;
        if end > datagram.len() {
            break;
        }
        if ext_type == EXT_SID {
            let sid = decode_sid(&datagram[p + 4..end]);
            return (!sid.is_empty()).then_some(sid);
        }
        p = end;
    }
    None
}

/// Decode libsrt's byte-swapped-per-word Stream ID encoding into a string. The
/// contents are word-aligned (NUL-padded), so each 4-byte group is reversed and
/// the trailing NUL padding is trimmed.
fn decode_sid(contents: &[u8]) -> String {
    let mut out = Vec::with_capacity(contents.len());
    for word in contents.chunks(4) {
        out.extend(word.iter().rev());
    }
    while out.last() == Some(&0) {
        out.pop();
    }
    String::from_utf8_lossy(&out).into_owned()
}

/// A deterministic-but-opaque SYN cookie derived from the caller's socket id.
/// A production listener would mix in the peer address and a per-boot secret;
/// for an unencrypted loss-tolerant ingest the anti-spoofing value is modest, so
/// a stable derivation keeps the responder dependency-free.
fn syn_cookie(socket_id: u32) -> u32 {
    socket_id
        .rotate_left(13)
        .wrapping_mul(0x9E37_79B1)
        .wrapping_add(0x5247_5421)
}

/// Build a listener response to a caller's handshake datagram, or `None` if the
/// datagram is not a parseable handshake or requests encryption.
///
/// For an `INDUCTION` request the response echoes the body with a freshly
/// derived SYN cookie installed; for `CONCLUSION` it echoes the body to confirm
/// agreement. The returned bytes are ready to send back to the caller.
pub fn respond(datagram: &[u8]) -> Option<Vec<u8>> {
    let hs = SrtHandshake::parse(datagram)?;
    if hs.wants_encryption() {
        return None; // encrypted SRT is out of scope; drop the handshake
    }
    let mut reply = datagram.to_vec();
    let cookie = match hs.handshake_type {
        HandshakeType::Induction => syn_cookie(hs.socket_id),
        HandshakeType::Conclusion => hs.cookie,
        _ => return None,
    };
    // Install the cookie at its offset (payload + 28).
    let at = SrtHandshake::PAYLOAD + 28;
    reply
        .get_mut(at..at + 4)?
        .copy_from_slice(&cookie.to_be_bytes());
    Some(reply)
}

/// Build a caller-side handshake control datagram (the **egress relay** uses
/// this to dial an SRT listener). Mirrors the 32-byte HSv5 body this crate's
/// [`respond`] listener understands.
///
/// Note: this is the minimal handshake subset (no `HSREQ`/`KMREQ` extension
/// blocks), so it interoperates with this crate's own listener; full extension
/// interop with third-party SRT stacks is a follow-up.
pub fn caller_handshake(
    req_type: HandshakeType,
    socket_id: u32,
    initial_sequence: u32,
    cookie: u32,
) -> Vec<u8> {
    let mut d = vec![0u8; 16]; // control header: control flag set, type 0x0000
    d[0] = 0x80;
    let mut body = vec![0u8; 32];
    body[0..4].copy_from_slice(&5u32.to_be_bytes()); // version HSv5
                                                     // bytes 4..8: encryption(0) + extension(0)
    body[8..12].copy_from_slice(&initial_sequence.to_be_bytes());
    body[12..16].copy_from_slice(&1500u32.to_be_bytes()); // MTU
    body[16..20].copy_from_slice(&8192u32.to_be_bytes()); // flow-window
    body[20..24].copy_from_slice(&req_type.to_u32().to_be_bytes());
    body[24..28].copy_from_slice(&socket_id.to_be_bytes());
    body[28..32].copy_from_slice(&cookie.to_be_bytes());
    d.extend_from_slice(&body);
    d
}

/// The caller's first handshake (`INDUCTION`, cookie 0).
pub fn caller_induction(socket_id: u32, initial_sequence: u32) -> Vec<u8> {
    caller_handshake(HandshakeType::Induction, socket_id, initial_sequence, 0)
}

/// The caller's second handshake (`CONCLUSION`, echoing the listener's cookie).
pub fn caller_conclusion(socket_id: u32, initial_sequence: u32, cookie: u32) -> Vec<u8> {
    caller_handshake(
        HandshakeType::Conclusion,
        socket_id,
        initial_sequence,
        cookie,
    )
}

// ---------------------------------------------------------------------------
// HSv5 KMREQ/KMRSP extension blocks (feature `srt-encrypt`).
//
// After the 32-byte HSv5 body, extension blocks follow as
// `ExtType(16) | ExtLen(16 words) | ExtContents`. The encrypted handshake
// carries the Key Material message in a KMREQ block (caller → listener); the
// listener echoes it as KMRSP to confirm.
// ---------------------------------------------------------------------------

/// SRT `SRT_CMD_KMREQ` extension type.
#[cfg(feature = "srt-encrypt")]
const EXT_KMREQ: u16 = 3;
/// SRT `SRT_CMD_KMRSP` extension type.
#[cfg(feature = "srt-encrypt")]
const EXT_KMRSP: u16 = 4;
/// Extension-field flag advertising a KMREQ block follows.
#[cfg(feature = "srt-encrypt")]
const HS_EXT_KM_FLAG: u16 = 0x0002;

/// The HSv5 "Encryption Field" code advertising the SEK length (2/3/4 for
/// AES-128/192/256), i.e. `key_len / 8`.
#[cfg(feature = "srt-encrypt")]
pub fn encryption_field(key_len: usize) -> u16 {
    (key_len / 8) as u16
}

/// Locate the KMREQ/KMRSP extension block's offset (the start of its `ExtType`
/// field) within a handshake datagram, if present.
#[cfg(feature = "srt-encrypt")]
fn km_ext_offset(datagram: &[u8]) -> Option<usize> {
    let mut p = SrtHandshake::PAYLOAD + 32; // past header + 32-byte body
    while p + 4 <= datagram.len() {
        let ext_type = u16::from_be_bytes([datagram[p], datagram[p + 1]]);
        let words = u16::from_be_bytes([datagram[p + 2], datagram[p + 3]]) as usize;
        let end = p + 4 + words * 4;
        if end > datagram.len() {
            break;
        }
        if ext_type == EXT_KMREQ || ext_type == EXT_KMRSP {
            return Some(p);
        }
        p = end;
    }
    None
}

/// The Key Material message bytes carried in a KMREQ/KMRSP block, if any.
#[cfg(feature = "srt-encrypt")]
pub fn km_extension(datagram: &[u8]) -> Option<&[u8]> {
    let p = km_ext_offset(datagram)?;
    let words = u16::from_be_bytes([datagram[p + 2], datagram[p + 3]]) as usize;
    datagram.get(p + 4..p + 4 + words * 4)
}

/// Build a caller `CONCLUSION` that requests encryption: it advertises the
/// cipher in the Encryption Field, flags a KMREQ extension, and appends the
/// Key Material message `km`.
#[cfg(feature = "srt-encrypt")]
pub fn caller_conclusion_encrypted(
    socket_id: u32,
    initial_sequence: u32,
    cookie: u32,
    key_len: usize,
    km: &[u8],
) -> Vec<u8> {
    let mut d = caller_conclusion(socket_id, initial_sequence, cookie);
    let body = SrtHandshake::PAYLOAD;
    // Encryption Field (body bytes 4..6) and Extension Field (6..8).
    d[body + 4..body + 6].copy_from_slice(&encryption_field(key_len).to_be_bytes());
    d[body + 6..body + 8].copy_from_slice(&HS_EXT_KM_FLAG.to_be_bytes());
    // KMREQ extension block: type | length-in-words | contents.
    d.extend_from_slice(&EXT_KMREQ.to_be_bytes());
    d.extend_from_slice(&((km.len() / 4) as u16).to_be_bytes());
    d.extend_from_slice(km);
    d
}

/// Listener response that also processes encryption. For `INDUCTION` it installs
/// the SYN cookie (encrypted or not). For an encrypted `CONCLUSION` it unwraps
/// the KMREQ with `passphrase`, returns the recovered key material, and echoes
/// the handshake as a KMRSP; an unencrypted conclusion behaves like [`respond`].
/// Returns `None` if the handshake is malformed or the passphrase is wrong.
#[cfg(feature = "srt-encrypt")]
pub fn respond_with_km(
    datagram: &[u8],
    passphrase: &[u8],
) -> Option<(Vec<u8>, Option<super::keymaterial::KeyMaterial>)> {
    let hs = SrtHandshake::parse(datagram)?;
    match hs.handshake_type {
        HandshakeType::Induction => {
            let mut reply = datagram.to_vec();
            let at = SrtHandshake::PAYLOAD + 28;
            reply
                .get_mut(at..at + 4)?
                .copy_from_slice(&syn_cookie(hs.socket_id).to_be_bytes());
            Some((reply, None))
        }
        HandshakeType::Conclusion if hs.wants_encryption() => {
            let km_bytes = km_extension(datagram)?;
            let km = super::keymaterial::KeyMaterial::parse(passphrase, km_bytes)?;
            // Echo as KMRSP: flip the extension type in place to confirm.
            let mut reply = datagram.to_vec();
            let off = km_ext_offset(&reply)?;
            reply[off..off + 2].copy_from_slice(&EXT_KMRSP.to_be_bytes());
            Some((reply, Some(km)))
        }
        HandshakeType::Conclusion => Some((datagram.to_vec(), None)),
        _ => None,
    }
}

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

    /// Build a control datagram with a handshake body for the given type/enc.
    fn handshake_datagram(req_type: u32, encryption: u16) -> Vec<u8> {
        let mut d = vec![0u8; 16]; // control header
        d[0] = 0x80; // control flag
        let mut body = vec![0u8; 32];
        body[0..4].copy_from_slice(&5u32.to_be_bytes()); // version 5
        body[4..6].copy_from_slice(&encryption.to_be_bytes());
        body[20..24].copy_from_slice(&req_type.to_be_bytes());
        body[24..28].copy_from_slice(&0xABCD_1234u32.to_be_bytes()); // socket id
        d.extend_from_slice(&body);
        d
    }

    #[test]
    fn parses_induction_handshake() {
        let d = handshake_datagram(1, 0);
        let hs = SrtHandshake::parse(&d).unwrap();
        assert_eq!(hs.version, 5);
        assert_eq!(hs.handshake_type, HandshakeType::Induction);
        assert_eq!(hs.socket_id, 0xABCD_1234);
        assert!(!hs.wants_encryption());
    }

    #[test]
    fn induction_response_installs_nonzero_cookie() {
        let d = handshake_datagram(1, 0);
        let reply = respond(&d).unwrap();
        let parsed = SrtHandshake::parse(&reply).unwrap();
        assert_ne!(parsed.cookie, 0, "cookie installed in induction response");
    }

    #[test]
    fn encrypted_handshake_is_rejected() {
        let d = handshake_datagram(1, 0x0002);
        assert!(respond(&d).is_none());
    }

    #[test]
    fn non_handshake_request_type_has_no_response() {
        let d = handshake_datagram(0, 0); // WAVEAHAND
        assert!(respond(&d).is_none());
    }

    #[test]
    fn caller_handshake_loops_through_listener() {
        // Full in-process loopback: caller INDUCTION → listener cookie →
        // caller CONCLUSION → listener agreement, all over the wire format.
        let induction = caller_induction(0x0BAD_F00D, 42);
        let hs = SrtHandshake::parse(&induction).unwrap();
        assert_eq!(hs.handshake_type, HandshakeType::Induction);
        assert_eq!(hs.version, 5);
        assert_eq!(hs.socket_id, 0x0BAD_F00D);

        let resp = respond(&induction).expect("listener induction reply");
        let cookie = SrtHandshake::parse(&resp).unwrap().cookie;
        assert_ne!(cookie, 0, "listener installed a cookie");

        let conclusion = caller_conclusion(0x0BAD_F00D, 42, cookie);
        let chs = SrtHandshake::parse(&conclusion).unwrap();
        assert_eq!(chs.handshake_type, HandshakeType::Conclusion);
        assert_eq!(chs.cookie, cookie, "caller echoes the cookie");

        let agree = respond(&conclusion).expect("listener conclusion reply");
        assert_eq!(SrtHandshake::parse(&agree).unwrap().cookie, cookie);
    }

    /// Encode `s` the way libsrt frames a Stream ID: NUL-pad to a 4-byte boundary,
    /// then byte-swap each 32-bit word, and append it as an `SRT_CMD_SID` block.
    fn datagram_with_sid(s: &str) -> Vec<u8> {
        let mut bytes = s.as_bytes().to_vec();
        while bytes.len() % 4 != 0 {
            bytes.push(0);
        }
        let mut swapped = Vec::with_capacity(bytes.len());
        for word in bytes.chunks(4) {
            swapped.extend(word.iter().rev());
        }
        let mut d = handshake_datagram(0xFFFF_FFFF, 0); // CONCLUSION
        d.extend_from_slice(&EXT_SID.to_be_bytes());
        d.extend_from_slice(&((swapped.len() / 4) as u16).to_be_bytes());
        d.extend_from_slice(&swapped);
        d
    }

    #[test]
    fn decodes_stream_id_extension() {
        // Word-aligned and non-aligned (padded) lengths both round-trip.
        let d = datagram_with_sid("live/cam01");
        assert_eq!(stream_id(&d).as_deref(), Some("live/cam01"));

        let d = datagram_with_sid("abcd");
        assert_eq!(stream_id(&d).as_deref(), Some("abcd"));

        // A libsrt access-control streamid form survives verbatim.
        let d = datagram_with_sid("#!::r=live/feed,m=publish");
        assert_eq!(stream_id(&d).as_deref(), Some("#!::r=live/feed,m=publish"));
    }

    #[test]
    fn absent_stream_id_is_none() {
        // A plain handshake with no extension blocks has no Stream ID.
        let d = handshake_datagram(0xFFFF_FFFF, 0);
        assert_eq!(stream_id(&d), None);
        // A truncated extension header is ignored, not a panic.
        let mut d = handshake_datagram(0xFFFF_FFFF, 0);
        d.extend_from_slice(&EXT_SID.to_be_bytes()); // type but no length/contents
        assert_eq!(stream_id(&d), None);
    }
}