flowscope 0.22.0

//! TLS record + handshake parsing on top of `tls-parser`.

use bytes::Bytes;
use tls_parser::{
    TlsExtension, TlsMessage, TlsMessageAlert, TlsMessageHandshake, parse_tls_extensions,
    parse_tls_plaintext,
};

use super::types::{
    TlsAlert, TlsAlertLevel, TlsClientHello, TlsConfig, TlsServerHello, TlsVersion,
};
use crate::error::{Error, Module};

/// Per-direction parser state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum DirState {
    /// Awaiting more record bytes.
    Reading,
    /// ChangeCipherSpec observed; subsequent records are encrypted
    /// and we stop trying to parse handshake data.
    Encrypted,
    /// Parser irrecoverably desynced (malformed input).
    Desynced,
}

/// Output of one parse step.
#[derive(Debug)]
pub(crate) enum ParseOutput {
    ClientHello(Box<TlsClientHello>),
    ServerHello(Box<TlsServerHello>),
    Alert(TlsAlert),
    /// One TLS 1.2 `Certificate` handshake message. The chain
    /// is the server's (or client's, in mutual TLS) cert list,
    /// each entry an X.509 DER blob. TLS 1.3 carries the cert
    /// chain encrypted in EncryptedExtensions — not surfaced.
    Certificate {
        chain: Vec<Bytes>,
    },
}

/// Try to advance the parser by exactly one TLS record. Every
/// handshake message found inside the record is pushed onto
/// `out` (one record can carry several — e.g. TLS 1.2's
/// ServerHello + Certificate + ServerKeyExchange + ServerHelloDone
/// chain).
///
/// Returns:
/// - `Ok(true)` — a record was consumed; caller should re-call.
/// - `Ok(false)` — need more bytes (or transitioned to Encrypted).
/// - `Err(_)` — desync.
pub(crate) fn step(
    state: &mut DirState,
    buffer: &mut Vec<u8>,
    is_initiator: bool,
    config: &TlsConfig,
    out: &mut Vec<ParseOutput>,
) -> crate::Result<bool> {
    if buffer.len() > config.max_buffer {
        *state = DirState::Desynced;
        buffer.clear();
        return Err(Error::buffer_overflow(Module::Tls, config.max_buffer));
    }
    if matches!(*state, DirState::Encrypted | DirState::Desynced) {
        return Ok(false);
    }
    if buffer.len() < 5 {
        return Ok(false);
    }
    // Peek the record length without consuming.
    let record_len = u16::from_be_bytes([buffer[3], buffer[4]]) as usize;
    let total = 5 + record_len;
    if buffer.len() < total {
        return Ok(false);
    }

    // Parse one record. The returned TlsPlaintext borrows from
    // `buffer`; we build owned events before releasing the borrow
    // and mutating the buffer.
    let mut became_encrypted = false;
    {
        let plaintext = match parse_tls_plaintext(&buffer[..total]) {
            Ok((_, p)) => p,
            Err(e) => {
                let msg = format!("{e:?}");
                *state = DirState::Desynced;
                buffer.clear();
                return Err(Error::parse(Module::Tls, msg));
            }
        };
        let record_version = TlsVersion::from_raw(plaintext.hdr.version.0);

        for msg in &plaintext.msg {
            match msg {
                TlsMessage::Handshake(h) => match h {
                    TlsMessageHandshake::ClientHello(ch) if is_initiator => {
                        let ev = build_client_hello(record_version, ch);
                        out.push(ParseOutput::ClientHello(Box::new(ev)));
                    }
                    TlsMessageHandshake::ServerHello(sh) if !is_initiator => {
                        let ev = build_server_hello(record_version, sh);
                        out.push(ParseOutput::ServerHello(Box::new(ev)));
                    }
                    TlsMessageHandshake::Certificate(c) => {
                        let chain: Vec<Bytes> = c
                            .cert_chain
                            .iter()
                            .map(|raw| Bytes::copy_from_slice(raw.data))
                            .collect();
                        out.push(ParseOutput::Certificate { chain });
                    }
                    _ => {}
                },
                TlsMessage::Alert(alert) => {
                    out.push(ParseOutput::Alert(build_alert(alert)));
                }
                TlsMessage::ChangeCipherSpec => {
                    became_encrypted = true;
                }
                TlsMessage::ApplicationData(_) | TlsMessage::Heartbeat(_) => {}
            }
        }
    }

    if became_encrypted {
        *state = DirState::Encrypted;
    }

    // Consume the record we just processed.
    let rest = buffer.split_off(total);
    *buffer = rest;
    Ok(true)
}

/// Build a [`TlsClientHello`] from a `tls-parser` ClientHello.
///
/// Shared with the QUIC parser (issue #82): a QUIC Initial carries
/// the same TLS 1.3 ClientHello in its CRYPTO stream, so the JA4
/// fingerprint path reuses this conversion instead of duplicating
/// the extension walk.
pub(crate) fn build_client_hello(
    record_version: TlsVersion,
    ch: &tls_parser::TlsClientHelloContents<'_>,
) -> TlsClientHello {
    let legacy_version = TlsVersion::from_raw(ch.version.0);
    let mut random = [0u8; 32];
    let n = ch.random.len().min(32);
    random[..n].copy_from_slice(&ch.random[..n]);
    let session_id = ch
        .session_id
        .map(Bytes::copy_from_slice)
        .unwrap_or_default();
    let cipher_suites: Vec<u16> = ch.ciphers.iter().map(|c| c.0).collect();
    let compression = Bytes::copy_from_slice(&ch.comp.iter().map(|c| c.0).collect::<Vec<u8>>());

    let mut sni: Option<String> = None;
    let mut alpn: Vec<String> = Vec::new();
    let mut supported_versions: Vec<TlsVersion> = Vec::new();
    let mut supported_groups: Vec<u16> = Vec::new();
    let mut key_share_groups: Vec<u16> = Vec::new();
    let mut extension_types: Vec<u16> = Vec::new();
    let mut ech_present = false;
    let mut ech_config_id: Option<u8> = None;

    if let Some(ext_bytes) = ch.ext
        && let Ok((_, exts)) = parse_tls_extensions(ext_bytes)
    {
        for ext in &exts {
            let id = extension_id(ext);
            extension_types.push(id);
            match ext {
                TlsExtension::SNI(items) => {
                    for (_kind, host) in items {
                        if let Ok(s) = std::str::from_utf8(host) {
                            sni = Some(s.to_string());
                            break;
                        }
                    }
                }
                TlsExtension::ALPN(protos) => {
                    for p in protos {
                        if let Ok(s) = std::str::from_utf8(p) {
                            alpn.push(s.to_string());
                        }
                    }
                }
                TlsExtension::SupportedVersions(vs) => {
                    for v in vs {
                        supported_versions.push(TlsVersion::from_raw(v.0));
                    }
                }
                TlsExtension::EllipticCurves(groups) => {
                    for g in groups {
                        supported_groups.push(g.0);
                    }
                }
                // key_share (51) and its pre-draft-23 codepoint (40)
                // carry the actual public shares; a post-quantum
                // hybrid share is what bloats the ClientHello.
                TlsExtension::KeyShare(body) | TlsExtension::KeyShareOld(body) => {
                    parse_key_share_groups(body, &mut key_share_groups);
                }
                TlsExtension::Unknown(_, body) if id == 0xfe0d => {
                    // ECH ClientHello outer (RFC draft §5.1):
                    //   1B  ECHClientHelloType    (0=outer, 1=inner)
                    //   1B  HPKE config_id
                    //   2B  HPKE KDF id
                    //   2B  HPKE AEAD id
                    //   2B  enc.len + enc bytes
                    //   2B  payload.len + payload bytes
                    if let [ty, cid, _, ..] = **body
                        && ty == 0
                    {
                        ech_present = true;
                        ech_config_id = Some(cid);
                    }
                }
                _ => {}
            }
        }
    }

    // A PQ hybrid in the actual key_share is the definitive signal;
    // fall back to the offered supported_groups (a client that
    // key-shares a PQ group always lists it there too, but a client
    // may advertise PQ capability without sharing yet).
    let pq_key_share = key_share_groups
        .iter()
        .chain(supported_groups.iter())
        .any(|g| crate::tls::is_pq_hybrid_group(*g));

    TlsClientHello {
        record_version,
        legacy_version,
        random,
        session_id,
        cipher_suites,
        compression,
        sni,
        alpn,
        supported_versions,
        supported_groups,
        extension_types,
        ech_present,
        ech_config_id,
        sni_is_outer: ech_present,
        key_share_groups,
        pq_key_share,
    }
}

/// Parse the `key_share` extension body (ClientHello form,
/// RFC 8446 §4.2.8) and push each entry's named group into `out`.
///
/// Body layout: `2B client_shares_len` then repeated
/// `2B group | 2B key_exchange_len | key_exchange`. Malformed /
/// truncated input stops the walk without erroring — the caller
/// treats a partial parse as "what we could read".
fn parse_key_share_groups(body: &[u8], out: &mut Vec<u16>) {
    // Skip the 2-byte client_shares vector length prefix.
    let Some(mut rest) = body.get(2..) else {
        return;
    };
    while rest.len() >= 4 {
        let group = u16::from_be_bytes([rest[0], rest[1]]);
        let ke_len = u16::from_be_bytes([rest[2], rest[3]]) as usize;
        out.push(group);
        let advance = 4 + ke_len;
        if rest.len() < advance {
            break;
        }
        rest = &rest[advance..];
    }
}

fn build_server_hello(
    record_version: TlsVersion,
    sh: &tls_parser::TlsServerHelloContents<'_>,
) -> TlsServerHello {
    let legacy_version = TlsVersion::from_raw(sh.version.0);
    let mut random = [0u8; 32];
    let n = sh.random.len().min(32);
    random[..n].copy_from_slice(&sh.random[..n]);
    let session_id = sh
        .session_id
        .map(Bytes::copy_from_slice)
        .unwrap_or_default();
    let cipher_suite = sh.cipher.0;
    let compression = sh.compression.0;

    let mut alpn: Option<String> = None;
    let mut supported_version: Option<TlsVersion> = None;
    // Extension order as seen on the wire (JA4S hashes them as-is).
    let mut extension_types: Vec<u16> = Vec::new();

    if let Some(ext_bytes) = sh.ext
        && let Ok((_, exts)) = parse_tls_extensions(ext_bytes)
    {
        for ext in &exts {
            extension_types.push(extension_id(ext));
            match ext {
                TlsExtension::ALPN(protos) => {
                    if let Some(p) = protos.first()
                        && let Ok(s) = std::str::from_utf8(p)
                    {
                        alpn = Some(s.to_string());
                    }
                }
                TlsExtension::SupportedVersions(vs) => {
                    if let Some(v) = vs.first() {
                        supported_version = Some(TlsVersion::from_raw(v.0));
                    }
                }
                _ => {}
            }
        }
    }

    TlsServerHello {
        record_version,
        legacy_version,
        random,
        session_id,
        cipher_suite,
        compression,
        alpn,
        supported_version,
        extension_types,
        // ECH retry_configs sit inside the encrypted
        // EncryptedExtensions; the plaintext-only observer can't
        // see them. Best-effort default `false` — consumers
        // wanting fidelity pair with a server-side keyed
        // implementation. Plan 144.
        ech_retry_configs: false,
    }
}

fn build_alert(alert: &TlsMessageAlert) -> TlsAlert {
    let level = match alert.severity.0 {
        1 => TlsAlertLevel::Warning,
        2 => TlsAlertLevel::Fatal,
        v => TlsAlertLevel::Other(v),
    };
    TlsAlert {
        level,
        description: alert.code.0,
    }
}

/// Map a TlsExtension to its IANA type ID. Used for fingerprinting
/// (JA3, JA4) and for the public `extension_types` field.
fn extension_id(ext: &TlsExtension<'_>) -> u16 {
    match ext {
        TlsExtension::SNI(_) => 0,
        TlsExtension::MaxFragmentLength(_) => 1,
        TlsExtension::StatusRequest(_) => 5,
        TlsExtension::EllipticCurves(_) => 10,
        TlsExtension::EcPointFormats(_) => 11,
        TlsExtension::SignatureAlgorithms(_) => 13,
        TlsExtension::Heartbeat(_) => 15,
        TlsExtension::ALPN(_) => 16,
        TlsExtension::SignedCertificateTimestamp(_) => 18,
        TlsExtension::Padding(_) => 21,
        TlsExtension::EncryptThenMac => 22,
        TlsExtension::ExtendedMasterSecret => 23,
        TlsExtension::SessionTicket(_) => 35,
        TlsExtension::PreSharedKey(_) => 41,
        TlsExtension::EarlyData(_) => 42,
        TlsExtension::SupportedVersions(_) => 43,
        TlsExtension::Cookie(_) => 44,
        TlsExtension::PskExchangeModes(_) => 45,
        TlsExtension::KeyShare(_) => 51,
        TlsExtension::NextProtocolNegotiation => 13172,
        TlsExtension::RenegotiationInfo(_) => 65281,
        TlsExtension::EncryptedServerName { .. } => 65486,
        TlsExtension::Grease(..) => 0x0a0a,
        // Unknown extensions: tls-parser keeps the raw type via Unknown.
        TlsExtension::Unknown(t, _) => t.0,
        // Variants we don't track explicitly map to their best-known
        // IANA IDs; for any new variants tls-parser adds, fall back to
        // a sentinel. Users who need exact IDs should walk extensions
        // themselves via parse_tls_extensions.
        _ => u16::MAX,
    }
}