smartproxy/rust/crates/rustproxy-passthrough/src/sni_parser.rs

//! ClientHello SNI extraction via manual byte parsing.
//! No TLS stack needed - we just parse enough of the ClientHello to extract the SNI.

/// Result of SNI extraction.
#[derive(Debug)]
pub enum SniResult {
    /// Successfully extracted SNI hostname.
    Found(String),
    /// TLS ClientHello detected but no SNI extension present.
    NoSni,
    /// Not a TLS ClientHello (plain HTTP or other protocol).
    NotTls,
    /// Need more data to determine.
    NeedMoreData,
}

/// Extract the SNI hostname from a TLS ClientHello message.
///
/// This parses just enough of the TLS record to find the SNI extension,
/// without performing any actual TLS operations.
pub fn extract_sni(data: &[u8]) -> SniResult {
    // Minimum TLS record header is 5 bytes
    if data.len() < 5 {
        return SniResult::NeedMoreData;
    }

    // Check for TLS record: content_type=22 (Handshake)
    if data[0] != 0x16 {
        return SniResult::NotTls;
    }

    // TLS version (major.minor) - accept any
    // data[1..2] = version

    // Record length
    let record_len = ((data[3] as usize) << 8) | (data[4] as usize);
    let _total_len = 5 + record_len;

    // We need at least the handshake header (5 TLS + 4 handshake = 9)
    if data.len() < 9 {
        return SniResult::NeedMoreData;
    }

    // Handshake type = 1 (ClientHello)
    if data[5] != 0x01 {
        return SniResult::NotTls;
    }

    // Handshake length (3 bytes) - informational, we parse incrementally
    let _handshake_len = ((data[6] as usize) << 16)
        | ((data[7] as usize) << 8)
        | (data[8] as usize);

    let hello = &data[9..];

    // ClientHello structure:
    // 2 bytes: client version
    // 32 bytes: random
    // 1 byte: session_id length + session_id
    let mut pos = 2 + 32; // skip version + random

    if pos >= hello.len() {
        return SniResult::NeedMoreData;
    }

    // Session ID
    let session_id_len = hello[pos] as usize;
    pos += 1 + session_id_len;

    if pos + 2 > hello.len() {
        return SniResult::NeedMoreData;
    }

    // Cipher suites
    let cipher_suites_len = ((hello[pos] as usize) << 8) | (hello[pos + 1] as usize);
    pos += 2 + cipher_suites_len;

    if pos + 1 > hello.len() {
        return SniResult::NeedMoreData;
    }

    // Compression methods
    let compression_len = hello[pos] as usize;
    pos += 1 + compression_len;

    if pos + 2 > hello.len() {
        // No extensions
        return SniResult::NoSni;
    }

    // Extensions length
    let extensions_len = ((hello[pos] as usize) << 8) | (hello[pos + 1] as usize);
    pos += 2;

    let extensions_end = pos + extensions_len;
    if extensions_end > hello.len() {
        // Partial extensions, try to parse what we have
    }

    // Parse extensions looking for SNI (type 0x0000)
    while pos + 4 <= hello.len() && pos < extensions_end {
        let ext_type = ((hello[pos] as u16) << 8) | (hello[pos + 1] as u16);
        let ext_len = ((hello[pos + 2] as usize) << 8) | (hello[pos + 3] as usize);
        pos += 4;

        if ext_type == 0x0000 {
            // SNI extension
            return parse_sni_extension(&hello[pos..(pos + ext_len).min(hello.len())], ext_len);
        }

        pos += ext_len;
    }

    SniResult::NoSni
}

/// Parse the SNI extension data.
fn parse_sni_extension(data: &[u8], _ext_len: usize) -> SniResult {
    if data.len() < 5 {
        return SniResult::NeedMoreData;
    }

    // Server name list length
    let _list_len = ((data[0] as usize) << 8) | (data[1] as usize);

    // Server name type (0 = hostname)
    if data[2] != 0x00 {
        return SniResult::NoSni;
    }

    // Hostname length
    let name_len = ((data[3] as usize) << 8) | (data[4] as usize);

    if data.len() < 5 + name_len {
        return SniResult::NeedMoreData;
    }

    match std::str::from_utf8(&data[5..5 + name_len]) {
        Ok(hostname) => SniResult::Found(hostname.to_lowercase()),
        Err(_) => SniResult::NoSni,
    }
}

/// Check if the initial bytes look like a TLS ClientHello.
pub fn is_tls(data: &[u8]) -> bool {
    data.len() >= 3 && data[0] == 0x16 && data[1] == 0x03
}

/// Extract the HTTP request path from initial data.
/// E.g., from "GET /foo/bar HTTP/1.1\r\n..." returns Some("/foo/bar").
pub fn extract_http_path(data: &[u8]) -> Option<String> {
    let text = std::str::from_utf8(data).ok()?;
    // Find first space (after method)
    let method_end = text.find(' ')?;
    let rest = &text[method_end + 1..];
    // Find end of path (next space before "HTTP/...")
    let path_end = rest.find(' ').unwrap_or(rest.len());
    let path = &rest[..path_end];
    // Strip query string for path matching
    let path = path.split('?').next().unwrap_or(path);
    if path.starts_with('/') {
        Some(path.to_string())
    } else {
        None
    }
}

/// Extract the HTTP Host header from initial data.
/// E.g., from "GET / HTTP/1.1\r\nHost: example.com\r\n..." returns Some("example.com").
pub fn extract_http_host(data: &[u8]) -> Option<String> {
    let text = std::str::from_utf8(data).ok()?;
    for line in text.split("\r\n") {
        if let Some(value) = line.strip_prefix("Host: ").or_else(|| line.strip_prefix("host: ")) {
            // Strip port if present
            let host = value.split(':').next().unwrap_or(value).trim();
            if !host.is_empty() {
                return Some(host.to_lowercase());
            }
        }
    }
    None
}

/// Check if the initial bytes look like HTTP.
pub fn is_http(data: &[u8]) -> bool {
    if data.len() < 4 {
        return false;
    }
    // Check for common HTTP methods
    let starts = [
        b"GET " as &[u8],
        b"POST",
        b"PUT ",
        b"HEAD",
        b"DELE",
        b"PATC",
        b"OPTI",
        b"CONN",
    ];
    starts.iter().any(|s| data.starts_with(s))
}

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

    #[test]
    fn test_not_tls() {
        let http_data = b"GET / HTTP/1.1\r\n";
        assert!(matches!(extract_sni(http_data), SniResult::NotTls));
    }

    #[test]
    fn test_too_short() {
        assert!(matches!(extract_sni(&[0x16, 0x03]), SniResult::NeedMoreData));
    }

    #[test]
    fn test_is_tls() {
        assert!(is_tls(&[0x16, 0x03, 0x01]));
        assert!(!is_tls(&[0x47, 0x45, 0x54])); // "GET"
    }

    #[test]
    fn test_is_http() {
        assert!(is_http(b"GET /"));
        assert!(is_http(b"POST /api"));
        assert!(!is_http(&[0x16, 0x03, 0x01]));
    }

    #[test]
    fn test_real_client_hello() {
        // A minimal TLS 1.2 ClientHello with SNI "example.com"
        let client_hello: Vec<u8> = build_test_client_hello("example.com");
        match extract_sni(&client_hello) {
            SniResult::Found(sni) => assert_eq!(sni, "example.com"),
            other => panic!("Expected Found, got {:?}", other),
        }
    }

    /// Build a minimal TLS ClientHello for testing.
    fn build_test_client_hello(hostname: &str) -> Vec<u8> {
        let hostname_bytes = hostname.as_bytes();

        // SNI extension
        let sni_ext_data = {
            let mut d = Vec::new();
            // Server name list length
            let name_entry_len = 3 + hostname_bytes.len(); // type(1) + len(2) + name
            d.push(((name_entry_len >> 8) & 0xFF) as u8);
            d.push((name_entry_len & 0xFF) as u8);
            // Host name type = 0
            d.push(0x00);
            // Host name length
            d.push(((hostname_bytes.len() >> 8) & 0xFF) as u8);
            d.push((hostname_bytes.len() & 0xFF) as u8);
            // Host name
            d.extend_from_slice(hostname_bytes);
            d
        };

        // Extension: type=0x0000 (SNI), length, data
        let sni_extension = {
            let mut e = Vec::new();
            e.push(0x00); e.push(0x00); // SNI type
            e.push(((sni_ext_data.len() >> 8) & 0xFF) as u8);
            e.push((sni_ext_data.len() & 0xFF) as u8);
            e.extend_from_slice(&sni_ext_data);
            e
        };

        // Extensions block
        let extensions = {
            let mut ext = Vec::new();
            ext.push(((sni_extension.len() >> 8) & 0xFF) as u8);
            ext.push((sni_extension.len() & 0xFF) as u8);
            ext.extend_from_slice(&sni_extension);
            ext
        };

        // ClientHello body
        let hello_body = {
            let mut h = Vec::new();
            // Client version TLS 1.2
            h.push(0x03); h.push(0x03);
            // Random (32 bytes)
            h.extend_from_slice(&[0u8; 32]);
            // Session ID length = 0
            h.push(0x00);
            // Cipher suites: length=2, one suite
            h.push(0x00); h.push(0x02);
            h.push(0x00); h.push(0x2F); // TLS_RSA_WITH_AES_128_CBC_SHA
            // Compression methods: length=1, null
            h.push(0x01); h.push(0x00);
            // Extensions
            h.extend_from_slice(&extensions);
            h
        };

        // Handshake: type=1 (ClientHello), length
        let handshake = {
            let mut hs = Vec::new();
            hs.push(0x01); // ClientHello
            // 3-byte length
            hs.push(((hello_body.len() >> 16) & 0xFF) as u8);
            hs.push(((hello_body.len() >> 8) & 0xFF) as u8);
            hs.push((hello_body.len() & 0xFF) as u8);
            hs.extend_from_slice(&hello_body);
            hs
        };

        // TLS record: type=0x16, version TLS 1.0, length
        let mut record = Vec::new();
        record.push(0x16); // Handshake
        record.push(0x03); record.push(0x01); // TLS 1.0
        record.push(((handshake.len() >> 8) & 0xFF) as u8);
        record.push((handshake.len() & 0xFF) as u8);
        record.extend_from_slice(&handshake);

        record
    }
}