rustpython-host_env 0.6.0

//! NSS-style TLS key log sink used by `SSLContext.keylog_filename`.

use alloc::sync::Arc;

use crate::fs::AppendLog;

pub const HEADER: &[u8] = b"# TLS secrets log file, generated by OpenSSL / Python\n";

/// Context-owned sink so existing and cached rustls configurations observe
/// changes to keylog_filename without rebuilding their session state.
#[derive(Debug, Default)]
pub struct KeyLog(pub AppendLog);

/// Each connection follows its current context. This routing state must not be
/// shared through the context's config cache or an explicit SSLSession.
#[derive(Debug)]
pub struct ConnectionKeyLog {
    sink: parking_lot::RwLock<Arc<KeyLog>>,
    master_secret: parking_lot::Mutex<Option<Vec<u8>>>,
}

impl ConnectionKeyLog {
    pub fn new(sink: Arc<KeyLog>) -> Self {
        Self {
            sink: parking_lot::RwLock::new(sink),
            master_secret: parking_lot::Mutex::new(None),
        }
    }

    pub fn set_sink(&self, sink: Arc<KeyLog>) {
        *self.sink.write() = sink;
    }

    pub fn master_secret(&self) -> Option<Vec<u8>> {
        self.master_secret.lock().clone()
    }
}

impl rustls::KeyLog for ConnectionKeyLog {
    fn log(&self, label: &str, client_random: &[u8], secret: &[u8]) {
        if label == "CLIENT_RANDOM" {
            *self.master_secret.lock() = Some(secret.to_vec());
        }
        self.sink.read().log(label, client_random, secret);
    }

    fn will_log(&self, label: &str) -> bool {
        label == "CLIENT_RANDOM" || self.sink.read().will_log(label)
    }
}

fn format_record(label: &str, client_random: &[u8], secret: &[u8]) -> Vec<u8> {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut line = Vec::with_capacity(label.len() + 2 * (client_random.len() + secret.len()) + 3);
    line.extend_from_slice(label.as_bytes());
    for bytes in [client_random, secret] {
        line.push(b' ');
        for byte in bytes {
            line.push(HEX[(byte >> 4) as usize]);
            line.push(HEX[(byte & 0xf) as usize]);
        }
    }
    line.push(b'\n');
    line
}

impl rustls::KeyLog for KeyLog {
    fn log(&self, label: &str, client_random: &[u8], secret: &[u8]) {
        // rustls has no error return from this callback. Opening the file fails
        // eagerly in the setter; subsequent writes are best effort.
        let _ = self.0.write(&format_record(label, client_random, secret));
    }

    fn will_log(&self, _label: &str) -> bool {
        self.0.enabled()
    }
}

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

    #[test]
    fn nss_record_has_two_lowercase_hex_fields() {
        assert_eq!(
            format_record("CLIENT_RANDOM", &[0x00, 0x1f, 0xa0, 0xff], &[0x09, 0xbc]),
            b"CLIENT_RANDOM 001fa0ff 09bc\n"
        );
    }

    #[test]
    fn disabled_by_default() {
        use rustls::KeyLog as _;
        assert!(!KeyLog::default().will_log("CLIENT_RANDOM"));
    }
}