ssl: add key logging with a shared host append sink (#8655)

* ssl: add key logging with a shared host append sink

Port the context-owned keylog sink design from Pyre, retaining CPython 3.14 path and context-switch semantics. Keep host file lifecycle in host_env and TLS formatting in stdlib.

Assisted-by: Codex:GPT-5

* ssl: use core and alloc imports for key logging

Assisted-by: Codex:GPT-5

crates/host_env/src/fs.rs

@@ -4,6 +4,9 @@ use std::{
     path::{Path, PathBuf},
 };
 
+mod append_log;
+pub use append_log::AppendLog;
+
 pub fn open(path: impl AsRef<Path>) -> io::Result<File> {
     File::open(path)
 }

crates/host_env/src/fs/append_log.rs

@@ -0,0 +1,144 @@
+use parking_lot::Mutex;
+use std::{
+    fs::{File, OpenOptions},
+    io::{self, Write},
+    path::Path,
+};
+
+/// A replaceable append-only log destination, shared by its owner's consumers.
+/// File replacement and complete record writes are serialized per destination.
+#[derive(Debug, Default)]
+pub struct AppendLog {
+    file: Mutex<Option<File>>,
+}
+
+impl AppendLog {
+    /// Close the old destination before opening its replacement. On failure the
+    /// log stays disabled. The header is written only to empty regular files.
+    pub fn set_path(&self, path: Option<&Path>, header: &[u8]) -> io::Result<()> {
+        let mut slot = self.file.lock();
+        *slot = None;
+        let Some(path) = path else {
+            return Ok(());
+        };
+        let mut file = OpenOptions::new().create(true).append(true).open(path)?;
+        let metadata = file.metadata()?;
+        if metadata.is_file() && metadata.len() == 0 {
+            file.write_all(header)?;
+            file.flush()?;
+        }
+        *slot = Some(file);
+        Ok(())
+    }
+
+    pub fn enabled(&self) -> bool {
+        self.file.lock().is_some()
+    }
+
+    /// Append and flush a complete record, or do nothing if disabled.
+    pub fn write(&self, record: &[u8]) -> io::Result<()> {
+        if let Some(file) = self.file.lock().as_mut() {
+            file.write_all(record)?;
+            file.flush()?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use core::sync::atomic::{AtomicUsize, Ordering};
+    use std::path::PathBuf;
+
+    struct TestDir(PathBuf);
+
+    impl TestDir {
+        fn new() -> Self {
+            static NEXT: AtomicUsize = AtomicUsize::new(0);
+            loop {
+                let path = std::env::temp_dir().join(format!(
+                    "rustpython-append-log-{}-{}",
+                    std::process::id(),
+                    NEXT.fetch_add(1, Ordering::Relaxed)
+                ));
+                match std::fs::create_dir(&path) {
+                    Ok(()) => return Self(path),
+                    Err(e) if e.kind() == io::ErrorKind::AlreadyExists => continue,
+                    Err(e) => panic!("create test directory: {e}"),
+                }
+            }
+        }
+    }
+
+    impl Drop for TestDir {
+        fn drop(&mut self) {
+            std::fs::remove_dir_all(&self.0).unwrap();
+        }
+    }
+
+    #[test]
+    fn append_replace_and_disable() {
+        let dir = TestDir::new();
+        let first = dir.0.join("first");
+        let second = dir.0.join("second");
+        let log = AppendLog::default();
+        assert!(!log.enabled());
+        log.write(b"ignored\n").unwrap();
+        log.set_path(Some(&first), b"header\n").unwrap();
+        log.write(b"one\n").unwrap();
+        log.set_path(Some(&first), b"header\n").unwrap();
+        log.write(b"two\n").unwrap();
+        log.set_path(Some(&second), b"header\n").unwrap();
+        log.write(b"three\n").unwrap();
+        log.set_path(None, b"header\n").unwrap();
+        log.write(b"ignored\n").unwrap();
+        assert!(!log.enabled());
+        assert_eq!(std::fs::read(first).unwrap(), b"header\none\ntwo\n");
+        assert_eq!(std::fs::read(second).unwrap(), b"header\nthree\n");
+    }
+
+    #[test]
+    fn failed_replacement_disables_the_old_destination() {
+        let dir = TestDir::new();
+        let path = dir.0.join("log");
+        let log = AppendLog::default();
+        log.set_path(Some(&path), b"header\n").unwrap();
+        assert!(
+            log.set_path(Some(&dir.0.join("missing/log")), b"header\n")
+                .is_err()
+        );
+        assert!(!log.enabled());
+        log.write(b"ignored\n").unwrap();
+        assert_eq!(std::fs::read(path).unwrap(), b"header\n");
+    }
+
+    #[test]
+    fn concurrent_records_are_complete() {
+        let dir = TestDir::new();
+        let path = dir.0.join("log");
+        let log = AppendLog::default();
+        log.set_path(Some(&path), b"").unwrap();
+        std::thread::scope(|scope| {
+            for n in 0..4 {
+                let log = &log;
+                scope.spawn(move || {
+                    for _ in 0..100 {
+                        log.write(format!("{n}\n").as_bytes()).unwrap();
+                    }
+                });
+            }
+        });
+        let contents = std::fs::read_to_string(path).unwrap();
+        assert_eq!(contents.lines().count(), 400);
+        for n in 0..4 {
+            assert_eq!(
+                contents
+                    .lines()
+                    .filter(|line| *line == n.to_string())
+                    .count(),
+                100
+            );
+        }
+    }
+}

crates/stdlib/src/ssl.rs

@@ -30,6 +30,8 @@ mod compat;
 // SSL exception types (shared with openssl backend)
 mod error;
 
+mod keylog;
+
 // Utilities for setting a Rustls cryptography provider.
 pub mod providers;
 
@@ -54,7 +56,8 @@ mod _ssl {
             },
             convert::IntoPyException,
             function::{
-                ArgBytesLike, ArgMemoryBuffer, Either, FuncArgs, OptionalArg, PyComparisonValue,
+                ArgBytesLike, ArgMemoryBuffer, Either, FsPath, FuncArgs, OptionalArg,
+                PyComparisonValue, PySetterValue,
             },
             stdlib::_warnings,
             types::{Comparable, Constructor, Hashable, PyComparisonOp, Representable},
@@ -732,6 +735,9 @@ mod _ssl {
         sni_callback: PyRwLock<Option<PyObjectRef>>,
         // Message callback for debugging (contains PyObjectRef - needs GC tracking)
         msg_callback: PyRwLock<Option<PyObjectRef>>,
+        keylog_filename: PyRwLock<Option<PyObjectRef>>,
+        #[pytraverse(skip)]
+        key_log: Arc<super::keylog::KeyLog>,
         // ECDH curve name for key exchange
         #[pytraverse(skip)]
         ecdh_curve: PyRwLock<Option<String>>,
@@ -1697,6 +1703,42 @@ mod _ssl {
             self.msg_callback.read().clone()
         }
 
+        #[pygetset]
+        fn keylog_filename(&self) -> Option<PyObjectRef> {
+            self.keylog_filename.read().clone()
+        }
+
+        #[pygetset(setter)]
+        fn set_keylog_filename(&self, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> {
+            let PySetterValue::Assign(value) = value else {
+                return Err(vm.new_attribute_error("attribute 'keylog_filename' cannot be deleted"));
+            };
+            // CPython disables the old destination even if path conversion or
+            // opening the replacement fails. Retain the original Python object
+            // (including PathLike and str subclasses) for the getter and GC.
+            let old_filename = self.keylog_filename.write().take();
+            self.key_log
+                .0
+                .set_path(None, super::keylog::HEADER)
+                .map_err(|e| e.into_pyexception(vm))?;
+            drop(old_filename);
+            if vm.is_none(&value) {
+                return Ok(());
+            }
+            let path = FsPath::try_from_object(vm, value.clone())?.to_path_buf(vm)?;
+            self.key_log
+                .0
+                .set_path(Some(&path), super::keylog::HEADER)
+                .map_err(|e| {
+                    let exc = e.into_pyexception(vm);
+                    let _ = exc.as_object().set_attr("filename", value.clone(), vm);
+                    exc
+                })?;
+            let old_filename = self.keylog_filename.write().replace(value);
+            drop(old_filename);
+            Ok(())
+        }
+
         #[pygetset(setter)]
         fn set__msg_callback(
             &self,
@@ -1886,6 +1928,7 @@ mod _ssl {
                     .ctx
                     .new_bytearray(Vec::with_capacity(TLS_RECORD_HEADER_SIZE))
                     .into(),
+                key_log: Arc::new(super::keylog::ConnectionKeyLog::new(zelf.key_log.clone())),
                 context: PyRwLock::new(zelf),
                 server_side: args.server_side,
                 server_hostname: PyRwLock::new(hostname),
@@ -1980,6 +2023,7 @@ mod _ssl {
                     .ctx
                     .new_bytearray(Vec::with_capacity(TLS_RECORD_HEADER_SIZE))
                     .into(),
+                key_log: Arc::new(super::keylog::ConnectionKeyLog::new(zelf.key_log.clone())),
                 context: PyRwLock::new(zelf),
                 server_side,
                 server_hostname: PyRwLock::new(hostname),
@@ -2299,6 +2343,8 @@ mod _ssl {
                 maximum_version: PyRwLock::new(max_version),
                 sni_callback: PyRwLock::new(None),
                 msg_callback: PyRwLock::new(None),
+                keylog_filename: PyRwLock::new(None),
+                key_log: Arc::new(super::keylog::KeyLog::default()),
                 ecdh_curve: PyRwLock::new(None),
                 ca_cert_count: PyRwLock::new(0),
                 x509_cert_count: PyRwLock::new(0),
@@ -2335,6 +2381,8 @@ mod _ssl {
         tls_record_header_buf: PyObjectRef,
         // SSL context
         context: PyRwLock<PyRef<PySSLContext>>,
+        #[pytraverse(skip)]
+        key_log: Arc<super::keylog::ConnectionKeyLog>,
         // Server-side or client-side
         #[pytraverse(skip)]
         server_side: bool,
@@ -3382,7 +3430,7 @@ mod _ssl {
             };
             drop(ctx);
 
-            let (config_arc, chain_builder) = if let Some(cached) = cached_config_arc {
+            let (mut config_arc, chain_builder) = if let Some(cached) = cached_config_arc {
                 cached
             } else {
                 let config =
@@ -3397,6 +3445,13 @@ mod _ssl {
             };
             *self.chain_builder.write() = Some(chain_builder);
 
+            // The first SNI connection only extracts ClientHello and is
+            // discarded after the callback. Log only the real handshake.
+            // Keep connection-specific routing out of the cached config.
+            if !use_sni_resolver || !self.is_first_sni_read() {
+                Arc::make_mut(&mut config_arc).key_log = self.key_log.clone();
+            }
+
             let conn = ServerConnection::new(config_arc).map_err(|e| {
                 vm.new_value_error(format!("Failed to create server connection: {e}"))
             })?;
@@ -3439,6 +3494,7 @@ mod _ssl {
                 // Check for pending context change (from SNI callback)
                 let pending_context = self.pending_context.write().take();
                 if let Some(new_ctx) = pending_context {
+                    self.key_log.set_sink(new_ctx.key_log.clone());
                     *self.context.write() = new_ctx;
                 }
 
@@ -3500,7 +3556,7 @@ mod _ssl {
 
                     let explicit_session = self.session.read().clone();
                     let session_store = Arc::new(PythonClientSessionStore::new(session_cache));
-                    let (config, chain_builder) = if let Some(session) = explicit_session {
+                    let (mut config, chain_builder) = if let Some(session) = explicit_session {
                         let session = session
                             .downcast_ref::<PySSLSession>()
                             .ok_or_else(|| vm.new_type_error("Value is not a SSLSession."))?;
@@ -3539,6 +3595,7 @@ mod _ssl {
                         };
                         create_client_config(config_options).map_err(|e| vm.new_value_error(e))?
                     };
+                    Arc::make_mut(&mut config).key_log = self.key_log.clone();
 
                     *self.chain_builder.write() = Some(chain_builder);
                     *self.client_config.write() = Some(config.clone());
@@ -4002,6 +4059,7 @@ mod _ssl {
 
         #[pygetset(setter)]
         fn set_context(&self, value: PyRef<PySSLContext>, _vm: &VirtualMachine) {
+            self.key_log.set_sink(value.key_log.clone());
             // Update context reference immediately
             // SSL_set_SSL_CTX allows context changes at any time,
             // even after handshake completion

crates/stdlib/src/ssl/keylog.rs

@@ -0,0 +1,80 @@
+use alloc::sync::Arc;
+use rustpython_host_env::fs::AppendLog;
+
+pub(super) 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(super) struct KeyLog(pub AppendLog);
+
+/// Each connection follows its current Python context. This routing state must
+/// not be shared through the context's config cache or an explicit SSLSession.
+#[derive(Debug)]
+pub(super) struct ConnectionKeyLog(parking_lot::RwLock<Arc<KeyLog>>);
+
+impl ConnectionKeyLog {
+    pub(super) fn new(sink: Arc<KeyLog>) -> Self {
+        Self(parking_lot::RwLock::new(sink))
+    }
+
+    pub(super) fn set_sink(&self, sink: Arc<KeyLog>) {
+        *self.0.write() = sink;
+    }
+}
+
+impl rustls::KeyLog for ConnectionKeyLog {
+    fn log(&self, label: &str, client_random: &[u8], secret: &[u8]) {
+        self.0.read().log(label, client_random, secret);
+    }
+
+    fn will_log(&self, label: &str) -> bool {
+        self.0.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"));
+    }
+}