From b0b2bcfc25eda92365472a5f794911398e2388a6 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sat, 5 Sep 2026 20:40:33 +0900 Subject: [PATCH 1/4] ssl: accept ClientHello before selecting the SNI context Replace the provisional server connection and replay flags with explicit handshake states. Retain fatal alerts across nonblocking retries and preserve fragmented input on timeout. Enable the three SNI alert regressions. Assisted-by: Codex:GPT-5 --- Lib/test/test_ssl.py | 3 - crates/stdlib/src/ssl.rs | 643 +++++++++++++---------------- crates/stdlib/src/ssl/compat.rs | 46 +-- crates/stdlib/src/ssl/handshake.rs | 119 ++++++ 4 files changed, 407 insertions(+), 404 deletions(-) create mode 100644 crates/stdlib/src/ssl/handshake.rs diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py index cd624696c02..f7ecabd853c 100644 --- a/Lib/test/test_ssl.py +++ b/Lib/test/test_ssl.py @@ -4455,7 +4455,6 @@ def servername_cb(ssl_sock, server_name, initial_context): self.check_common_name(stats, SIGNED_CERTFILE_HOSTNAME) self.assertEqual(calls, []) - @unittest.expectedFailure # TODO: RUSTPYTHON; + TLSV1_ALERT_ACCESS_DENIED def test_sni_callback_alert(self): # Returning a TLS alert is reflected to the connecting client server_context, other_context, client_context = self.sni_contexts() @@ -4469,7 +4468,6 @@ def cb_returning_alert(ssl_sock, server_name, initial_context): sni_name='supermessage') self.assertEqual(cm.exception.reason, 'TLSV1_ALERT_ACCESS_DENIED') - @unittest.expectedFailure # TODO: RUSTPYTHON def test_sni_callback_raising(self): # Raising fails the connection with a TLS handshake failure alert. server_context, other_context, client_context = self.sni_contexts() @@ -4489,7 +4487,6 @@ def cb_raising(ssl_sock, server_name, initial_context): self.assertRegex(cm.exception.reason, regex) self.assertEqual(catch.unraisable.exc_type, ZeroDivisionError) - @unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'SSLEOFError' object has no attribute 'reason' def test_sni_callback_wrong_return_type(self): # Returning the wrong return type terminates the TLS connection # with an internal error alert. diff --git a/crates/stdlib/src/ssl.rs b/crates/stdlib/src/ssl.rs index f22d8b25367..7fac3e7b133 100644 --- a/crates/stdlib/src/ssl.rs +++ b/crates/stdlib/src/ssl.rs @@ -30,6 +30,7 @@ mod compat; // SSL exception types (shared with openssl backend) mod error; +mod handshake; mod keylog; // Utilities for setting a Rustls cryptography provider. @@ -88,11 +89,11 @@ mod _ssl { use parking_lot::{Mutex as ParkingMutex, RwLock as ParkingRwLock}; use pem_rfc7468::{LineEnding, encode_string}; use rustls::{ - ClientConnection, Connection, HandshakeKind, RootCertStore, ServerConnection, + ClientConnection, Connection, HandshakeKind, RootCertStore, client::{ClientSessionMemoryCache, ClientSessionStore}, crypto::SupportedKxGroup, pki_types::{CertificateDer, CertificateRevocationListDer, PrivateKeyDer, ServerName}, - server::{ClientHello, ResolvesServerCert}, + server::{Accepted, ResolvesServerCert}, sign::CertifiedKey, version::{TLS12, TLS13}, }; @@ -113,10 +114,11 @@ mod _ssl { ssl_do_handshake, }; + use super::handshake::HandshakeState; use super::providers::CryptoExt; // Type aliases for better readability - // Additional type alias for certificate/key pairs (SessionCache and SniCertName defined below) + // Additional type alias for certificate/key pairs (SessionCache defined below) /// Certificate and private key pair used in SSL contexts type CertKeyPair = (Arc, PrivateKeyDer<'static>); @@ -425,30 +427,6 @@ mod _ssl { Ok(()) } - // SNI certificate resolver that uses shared mutable state - // The Python SNI callback updates this state, and resolve() reads from it - #[derive(Debug)] - struct SniCertResolver { - // SNI state: (certificate, server_name) - sni_state: Arc>, - } - - impl ResolvesServerCert for SniCertResolver { - fn resolve(&self, client_hello: ClientHello<'_>) -> Option> { - let mut state = self.sni_state.lock(); - - // Extract and store SNI from client hello for later use - if let Some(sni) = client_hello.server_name() { - state.1 = Some(sni.to_string()); - } else { - state.1 = None; - } - - // Return the current certificate (may have been updated by Python callback) - Some(state.0.clone()) - } - } - // Session data structure for tracking TLS sessions #[derive(Debug, Clone)] struct SessionData { @@ -486,9 +464,6 @@ mod _ssl { // Type alias to simplify complex session cache type type SessionCache = Arc, Arc>>>>; - // Type alias for SNI state - type SniCertName = (Arc, Option); - // SESSION EMULATION IMPLEMENTATION // // IMPORTANT: This is an EMULATION of CPython's SSL session management. @@ -1933,7 +1908,7 @@ mod _ssl { server_side: args.server_side, server_hostname: PyRwLock::new(hostname), connection: PyMutex::new(None), - handshake_done: PyMutex::new(false), + handshake_state: PyMutex::new(HandshakeState::new(args.server_side)), session_was_reused: PyMutex::new(false), owner: PyRwLock::new( args.owner @@ -1948,10 +1923,6 @@ mod _ssl { client_session_store: PyRwLock::new(None), incoming_bio: None, outgoing_bio: None, - sni_state: PyRwLock::new(None), - pending_context: PyRwLock::new(None), - client_hello_buffer: PyMutex::new(None), - sni_callback_processed: PyMutex::new(false), shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), @@ -2028,7 +1999,7 @@ mod _ssl { server_side, server_hostname: PyRwLock::new(hostname), connection: PyMutex::new(None), - handshake_done: PyMutex::new(false), + handshake_state: PyMutex::new(HandshakeState::new(server_side)), session_was_reused: PyMutex::new(false), owner: PyRwLock::new( args.owner @@ -2043,10 +2014,6 @@ mod _ssl { client_session_store: PyRwLock::new(None), incoming_bio: Some(args.incoming), outgoing_bio: Some(args.outgoing), - sni_state: PyRwLock::new(None), - pending_context: PyRwLock::new(None), - client_hello_buffer: PyMutex::new(None), - sni_callback_processed: PyMutex::new(false), shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), @@ -2392,9 +2359,8 @@ mod _ssl { // TLS connection state #[pytraverse(skip)] connection: PyMutex>, - // Handshake completed flag - #[pytraverse(skip)] - handshake_done: PyMutex, + // Includes the saved exception while a rejected handshake sends its alert. + handshake_state: PyMutex, // Session was reused (for session resumption tracking) #[pytraverse(skip)] session_was_reused: PyMutex, @@ -2416,17 +2382,6 @@ mod _ssl { // MemoryBIO mode (optional) incoming_bio: Option>, outgoing_bio: Option>, - // SNI certificate resolver state (for server-side only) - #[pytraverse(skip)] - sni_state: PyRwLock>>>, - // Pending context change (for SNI callback deferred handling) - pending_context: PyRwLock>>, - // Buffer to store ClientHello for connection recreation - #[pytraverse(skip)] - client_hello_buffer: PyMutex>>, - // Whether the Python SNI callback has already been run for this handshake - #[pytraverse(skip)] - sni_callback_processed: PyMutex, // Shutdown state for tracking close-notify exchange #[pytraverse(skip)] shutdown_state: PyMutex, @@ -2646,7 +2601,7 @@ mod _ssl { } fn complete_handshake(&self, vm: &VirtualMachine) { - *self.handshake_done.lock() = true; + *self.handshake_state.lock() = HandshakeState::Connected; // Check if session was resumed - get value and release lock immediately let was_resumed = self @@ -2731,146 +2686,199 @@ mod _ssl { sock_wait(&socket, wait_kind, timeout, vm).map_err(|e| e.into_pyexception(vm)) } - // SNI (Server Name Indication) Helper Methods: - // These methods support the server-side handshake SNI callback mechanism - - /// Check if this is the first read during handshake (for SNI callback) - /// Returns true until the SNI callback has been processed. - pub(crate) fn is_first_sni_read(&self) -> bool { - !*self.sni_callback_processed.lock() + fn handshake_completed(&self) -> bool { + matches!(*self.handshake_state.lock(), HandshakeState::Connected) } - /// Check if SNI callback is configured - pub(crate) fn has_sni_callback(&self) -> bool { - // Nested read locks are safe - self.context.read().sni_callback.read().is_some() + fn reject_handshake(&self, error: PyBaseExceptionRef, bytes: Vec) { + *self.handshake_state.lock() = HandshakeState::SendingAlert { + error, + bytes, + sent: 0, + }; } - /// Save ClientHello data for potential connection recreation. - pub(crate) fn save_client_hello_from_bytes(&self, bytes_data: &[u8]) { - let mut buffer = self.client_hello_buffer.lock(); - match buffer.as_mut() { - Some(existing) => existing.extend_from_slice(bytes_data), - None => *buffer = Some(bytes_data.to_vec()), + /// Receive ClientHello and select the context before rustls can generate + /// ServerHello. Keep protocol progress across non-blocking retries. + fn accept_client_hello(&self, vm: &VirtualMachine) -> PyResult> { + loop { + // Never hold the state lock across Python callbacks or I/O. + // Reentrant handshake calls see an in-progress operation. + let state = core::mem::replace( + &mut *self.handshake_state.lock(), + HandshakeState::CallingSni, + ); + match state { + HandshakeState::WaitingForClientHello(mut acceptor) => { + match acceptor.accept() { + Ok(Some(accepted)) => { + let name = accepted.client_hello().server_name().map(str::to_owned); + if let Err((description, error)) = + self.invoke_sni_callback(name.as_deref(), vm) + { + self.reject_handshake( + error, + super::handshake::sni_alert(description), + ); + continue; + } + return Ok(Some(accepted)); + } + Err((error, mut alert)) => { + let mut bytes = Vec::new(); + // Vec writes cannot fail. + let _ = alert.write_all(&mut bytes); + let error = if matches!(error, rustls::Error::InvalidMessage(_)) { + SslError::PreauthData + } else { + SslError::from_rustls(error) + }; + self.reject_handshake(error.into_py_err(vm), bytes); + continue; + } + Ok(None) => {} + } + let read = (|| { + if self.sock_wait_for_io_impl(SockWaitKind::Read, vm)? { + return Err(timeout_error_msg( + vm, + "The handshake operation timed out".to_owned(), + ) + .upcast()); + } + let bytes = self.sock_recv_at_most_one_tls_record(vm).map_err(|e| { + if is_blocking_io_error(&e, vm) { + create_ssl_want_read_error(vm).upcast() + } else { + e + } + })?; + if bytes.is_empty() { + return Err( + if self.is_bio_mode() + && !self.incoming_bio.as_ref().is_some_and(|bio| bio.eof()) + { + create_ssl_want_read_error(vm).upcast() + } else { + SslError::Eof.into_py_err(vm) + }, + ); + } + super::handshake::feed_acceptor(&mut acceptor, bytes.as_bytes()) + .map_err(|e| e.into_pyexception(vm))?; + Ok(()) + })(); + *self.handshake_state.lock() = + HandshakeState::WaitingForClientHello(acceptor); + read?; + } + HandshakeState::SendingAlert { + error, + bytes, + mut sent, + } => { + let write = (|| { + if sent == bytes.len() { + return Ok(()); + } + if self.sock_wait_for_io_impl(SockWaitKind::Write, vm)? { + return Err(timeout_error_msg( + vm, + "The handshake operation timed out".to_owned(), + ) + .upcast()); + } + let result = self.sock_send(&bytes[sent..], vm).map_err(|e| { + if is_blocking_io_error(&e, vm) { + create_ssl_want_write_error(vm).upcast() + } else { + e + } + })?; + let count = result.try_to_value::(vm)?; + if count == 0 { + return Err(create_ssl_want_write_error(vm).upcast()); + } + sent += count; + Ok(()) + })(); + let done = sent == bytes.len(); + *self.handshake_state.lock() = HandshakeState::SendingAlert { + error: error.clone(), + bytes, + sent, + }; + write?; + if done { + return Err(error); + } + } + state @ (HandshakeState::Handshaking | HandshakeState::Connected) => { + *self.handshake_state.lock() = state; + return Ok(None); + } + HandshakeState::CallingSni => { + return Err(vm.new_value_error("handshake already in progress")); + } + } } } - /// Get the extracted SNI name from resolver - pub(crate) fn get_extracted_sni_name(&self) -> Option { - // Clone the Arc option to avoid nested lock (sni_state.read -> arc.lock) - let sni_state_opt = self.sni_state.read().clone(); - sni_state_opt.as_ref().and_then(|arc| arc.lock().1.clone()) - } - - /// Invoke the Python SNI callback - pub(crate) fn invoke_sni_callback( + /// Invoke SNI without any connection or state locks held. The alert is + /// returned separately from the local exception so the peer gets the + /// correct wire error before this operation reports failure. + fn invoke_sni_callback( &self, sni_name: Option<&str>, vm: &VirtualMachine, - ) -> PyResult<()> { - // The callback may have been cleared (sni_callback = None) between the - // handshake deciding to invoke it and this point. A concurrent removal - // is not an error: there is simply nothing to run. + ) -> Result<(), (u8, PyBaseExceptionRef)> { let callback = self.context.read().sni_callback.read().clone(); let Some(callback) = callback else { return Ok(()); }; - + let failed = || { + SslError::create_ssl_error_with_reason( + vm, + Some("SSL"), + "CALLBACK_FAILED", + "[SSL: CALLBACK_FAILED] callback failed", + ) + }; let ssl_sock = self .owner .read() .as_ref() .and_then(|owner| owner.upgrade()) .ok_or_else(|| { - super::compat::SslError::create_ssl_error_with_reason( - vm, - Some("SSL"), - "PARSE_TLSEXT", - "[SSL: PARSE_TLSEXT] SNI callback owner is no longer available", + ( + 80, + SslError::create_ssl_error_with_reason( + vm, + Some("SSL"), + "PARSE_TLSEXT", + "[SSL: PARSE_TLSEXT] SNI callback owner is no longer available", + ), ) })?; - let server_name_py: PyObjectRef = match sni_name { - Some(name) => vm.ctx.new_str(name.to_string()).into(), - None => vm.ctx.none(), - }; - let initial_context: PyObjectRef = self.context.read().clone().into(); - - // catches exceptions from the callback and reports them as unraisable - let result = match callback.call((ssl_sock, server_name_py, initial_context), vm) { - Ok(result) => result, + let server_name: PyObjectRef = + sni_name.map_or_else(|| vm.ctx.none(), |name| vm.ctx.new_str(name).into()); + let context: PyObjectRef = self.context.read().clone().into(); + let result = callback + .call((ssl_sock, server_name, context), vm) + .map_err(|exc| { + vm.run_unraisable(exc, Some("in ssl servername callback".to_owned()), callback); + (40, failed()) + })?; + if vm.is_none(&result) { + return Ok(()); + } + match result.try_to_value::(vm) { + Ok(alert) => Err((alert as u8, failed())), Err(exc) => { - vm.run_unraisable( - exc, - Some("in ssl servername callback".to_owned()), - callback.clone(), - ); - // Return SSL error like SSL_TLSEXT_ERR_ALERT_FATAL - let ssl_exc: PyBaseExceptionRef = vm - .new_os_subtype_error( - PySSLError::class(&vm.ctx).to_owned(), - None, - "SNI callback raised exception", - ) - .upcast(); - let _ = ssl_exc.as_object().set_attr( - "reason", - vm.ctx.new_str("TLSV1_ALERT_INTERNAL_ERROR"), - vm, - ); - return Err(ssl_exc); - } - }; - - // Check return value type (must be None or integer) - if !vm.is_none(&result) { - // Try to convert to integer - if result.try_to_value::(vm).is_err() { - // Type conversion failed - raise TypeError as unraisable - let type_error = vm.new_type_error(format!( - "servername callback must return None or an integer, not '{}'", - result.class().name() - )); - vm.run_unraisable(type_error, None, result); - - // Return SSL error with reason set to TLSV1_ALERT_INTERNAL_ERROR - // - // RUSTLS API LIMITATION: - // We cannot send a TLS InternalError alert to the client here because: - // 1. Rustls does not provide a public API like send_fatal_alert() - // 2. This method is called AFTER dropping the connection lock (to prevent deadlock) - // 3. By the time we detect the error, the connection is no longer available - // - // CPython/OpenSSL behavior: - // - SNI callback runs inside SSL_do_handshake with connection active - // - Sets *al = SSL_AD_INTERNAL_ERROR - // - OpenSSL automatically sends alert before returning - // - // RustPython/Rustls behavior: - // - SNI callback runs after dropping connection lock (deadlock prevention) - // - Exception has _reason='TLSV1_ALERT_INTERNAL_ERROR' for error reporting - // - TCP connection closes without sending TLS alert to client - // - // If rustls adds send_fatal_alert() API in the future, we should: - // - Re-acquire connection lock after callback - // - Call: connection.send_fatal_alert(AlertDescription::InternalError) - // - Then close connection - let exc: PyBaseExceptionRef = vm - .new_os_subtype_error( - PySSLError::class(&vm.ctx).to_owned(), - None, - "SNI callback returned invalid type", - ) - .upcast(); - let _ = exc.as_object().set_attr( - "reason", - vm.ctx.new_str("TLSV1_ALERT_INTERNAL_ERROR"), - vm, - ); - return Err(exc); + vm.run_unraisable(exc, None, result); + Err((80, failed())) } } - - Ok(()) } // Helper to call socket methods, bypassing any SSL wrapper @@ -2905,62 +2913,48 @@ mod _ssl { .expect("BUG: tls_record_header_buf is not PyByteArray"); let buf_len = tls_record_header_buf.borrow_buf().len(); - - let (mut with_header, mut remaining_record_body_len) = - if buf_len < TLS_RECORD_HEADER_SIZE { - // We do not have a full TLS record header, start receiving one. - let bytes_obj = self.sock_recv(TLS_RECORD_HEADER_SIZE - buf_len, vm)?; - let bytes = obj_to_bytes(bytes_obj)?; - - let mut buf = tls_record_header_buf.borrow_buf_mut(); - buf.extend_from_slice(bytes.as_bytes()); - - if buf.len() < TLS_RECORD_HEADER_SIZE { - return Ok(bytes); - } - - // Parse the remaining length. - let record_body_len = u16::from_be_bytes([buf[3], buf[4]]); - // Validity of length value will be checked by rustls. - - // Zero-length TLS record. - if record_body_len == 0 { - buf.clear(); - return Ok(bytes); + let (mut collected, remaining) = if buf_len < TLS_RECORD_HEADER_SIZE { + let bytes = obj_to_bytes(self.sock_recv(TLS_RECORD_HEADER_SIZE - buf_len, vm)?)?; + let mut header = tls_record_header_buf.borrow_buf_mut(); + header.extend_from_slice(bytes.as_bytes()); + if header.len() < TLS_RECORD_HEADER_SIZE { + return Ok(bytes); + } + let remaining = u16::from_be_bytes([header[3], header[4]]); + if remaining == 0 { + header.clear(); + return Ok(bytes); + } + (Some(bytes.as_bytes().to_vec()), remaining) + } else { + let header = tls_record_header_buf.borrow_buf(); + (None, u16::from_be_bytes([header[3], header[4]])) + }; + let bytes = match self.sock_recv(remaining as usize, vm) { + Ok(bytes) => obj_to_bytes(bytes)?, + Err(error) => { + // Header bytes already consumed must reach the TLS parser, + // even if the body is not yet available on a nonblocking socket. + if (is_blocking_io_error(&error, vm) + || error.fast_isinstance(vm.ctx.exceptions.timeout_error)) + && let Some(header) = collected + { + return Ok(vm.ctx.new_bytes(header)); } - - let mut bytes_vec = bytes.as_bytes().to_vec(); - bytes_vec.reserve(record_body_len as usize); - (Some(bytes_vec), record_body_len) - } else { - let buf = tls_record_header_buf.borrow_buf(); - let remaining_record_body_len = u16::from_be_bytes([buf[3], buf[4]]); - (None, remaining_record_body_len) - }; - - // We have full record header and are in a process of receiving a record. - let bytes_obj = self.sock_recv(remaining_record_body_len as usize, vm)?; - let bytes = obj_to_bytes(bytes_obj)?; - - if let Some(with_header) = with_header.as_mut() { - with_header.extend_from_slice(bytes.as_bytes()); - } - - let mut buf = tls_record_header_buf.borrow_buf_mut(); - remaining_record_body_len -= bytes.len() as u16; - if remaining_record_body_len == 0 { - // Record received completely, need to start a new one beginning with its header. - buf.clear(); + return Err(error); + } + }; + let mut header = tls_record_header_buf.borrow_buf_mut(); + let remaining = remaining - bytes.len() as u16; + if remaining == 0 { + header.clear(); } else { - // Update remaining length in the header. - buf.as_mut_slice()[3..5].copy_from_slice(&remaining_record_body_len.to_be_bytes()); + header[3..5].copy_from_slice(&remaining.to_be_bytes()); } - - if let Some(with_header) = with_header { - Ok(vm.ctx.new_bytes(with_header)) - } else { - Ok(bytes) + if let Some(collected) = collected.as_mut() { + collected.extend_from_slice(bytes.as_bytes()); } + Ok(collected.map_or(bytes, |bytes| vm.ctx.new_bytes(bytes))) } /// Socket send - just sends data, caller must handle pending flush @@ -3278,10 +3272,11 @@ mod _ssl { /// - ALPN protocol negotiation /// - Session resumption configuration /// - /// Returns the configured ServerConnection. + /// Continue the accepted ClientHello using the context selected by SNI. fn initialize_server_connection( &self, conn_guard: &mut Option, + accepted: Accepted, vm: &VirtualMachine, ) -> PyResult<()> { let ctx = self.context.read(); @@ -3332,55 +3327,11 @@ mod _ssl { false }; - // Check if SNI callback is set - let sni_callback = ctx.sni_callback.read().clone(); - let use_sni_resolver = sni_callback.is_some(); - - // Create SNI state if needed (to be stored in PySSLSocket later) - // For SNI, use the first cert_key pair as the initial certificate - let sni_state: Option>> = if use_sni_resolver { - // Use first cert_key as initial certificate for SNI - // Extract CertifiedKey from tuple - let (first_cert_key, _) = &cert_keys_clone[0]; - let first_cert_key = first_cert_key.clone(); - - // Check if we already have existing SNI state (from previous connection) - let existing_sni_state = self.sni_state.read().clone(); - - if let Some(sni_state_arc) = existing_sni_state { - // Reuse existing Arc and update its contents - // This is crucial: rustls SniCertResolver holds references to this Arc - let mut state = sni_state_arc.lock(); - state.0 = first_cert_key; - state.1 = None; // Reset SNI name for new connection - drop(state); - - // Return the existing Arc (not a new one!) - Some(sni_state_arc) - } else { - // First connection: create new SNI state - Some(Arc::new(ParkingMutex::new((first_cert_key, None)))) - } - } else { - None - }; - - // Determine which cert resolver to use - // Priority: SNI > Multi-cert/Single-cert via MultiCertResolver - let cert_resolver: Option> = if use_sni_resolver { - // SNI takes precedence - use first cert_key for initial setup - sni_state.as_ref().map(|sni_state_arc| { - Arc::new(SniCertResolver { - sni_state: sni_state_arc.clone(), - }) as Arc - }) - } else { - // Use MultiCertResolver for all cases (single or multiple certs) - // Extract CertifiedKey from tuples for MultiCertResolver - let cert_keys_only: Vec> = - cert_keys_clone.iter().map(|(ck, _)| ck.clone()).collect(); - Some(Arc::new(MultiCertResolver::new(cert_keys_only))) - }; + // Certificate selection uses the actual ClientHello after SNI has + // selected the context, including multi-certificate configurations. + let cert_keys_only = cert_keys_clone.iter().map(|(ck, _)| ck.clone()).collect(); + let cert_resolver: Option> = + Some(Arc::new(MultiCertResolver::new(cert_keys_only))); // Extract cert_chain and private_key from first cert_key // @@ -3421,8 +3372,10 @@ mod _ssl { // A capath may gain hashed entries between connections, so only // cache configurations whose trust store is context-owned. + // A deferred verifier owns this connection's error slot, so it + // must never be reused by another connection. let cache_server_config = - !use_sni_resolver && ctx.capath_state.read().directories.is_empty(); + !use_deferred_validation && ctx.capath_state.read().directories.is_empty(); let cached_config_arc = if cache_server_config { ctx.server_config.read().clone() } else { @@ -3445,36 +3398,22 @@ 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}")) - })?; - - *conn_guard = Some(Connection::Server(conn)); - - // If ClientHello buffer exists (from SNI callback), re-inject it - if let Some(ref hello_data) = *self.client_hello_buffer.lock() - && let Some(Connection::Server(ref mut server)) = *conn_guard - { - let mut cursor = std::io::Cursor::new(hello_data.as_slice()); - let _ = server.read_tls(&mut cursor); - - // Process the re-injected ClientHello - let _ = server.process_new_packets(); - - // DON'T clear buffer - keep it to prevent callback from being invoked again - // The buffer being non-empty signals that SNI callback was already processed - } - - // Store SNI state if we're using SNI resolver - if let Some(sni_state_arc) = sni_state { - *self.sni_state.write() = Some(sni_state_arc); + // The cache contains only context-owned configuration, never a + // particular connection's mutable keylog routing. + Arc::make_mut(&mut config_arc).key_log = self.key_log.clone(); + match accepted.into_connection(config_arc) { + Ok(conn) => { + *conn_guard = Some(Connection::Server(conn)); + *self.handshake_state.lock() = HandshakeState::Handshaking; + } + Err((error, mut alert)) => { + let mut bytes = Vec::new(); + alert + .write_all(&mut bytes) + .map_err(|e| e.into_pyexception(vm))?; + self.reject_handshake(SslError::from_rustls(error).into_py_err(vm), bytes); + self.accept_client_hello(vm)?; + } } Ok(()) @@ -3483,24 +3422,36 @@ mod _ssl { #[pymethod] fn do_handshake(&self, vm: &VirtualMachine) -> PyResult<()> { // Check if handshake already done - if *self.handshake_done.lock() { + if self.handshake_completed() { return Ok(()); } + let accepted = if self.server_side { + self.accept_client_hello(vm)? + } else { + None + }; let mut conn_guard = self.connection.lock(); // Initialize connection if not already done if conn_guard.is_none() { - // 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; - } - if self.server_side { - // Server-side connection - delegate to helper method - self.initialize_server_connection(&mut conn_guard, vm)?; + if let Err(error) = self.initialize_server_connection( + &mut conn_guard, + accepted + .ok_or_else(|| vm.new_value_error("TLS connection is not available"))?, + vm, + ) { + drop(conn_guard); + if matches!( + *self.handshake_state.lock(), + HandshakeState::SendingAlert { .. } + ) { + return Err(error); + } + self.reject_handshake(error, super::handshake::sni_alert(40)); + return self.accept_client_hello(vm).map(|_| ()); + } } else { // Client-side connection let ctx = self.context.read(); @@ -3612,46 +3563,11 @@ mod _ssl { // Perform the actual handshake by exchanging data with the socket/BIO let conn = conn_guard.as_mut().expect("unreachable"); - let is_client = matches!(conn, Connection::Client(_)); let handshake_result = ssl_do_handshake(conn, self, vm); drop(conn_guard); - - if is_client { - // CLIENT is simple - no SNI callback handling needed - handshake_result.map_err(|e| e.into_py_err(vm))?; - self.complete_handshake(vm); - Ok(()) - } else { - // Use OpenSSL-compatible handshake for server - // Handle SNI callback restart - match handshake_result { - Ok(()) => { - // Handshake completed successfully - self.complete_handshake(vm); - Ok(()) - } - Err(SslError::SniCallbackRestart) => { - // SNI detected - need to call callback and recreate connection - - // Get the SNI name that was extracted (may be None if client didn't send SNI) - let sni_name = self.get_extracted_sni_name(); - - // Now safe to call Python callback (no locks held) - self.invoke_sni_callback(sni_name.as_deref(), vm)?; - *self.sni_callback_processed.lock() = true; - - // Clear connection to trigger recreation - *self.connection.lock() = None; - - // Recursively call do_handshake to recreate with new context - self.do_handshake(vm) - } - Err(e) => { - // Other errors - convert to Python exception - Err(e.into_py_err(vm)) - } - } - } + handshake_result.map_err(|e| e.into_py_err(vm))?; + self.complete_handshake(vm); + Ok(()) } #[pymethod] @@ -3696,7 +3612,7 @@ mod _ssl { // Ensure handshake is done - if not, complete it first // This matches OpenSSL behavior where SSL_read() auto-completes handshake - if !*self.handshake_done.lock() { + if !self.handshake_completed() { self.do_handshake(vm)?; } @@ -3880,7 +3796,7 @@ mod _ssl { } // Ensure handshake is done (SSL_write auto-completes handshake) - if !*self.handshake_done.lock() { + if !self.handshake_completed() { self.do_handshake(vm)?; } @@ -3933,7 +3849,7 @@ mod _ssl { let binary = args.binary_form.unwrap_or(false); // Check if handshake is complete - if !*self.handshake_done.lock() { + if !self.handshake_completed() { return Err(vm.new_value_error("handshake not done yet")); } @@ -4064,9 +3980,6 @@ mod _ssl { // SSL_set_SSL_CTX allows context changes at any time, // even after handshake completion *self.context.write() = value; - - // Clear pending context as we've applied the change - *self.pending_context.write() = None; } #[pygetset] @@ -4081,7 +3994,7 @@ mod _ssl { vm: &VirtualMachine, ) -> PyResult<()> { // Check if handshake is already done - if *self.handshake_done.lock() { + if self.handshake_completed() { return Err( vm.new_value_error("Cannot set server_hostname on socket after handshake") ); @@ -4125,7 +4038,7 @@ mod _ssl { } // Check if handshake is already done - if *self.handshake_done.lock() { + if self.handshake_completed() { return Err(vm.new_value_error("Cannot set session after handshake.")); } @@ -4192,6 +4105,14 @@ mod _ssl { #[pymethod] fn shutdown(&self, vm: &VirtualMachine) -> PyResult { + if !self.handshake_completed() { + return Err(SslError::create_ssl_error_with_reason( + vm, + Some("SSL"), + "SHUTDOWN_WHILE_IN_INIT", + "[SSL: SHUTDOWN_WHILE_IN_INIT] shutdown while in init", + )); + } // Check current shutdown state let current_state = *self.shutdown_state.lock(); @@ -4552,7 +4473,7 @@ mod _ssl { } // Check if handshake completed - if !*self.handshake_done.lock() { + if !self.handshake_completed() { return None; } @@ -4586,7 +4507,7 @@ mod _ssl { } // Check if handshake has been completed - if !*self.handshake_done.lock() { + if !self.handshake_completed() { return Err(vm.new_value_error( "Handshake must be completed before post-handshake authentication", )); diff --git a/crates/stdlib/src/ssl/compat.rs b/crates/stdlib/src/ssl/compat.rs index 8371ff1e4e9..d4ec6d65932 100644 --- a/crates/stdlib/src/ssl/compat.rs +++ b/crates/stdlib/src/ssl/compat.rs @@ -198,8 +198,6 @@ pub(super) enum SslError { Io(std::io::Error), /// Timeout error (socket.timeout) Timeout(String), - /// SNI callback triggered - need to restart handshake - SniCallbackRestart, /// Python exception (pass through directly) Py(PyBaseExceptionRef), /// TLS alert received with OpenSSL-compatible error code @@ -387,10 +385,7 @@ impl SslError { ) .upcast(), Self::Io(err) => err.into_pyexception(vm), - Self::SniCallbackRestart => { - // This should be handled at PySSLSocket level - unreachable!("SniCallbackRestart should not reach Python layer") - } + Self::Py(exc) => exc, Self::AlertReceived { lib, reason } => { Self::create_ssl_error_from_codes(vm, lib, reason) @@ -1066,18 +1061,12 @@ fn handshake_read_data( conn: &mut Connection, socket: &PySSLSocket, is_bio: bool, - is_server: bool, vm: &VirtualMachine, -) -> SslResult<(bool, bool)> { +) -> SslResult { if !conn.wants_read() { - return Ok((false, false)); + return Ok(false); } - // SERVER-SPECIFIC: Check if this is before the SNI callback. - // sock_recv() may return only part of a TLS record, so keep capturing - // ClientHello fragments until process_new_packets() has produced a response. - let is_first_sni_read = is_server && socket.has_sni_callback() && socket.is_first_sni_read(); - // Wait for data in socket mode if !is_bio { let timed_out = socket @@ -1106,26 +1095,17 @@ fn handshake_read_data( return Err(SslError::WantRead); } if !conn.wants_write() && e.fast_isinstance(vm.ctx.exceptions.timeout_error) { - return Ok((false, false)); + return Ok(false); } return Err(SslError::Py(e)); } } }; - // SERVER-SPECIFIC: Save ClientHello fragments for potential connection recreation. - if is_first_sni_read { - // Extract bytes from PyObjectRef - use rustpython_vm::builtins::PyBytes; - if let Some(bytes_obj) = data_obj.downcast_ref::() { - socket.save_client_hello_from_bytes(bytes_obj.as_bytes()); - } - } - // Feed data to rustls ssl_read_tls_records(conn, data_obj, is_bio, vm)?; - Ok((true, is_first_sni_read)) + Ok(true) } /// Handle handshake completion for server-side TLS 1.3 @@ -1261,8 +1241,7 @@ pub(super) fn ssl_do_handshake( made_progress |= write_progress; // Read TLS handshake data from socket/BIO - let (read_progress, is_first_sni_read) = - handshake_read_data(conn, socket, is_bio, is_server, vm)?; + let read_progress = handshake_read_data(conn, socket, is_bio, vm)?; made_progress |= read_progress; // Process TLS packets (state machine) @@ -1292,19 +1271,6 @@ pub(super) fn ssl_do_handshake( return Err(SslError::from_rustls(e)); } - // SERVER-SPECIFIC: Check SNI callback after processing packets. - // A partial TLS record can be read without producing any handshake - // response. Wait until rustls has processed a complete ClientHello. - if is_server && is_first_sni_read && socket.has_sni_callback() && conn.wants_write() { - // IMPORTANT: Do NOT call the callback here! - // The connection lock is still held, which would cause deadlock. - // Return SniCallbackRestart to signal do_handshake to: - // 1. Drop conn_guard - // 2. Call the callback (with Some(name) or None) - // 3. Restart handshake - return Err(SslError::SniCallbackRestart); - } - // Check if handshake is complete and handle post-handshake processing // CRITICAL: We do NOT check wants_read() - this matches CPython/OpenSSL behavior! if handle_handshake_complete(conn, socket, is_server, vm)? { diff --git a/crates/stdlib/src/ssl/handshake.rs b/crates/stdlib/src/ssl/handshake.rs new file mode 100644 index 00000000000..414fa5e1e55 --- /dev/null +++ b/crates/stdlib/src/ssl/handshake.rs @@ -0,0 +1,119 @@ +use core::fmt; +use rustls::server::Acceptor; +use rustpython_vm::{ + builtins::PyBaseExceptionRef, + object::{Traverse, TraverseFn}, +}; + +/// Server configuration is selected only after receiving ClientHello and +/// invoking SNI. A rejected handshake stays terminal while its alert drains. +pub(super) enum HandshakeState { + WaitingForClientHello(Acceptor), + CallingSni, + Handshaking, + Connected, + SendingAlert { + error: PyBaseExceptionRef, + bytes: Vec, + sent: usize, + }, +} + +impl HandshakeState { + pub(super) fn new(server_side: bool) -> Self { + if server_side { + Self::WaitingForClientHello(Acceptor::default()) + } else { + Self::Handshaking + } + } +} + +impl fmt::Debug for HandshakeState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::WaitingForClientHello(_) => "WaitingForClientHello", + Self::CallingSni => "CallingSni", + Self::Handshaking => "Handshaking", + Self::Connected => "Connected", + Self::SendingAlert { .. } => "SendingAlert", + }) + } +} + +// Only the saved Python exception can contain GC references. Acceptor and the +// encoded TLS alert contain Rust-owned protocol data. +unsafe impl Traverse for HandshakeState { + fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) { + if let Self::SendingAlert { error, .. } = self { + error.traverse(tracer_fn); + } + } +} + +/// SNI rejection happens before ServerHello, so the fatal alert is plaintext +/// even when the ClientHello offers TLS 1.3 (RFC 8446 section 5.1). +pub(super) fn sni_alert(description: u8) -> Vec { + vec![21, 3, 3, 0, 2, 2, description] +} + +pub(super) fn feed_acceptor(acceptor: &mut Acceptor, bytes: &[u8]) -> std::io::Result<()> { + let mut reader = std::io::Cursor::new(bytes); + while reader.position() < bytes.len() as u64 { + if acceptor.read_tls(&mut reader)? == 0 { + return Err(std::io::ErrorKind::UnexpectedEof.into()); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::sync::Arc; + + fn large_client_hello() -> Vec { + let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); + let mut config = rustls::ClientConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions() + .unwrap() + .with_root_certificates(rustls::RootCertStore::empty()) + .with_no_client_auth(); + config.alpn_protocols = (0..70).map(|i| format!("{i:0100}").into_bytes()).collect(); + let mut client = + rustls::ClientConnection::new(Arc::new(config), "localhost".try_into().unwrap()) + .unwrap(); + let mut bytes = Vec::new(); + client.write_tls(&mut bytes).unwrap(); + assert!(bytes.len() > 4096); + bytes + } + + #[test] + fn consumes_large_client_hello_completely() { + let mut acceptor = Acceptor::default(); + feed_acceptor(&mut acceptor, &large_client_hello()).unwrap(); + let accepted = acceptor.accept().unwrap().unwrap(); + assert_eq!(accepted.client_hello().server_name(), Some("localhost")); + assert_eq!(accepted.client_hello().alpn().unwrap().count(), 70); + } + + #[test] + fn waits_for_every_fragment_of_client_hello() { + let hello = large_client_hello(); + let mut acceptor = Acceptor::default(); + for byte in &hello[..hello.len() - 1] { + feed_acceptor(&mut acceptor, &[*byte]).unwrap(); + assert!(acceptor.accept().unwrap().is_none()); + } + feed_acceptor(&mut acceptor, &hello[hello.len() - 1..]).unwrap(); + assert!(acceptor.accept().unwrap().is_some()); + } + + #[test] + fn fatal_sni_alert_is_one_plaintext_record() { + assert_eq!(sni_alert(49), [21, 3, 3, 0, 2, 2, 49]); + assert_eq!(sni_alert(40), [21, 3, 3, 0, 2, 2, 40]); + assert_eq!(sni_alert(80), [21, 3, 3, 0, 2, 2, 80]); + } +} From dd0a184b541c797c8a6a9f9f30d3d7d3177d851b Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sat, 5 Sep 2026 20:53:35 +0900 Subject: [PATCH 2/4] ssl: share socket and MemoryBIO transport progress Queue drained TLS records before fallible output and use one handshake record loop for both transports. Assisted-by: Codex:GPT-5 --- crates/stdlib/src/ssl.rs | 356 ++++++++++------------------- crates/stdlib/src/ssl/compat.rs | 390 ++++---------------------------- 2 files changed, 173 insertions(+), 573 deletions(-) diff --git a/crates/stdlib/src/ssl.rs b/crates/stdlib/src/ssl.rs index 7fac3e7b133..5f2c74f2592 100644 --- a/crates/stdlib/src/ssl.rs +++ b/crates/stdlib/src/ssl.rs @@ -1896,9 +1896,11 @@ mod _ssl { // Create _SSLSocket instance let ssl_socket = PySSLSocket { - sock: args.sock.clone(), - sock_send_method: socket_class.get_attr("send", vm)?, - sock_recv_method: socket_class.get_attr("recv", vm)?, + io: SocketOrBio::Socket { + sock: args.sock.clone(), + send: socket_class.get_attr("send", vm)?, + recv: socket_class.get_attr("recv", vm)?, + }, tls_record_header_buf: vm .ctx .new_bytearray(Vec::with_capacity(TLS_RECORD_HEADER_SIZE)) @@ -1921,8 +1923,6 @@ mod _ssl { chain_builder: PyRwLock::new(None), verified_chain: PyRwLock::new(None), client_session_store: PyRwLock::new(None), - incoming_bio: None, - outgoing_bio: None, shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), @@ -1985,10 +1985,10 @@ mod _ssl { // Create _SSLSocket instance with BIO mode let ssl_socket = PySSLSocket { - // No socket in BIO mode - sock: vm.ctx.none(), - sock_send_method: vm.ctx.none(), - sock_recv_method: vm.ctx.none(), + io: SocketOrBio::Bio { + incoming: args.incoming, + outgoing: args.outgoing, + }, tls_record_header_buf: vm .ctx @@ -2012,8 +2012,6 @@ mod _ssl { chain_builder: PyRwLock::new(None), verified_chain: PyRwLock::new(None), client_session_store: PyRwLock::new(None), - incoming_bio: Some(args.incoming), - outgoing_bio: Some(args.outgoing), shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), @@ -2330,19 +2328,74 @@ mod _ssl { } } + #[derive(Debug)] + enum SocketOrBio { + Socket { + sock: PyObjectRef, + send: PyObjectRef, + recv: PyObjectRef, + }, + Bio { + incoming: PyRef, + outgoing: PyRef, + }, + } + + unsafe impl crate::vm::object::Traverse for SocketOrBio { + fn traverse(&self, tracer: &mut crate::vm::object::TraverseFn<'_>) { + match self { + Self::Socket { sock, send, recv } => { + sock.traverse(tracer); + send.traverse(tracer); + recv.traverse(tracer); + } + Self::Bio { incoming, outgoing } => { + incoming.traverse(tracer); + outgoing.traverse(tracer); + } + } + } + } + + impl SocketOrBio { + fn socket_object(&self, vm: &VirtualMachine) -> PyObjectRef { + match self { + Self::Socket { sock, .. } => sock.clone(), + Self::Bio { .. } => vm.ctx.none(), + } + } + fn incoming(&self) -> Option> { + match self { + Self::Bio { incoming, .. } => Some(incoming.clone()), + _ => None, + } + } + fn recv(&self, size: usize, vm: &VirtualMachine) -> PyResult { + match self { + Self::Socket { sock, recv, .. } => recv.call((sock.clone(), size), vm), + Self::Bio { incoming, .. } => { + incoming.as_object().get_attr("read", vm)?.call((size,), vm) + } + } + } + fn send(&self, data: &[u8], vm: &VirtualMachine) -> PyResult { + let bytes = vm.ctx.new_bytes(data.to_vec()); + match self { + Self::Socket { sock, send, .. } => send.call((sock.clone(), bytes), vm), + Self::Bio { outgoing, .. } => outgoing + .as_object() + .get_attr("write", vm)? + .call((bytes,), vm), + } + } + } + // SSLSocket - represents a TLS-wrapped socket #[pyattr] #[pyclass(name = "_SSLSocket", module = "ssl", traverse)] #[derive(Debug, PyPayload)] pub(crate) struct PySSLSocket { - // Underlying socket - sock: PyObjectRef, - // Cached socket.socket.send - #[pytraverse(skip)] - sock_send_method: PyObjectRef, - // Cached socket.socket.recv - #[pytraverse(skip)] - sock_recv_method: PyObjectRef, + io: SocketOrBio, // Header of currently read TLS record. #[pytraverse(skip)] tls_record_header_buf: PyObjectRef, @@ -2379,9 +2432,6 @@ mod _ssl { // Per-connection store containing the session selected by this connection. #[pytraverse(skip)] client_session_store: PyRwLock>>, - // MemoryBIO mode (optional) - incoming_bio: Option>, - outgoing_bio: Option>, // Shutdown state for tracking close-notify exchange #[pytraverse(skip)] shutdown_state: PyMutex, @@ -2418,12 +2468,16 @@ mod _ssl { impl PySSLSocket { // Check if this is BIO mode pub(crate) fn is_bio_mode(&self) -> bool { - self.incoming_bio.is_some() && self.outgoing_bio.is_some() + matches!(self.io, SocketOrBio::Bio { .. }) } // Get incoming BIO reference (for EOF checking) pub(crate) fn incoming_bio(&self) -> Option { - self.incoming_bio.as_ref().map(|bio| bio.clone().into()) + self.io.incoming().map(|bio| bio.into()) + } + + pub(crate) fn transport_eof(&self) -> bool { + self.io.incoming().is_none_or(|bio| bio.eof()) } // Check for deferred certificate verification errors (TLS 1.3) @@ -2446,7 +2500,11 @@ mod _ssl { } // Get timeout from socket - let timeout_obj = self.sock.get_attr("gettimeout", vm)?.call((), vm)?; + let timeout_obj = self + .io + .socket_object(vm) + .get_attr("gettimeout", vm)? + .call((), vm)?; // timeout can be None (blocking), 0.0 (non-blocking), or positive float if vm.is_none(&timeout_obj) { @@ -2651,7 +2709,7 @@ mod _ssl { } // Use select with the effective timeout - let py_socket: PyRef = self.sock.clone().try_into_value(vm)?; + let py_socket: PyRef = self.io.socket_object(vm).try_into_value(vm)?; let socket = py_socket .sock() .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?; @@ -2678,7 +2736,7 @@ mod _ssl { return Ok(false); } - let py_socket: PyRef = self.sock.clone().try_into_value(vm)?; + let py_socket: PyRef = self.io.socket_object(vm).try_into_value(vm)?; let socket = py_socket .sock() .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?; @@ -2756,7 +2814,7 @@ mod _ssl { if bytes.is_empty() { return Err( if self.is_bio_mode() - && !self.incoming_bio.as_ref().is_some_and(|bio| bio.eof()) + && !self.io.incoming().as_ref().is_some_and(|bio| bio.eof()) { create_ssl_want_read_error(vm).upcast() } else { @@ -2881,17 +2939,10 @@ mod _ssl { } } - // Helper to call socket methods, bypassing any SSL wrapper + // Transport dispatch is shared by handshake, application data, alerts + // and shutdown, bypassing SSLSocket overrides that would reenter TLS. pub(crate) fn sock_recv(&self, size: usize, vm: &VirtualMachine) -> PyResult { - // In BIO mode, read from incoming BIO (flags not supported) - if let Some(ref bio) = self.incoming_bio { - let bio_obj: PyObjectRef = bio.clone().into(); - let read_method = bio_obj.get_attr("read", vm)?; - return read_method.call((vm.ctx.new_int(size),), vm); - } - - self.sock_recv_method - .call((self.sock.clone(), vm.ctx.new_int(size)), vm) + self.io.recv(size, vm) } // Helper to receive data for at most one TLS record. @@ -2957,217 +3008,58 @@ mod _ssl { Ok(collected.map_or(bytes, |bytes| vm.ctx.new_bytes(bytes))) } - /// Socket send - just sends data, caller must handle pending flush - /// Use flush_pending_tls_output before this if ordering is important pub(crate) fn sock_send(&self, data: &[u8], vm: &VirtualMachine) -> PyResult { - // In BIO mode, write to outgoing BIO - if let Some(ref bio) = self.outgoing_bio { - let bio_obj: PyObjectRef = bio.clone().into(); - let write_method = bio_obj.get_attr("write", vm)?; - return write_method.call((vm.ctx.new_bytes(data.to_vec()),), vm); - } - - self.sock_send_method - .call((self.sock.clone(), vm.ctx.new_bytes(data.to_vec())), vm) + self.io.send(data, vm) } - /// Flush any pending TLS output data to the socket - /// Optional deadline parameter allows respecting a read deadline during flush + /// Resume pending output identically for socket and MemoryBIO transports. pub(crate) fn flush_pending_tls_output( &self, vm: &VirtualMachine, deadline: Option, ) -> PyResult<()> { let mut pending = self.pending_tls_output.lock(); - if pending.is_empty() { - return Ok(()); - } - - let socket_timeout = self.get_socket_timeout(vm)?; - let is_non_blocking = socket_timeout.is_some_and(|t| t.is_zero()); - - let mut sent_total = 0; - - while sent_total < pending.len() { - // Calculate timeout: use deadline if provided, otherwise use socket timeout - let timeout_to_use = if let Some(dl) = deadline { - let now = std::time::Instant::now(); - if now >= dl { - // Deadline already passed - *pending = pending[sent_total..].to_vec(); - return Err( - timeout_error_msg(vm, "The operation timed out".to_string()).upcast() - ); - } - Some(dl - now) + while !pending.is_empty() { + let timeout = if let Some(deadline) = deadline { + Some( + deadline + .checked_duration_since(std::time::Instant::now()) + .ok_or_else(|| { + timeout_error_msg(vm, "The write operation timed out".to_owned()) + .upcast() + })?, + ) } else { - socket_timeout + self.get_socket_timeout(vm)? }; - - // Use sock_wait directly with calculated timeout - let py_socket: PyRef = self.sock.clone().try_into_value(vm)?; - let socket = py_socket - .sock() - .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?; - let timed_out = sock_wait(&socket, SockWaitKind::Write, timeout_to_use, vm)?; - - if timed_out { - // Keep unsent data in pending buffer - *pending = pending[sent_total..].to_vec(); - if is_non_blocking { - return Err(create_ssl_want_write_error(vm).upcast()); - } + if self.sock_wait_for_io_with_timeout(SockWaitKind::Write, timeout, vm)? { return Err( - timeout_error_msg(vm, "The write operation timed out".to_string()).upcast(), + timeout_error_msg(vm, "The write operation timed out".to_owned()).upcast(), ); } - - match self.sock_send(&pending[sent_total..], vm) { - Ok(result) => { - let sent: usize = result.try_to_value::(vm)?.try_into().unwrap_or(0); - if sent == 0 { - if is_non_blocking { - // Keep unsent data in pending buffer - *pending = pending[sent_total..].to_vec(); - return Err(create_ssl_want_write_error(vm).upcast()); - } - // Socket said ready but sent 0 bytes - retry - continue; - } - sent_total += sent; - } - Err(e) => { - if is_blocking_io_error(&e, vm) { - if is_non_blocking { - // Keep unsent data in pending buffer - *pending = pending[sent_total..].to_vec(); - return Err(create_ssl_want_write_error(vm).upcast()); - } - continue; - } - // Keep unsent data in pending buffer for other errors too - *pending = pending[sent_total..].to_vec(); - return Err(e); + let result = self.sock_send(&pending, vm).map_err(|e| { + if is_blocking_io_error(&e, vm) { + create_ssl_want_write_error(vm).upcast() + } else { + e } + })?; + let count = result.try_to_value::(vm)?; + if count == 0 { + return Err(create_ssl_want_write_error(vm).upcast()); } + pending.drain(..count); } - - // All data sent successfully - pending.clear(); Ok(()) } - /// Send TLS output data to socket, saving unsent bytes to pending buffer - /// This prevents data loss when rustls' write_tls() drains its internal buffer - /// but the socket cannot accept all the data immediately + /// Queue TLS output through the same partial-write path used by data I/O. fn send_tls_output(&self, buf: Vec, vm: &VirtualMachine) -> PyResult<()> { - if buf.is_empty() { - return Ok(()); - } - - let timeout = self.get_socket_timeout(vm)?; - let is_non_blocking = timeout.is_some_and(|t| t.is_zero()); - - let mut sent_total = 0; - while sent_total < buf.len() { - let timed_out = self.sock_wait_for_io_impl(SockWaitKind::Write, vm)?; - if timed_out { - // Save unsent data to pending buffer - self.pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err( - timeout_error_msg(vm, "The write operation timed out".to_string()).upcast(), - ); - } - - match self.sock_send(&buf[sent_total..], vm) { - Ok(result) => { - let sent: usize = result.try_to_value::(vm)?.try_into().unwrap_or(0); - if sent == 0 { - if is_non_blocking { - // Save unsent data to pending buffer - self.pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(create_ssl_want_write_error(vm).upcast()); - } - continue; - } - sent_total += sent; - } - Err(e) => { - if is_blocking_io_error(&e, vm) { - if is_non_blocking { - // Save unsent data to pending buffer - self.pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(create_ssl_want_write_error(vm).upcast()); - } - continue; - } - // Save unsent data for other errors too - self.pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(e); - } - } - } - - Ok(()) + super::compat::send_all_bytes(self, buf, vm, None).map_err(|e| e.into_py_err(vm)) } - /// Flush all pending TLS output data, respecting socket timeout - /// Used during handshake completion and shutdown() to ensure all data is sent pub(crate) fn blocking_flush_all_pending(&self, vm: &VirtualMachine) -> PyResult<()> { - // Get socket timeout to respect during flush - let timeout = self.get_socket_timeout(vm)?; - if timeout.is_some_and(|t| t.is_zero()) { - return self.flush_pending_tls_output(vm, None); - } - - loop { - let pending_data = { - let pending = self.pending_tls_output.lock(); - if pending.is_empty() { - return Ok(()); - } - pending.clone() - }; - - // Wait for socket to be writable, respecting socket timeout - let py_socket: PyRef = self.sock.clone().try_into_value(vm)?; - let socket = py_socket - .sock() - .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?; - let timed_out = sock_wait(&socket, SockWaitKind::Write, timeout, vm)?; - - if timed_out { - return Err( - timeout_error_msg(vm, "The write operation timed out".to_string()).upcast(), - ); - } - - // Try to send pending data (use raw to avoid recursion) - match self.sock_send(&pending_data, vm) { - Ok(result) => { - let sent: usize = result.try_to_value::(vm)?.try_into().unwrap_or(0); - if sent > 0 { - let mut pending = self.pending_tls_output.lock(); - pending.drain(..sent); - } - // If sent == 0, loop will retry with sock_wait - } - Err(e) => { - if is_blocking_io_error(&e, vm) { - continue; - } - return Err(e); - } - } - } + self.flush_pending_tls_output(vm, None) } // Helper function to convert Python PROTO_* constants to rustls versions @@ -4121,7 +4013,7 @@ mod _ssl { if self.is_bio_mode() { return Ok(vm.ctx.none()); } - return Ok(self.sock.clone()); + return Ok(self.io.socket_object(vm)); } // Get connection @@ -4183,7 +4075,7 @@ mod _ssl { // Check socket timeout mode let timeout_mode = if !is_bio { // Get socket timeout - match self.sock.get_attr("gettimeout", vm) { + match self.io.socket_object(vm).get_attr("gettimeout", vm) { Ok(method) => match method.call((), vm) { Ok(timeout) => { if vm.is_none(&timeout) { @@ -4226,7 +4118,7 @@ mod _ssl { let _ = self.flush_pending_tls_output(vm, None); *self.shutdown_state.lock() = ShutdownState::Completed; *self.connection.lock() = None; - return Ok(self.sock.clone()); + return Ok(self.io.socket_object(vm)); } _ => { // Blocking or timeout mode: wait for peer's close_notify. @@ -4322,7 +4214,7 @@ mod _ssl { // Shutdown complete *self.shutdown_state.lock() = ShutdownState::Completed; *self.connection.lock() = None; - return Ok(self.sock.clone()); + return Ok(self.io.socket_object(vm)); } } } @@ -4347,7 +4239,7 @@ mod _ssl { if is_bio { return Ok(vm.ctx.none()); } - Ok(self.sock.clone()) + Ok(self.io.socket_object(vm)) } // Helper: Write all pending TLS data (including close_notify) to outgoing buffer/BIO @@ -4388,7 +4280,7 @@ mod _ssl { // data. During STARTTLS, after close_notify exchange, the socket // transitions to cleartext. Without peeking, sock_recv may consume // cleartext data meant for the application after unwrap(). - if self.incoming_bio.is_none() { + if self.io.incoming().is_none() { return Ok(self.try_read_close_notify_socket(conn, vm)); } @@ -4399,7 +4291,7 @@ mod _ssl { let data = bytes.as_bytes(); if data.is_empty() { - if let Some(ref bio) = self.incoming_bio { + if let Some(ref bio) = self.io.incoming() { // BIO mode: check if EOF was signaled via write_eof() let bio_obj: PyObjectRef = bio.clone().into(); let eof_attr = bio_obj.get_attr("eof", vm)?; diff --git a/crates/stdlib/src/ssl/compat.rs b/crates/stdlib/src/ssl/compat.rs index d4ec6d65932..b514be2cf49 100644 --- a/crates/stdlib/src/ssl/compat.rs +++ b/crates/stdlib/src/ssl/compat.rs @@ -42,8 +42,9 @@ use super::_ssl::{ // Import error types and helper functions from error module use super::error::{ - PySSLCertVerificationError, PySSLError, create_ssl_eof_error, create_ssl_syscall_error, - create_ssl_want_read_error, create_ssl_want_write_error, create_ssl_zero_return_error, + PySSLCertVerificationError, PySSLError, PySSLWantWriteError, create_ssl_eof_error, + create_ssl_syscall_error, create_ssl_want_read_error, create_ssl_want_write_error, + create_ssl_zero_return_error, }; // OpenSSL Constants: @@ -849,100 +850,25 @@ pub(super) fn is_blocking_io_error(err: &Py, vm: &VirtualMachin /// until all data is sent. For non-blocking sockets, returns WantWrite /// if no progress can be made. /// Optional deadline parameter allows respecting a read deadline during flush. -fn send_all_bytes( +pub(super) fn send_all_bytes( socket: &PySSLSocket, buf: Vec, vm: &VirtualMachine, deadline: Option, ) -> SslResult<()> { - // First, flush any previously pending TLS data with deadline + // Retain newly drained records before a fallible flush of earlier output. + socket.pending_tls_output.lock().extend_from_slice(&buf); socket .flush_pending_tls_output(vm, deadline) - .map_err(SslError::Py)?; - - if buf.is_empty() { - return Ok(()); - } - - let mut sent_total = 0; - while sent_total < buf.len() { - // Check deadline before each send attempt - if let Some(dl) = deadline - && std::time::Instant::now() >= dl - { - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::Timeout("The operation timed out".to_string())); - } - - // Wait for socket to be writable before sending - let timed_out = if let Some(dl) = deadline { - let now = std::time::Instant::now(); - if now >= dl { - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::Timeout( - "The write operation timed out".to_string(), - )); - } - socket - .sock_wait_for_io_with_timeout(SockWaitKind::Write, Some(dl - now), vm) - .map_err(SslError::Py)? - } else { - socket - .sock_wait_for_io_impl(SockWaitKind::Write, vm) - .map_err(SslError::Py)? - }; - if timed_out { - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::Timeout( - "The write operation timed out".to_string(), - )); - } - - match socket.sock_send(&buf[sent_total..], vm) { - Ok(result) => { - let sent: usize = result - .try_to_value::(vm) - .map_err(SslError::Py)? - .try_into() - .map_err(|_| SslError::Syscall("Invalid send return value".to_string()))?; - if sent == 0 { - // No progress - save unsent bytes to pending buffer - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::WantWrite); - } - sent_total += sent; - } - Err(e) => { - if is_blocking_io_error(&e, vm) { - // Save unsent bytes to pending buffer - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::WantWrite); - } - // For other errors, also save unsent bytes - socket - .pending_tls_output - .lock() - .extend_from_slice(&buf[sent_total..]); - return Err(SslError::Py(e)); + .map_err(|error| { + if error.fast_isinstance(PySSLWantWriteError::class(&vm.ctx)) { + SslError::WantWrite + } else if error.fast_isinstance(vm.ctx.exceptions.timeout_error) { + SslError::Timeout("The write operation timed out".to_owned()) + } else { + SslError::Py(error) } - } - } - Ok(()) + }) } // Handshake Helper Functions @@ -954,7 +880,6 @@ fn send_all_bytes( fn handshake_write_loop( conn: &mut Connection, socket: &PySSLSocket, - force_initial_write: bool, vm: &VirtualMachine, ) -> SslResult { let mut made_progress = false; @@ -965,12 +890,7 @@ fn handshake_write_loop( .flush_pending_tls_output(vm, None) .map_err(SslError::Py)?; - while conn.wants_write() || force_initial_write { - if force_initial_write && !conn.wants_write() { - // No data to write on first iteration - break to avoid infinite loop - break; - } - + while conn.wants_write() { let mut buf = Vec::new(); let written = conn .write_tls(&mut buf as &mut dyn std::io::Write) @@ -1023,7 +943,11 @@ fn recv_at_most_one_tls_record( } })?; if bytes.is_empty() { - Err(SslError::Eof) + Err(if socket.is_bio_mode() && !socket.transport_eof() { + SslError::WantRead + } else { + SslError::Eof + }) } else { Ok(bytes.into()) } @@ -1060,128 +984,18 @@ fn recv_at_most_one_tls_record_for_data( fn handshake_read_data( conn: &mut Connection, socket: &PySSLSocket, - is_bio: bool, - vm: &VirtualMachine, -) -> SslResult { - if !conn.wants_read() { - return Ok(false); - } - - // Wait for data in socket mode - if !is_bio { - let timed_out = socket - .sock_wait_for_io_impl(SockWaitKind::Read, vm) - .map_err(SslError::Py)?; - - if timed_out { - // This should rarely happen now - only if socket itself has a timeout - // and we're waiting for required handshake data - return Err(SslError::Timeout("timed out".to_string())); - } - } - - let data_obj = if !is_bio { - // In socket mode, read one TLS record at a time to avoid consuming - // application data that may arrive alongside the final handshake - // record. This matches OpenSSL's default (no read-ahead) behaviour - // and keeps remaining data in the kernel buffer where select() can - // detect it. - recv_at_most_one_tls_record(socket, vm)? - } else { - match socket.sock_recv(SSL3_RT_MAX_PACKET_SIZE, vm) { - Ok(d) => d, - Err(e) => { - if is_blocking_io_error(&e, vm) { - return Err(SslError::WantRead); - } - if !conn.wants_write() && e.fast_isinstance(vm.ctx.exceptions.timeout_error) { - return Ok(false); - } - return Err(SslError::Py(e)); - } - } - }; - - // Feed data to rustls - ssl_read_tls_records(conn, data_obj, is_bio, vm)?; - - Ok(true) -} - -/// Handle handshake completion for server-side TLS 1.3 -/// -/// Tries to send NewSessionTicket in non-blocking mode to avoid deadlocks. -/// Returns true if handshake is complete and we should exit. -fn handle_handshake_complete( - conn: &mut Connection, - socket: &PySSLSocket, - _is_server: bool, vm: &VirtualMachine, -) -> SslResult { - if conn.is_handshaking() { - return Ok(false); // Not complete yet - } - - // Handshake is complete! - // - // Different behavior for BIO mode vs socket mode: - // - // BIO mode (CPython-compatible): - // - Python code calls outgoing.read() to get pending data - // - We just return here and let Python handle the data - // - // Socket mode (rustls-specific): - // - OpenSSL automatically writes to socket in SSL_do_handshake() - // - We must explicitly call write_tls() to send pending data - // - Without this, client hangs waiting for server's NewSessionTicket - - if socket.is_bio_mode() { - // BIO mode: Write pending data to outgoing BIO (one-time drain) - // Python's ssl_io_loop will read from outgoing BIO - if conn.wants_write() { - // Call write_tls ONCE to drain pending data - // Do NOT loop on wants_write() - avoid infinite loop/deadlock - let tls_data = ssl_write_tls_records(conn)?; - if !tls_data.is_empty() { - send_all_bytes(socket, tls_data, vm, None)?; - } - - // IMPORTANT: Don't check wants_write() again! - // Python's ssl_io_loop will call do_handshake() again if needed - } - } else if conn.wants_write() { - // Send all pending data (e.g., TLS 1.3 NewSessionTicket) to socket - // Must drain ALL rustls buffer - don't break on WantWrite - while conn.wants_write() { - let tls_data = ssl_write_tls_records(conn)?; - if tls_data.is_empty() { - break; - } - match send_all_bytes(socket, tls_data, vm, None) { - Ok(()) => {} - Err(SslError::WantWrite) => { - // Socket buffer full, data saved to pending_tls_output - // Flush pending and continue draining rustls buffer - socket - .blocking_flush_all_pending(vm) - .map_err(SslError::Py)?; - } - Err(e) => return Err(e), - } - } - } - - // CRITICAL: Ensure all pending TLS data is sent before returning - // TLS 1.3 Finished must reach server before handshake is considered complete - // Without this, server may not process application data - if !socket.is_bio_mode() { - // Flush pending_tls_output to ensure all TLS data reaches the server - socket - .blocking_flush_all_pending(vm) - .map_err(SslError::Py)?; +) -> SslResult<()> { + if socket + .sock_wait_for_io_impl(SockWaitKind::Read, vm) + .map_err(SslError::Py)? + { + return Err(SslError::Timeout( + "The handshake operation timed out".to_owned(), + )); } - - Ok(true) + let data = recv_at_most_one_tls_record(socket, vm)?; + ssl_read_tls_records(conn, data, socket.is_bio_mode(), vm) } /// Try to read plaintext data from TLS connection buffer @@ -1221,137 +1035,31 @@ pub(super) fn ssl_do_handshake( socket: &PySSLSocket, vm: &VirtualMachine, ) -> SslResult<()> { - // Check if handshake is already done - if !conn.is_handshaking() { - return Ok(()); - } - - let is_bio = socket.is_bio_mode(); - let is_server = matches!(conn, Connection::Server(_)); - let mut first_iteration = true; // Track if this is the first loop iteration loop { - let mut made_progress = false; - - // IMPORTANT: In BIO mode, force initial write even if wants_write() is false - // rustls requires write_tls() to be called to generate ClientHello/ServerHello - let force_initial_write = is_bio && first_iteration; - - // Write TLS handshake data to socket/BIO - let write_progress = handshake_write_loop(conn, socket, force_initial_write, vm)?; - made_progress |= write_progress; - - // Read TLS handshake data from socket/BIO - let read_progress = handshake_read_data(conn, socket, is_bio, vm)?; - made_progress |= read_progress; - - // Process TLS packets (state machine) - if let Err(e) = conn.process_new_packets() { - // Send close_notify on error - if !is_bio { - conn.send_close_notify(); - // Flush any pending TLS data before sending close_notify - let _ = socket.flush_pending_tls_output(vm, None); - // Actually send the close_notify alert using send_all_bytes - // for proper partial send handling and retry logic - if let Ok(alert_data) = ssl_write_tls_records(conn) - && !alert_data.is_empty() - { - let _ = send_all_bytes(socket, alert_data, vm, None); - } - } - - // InvalidMessage during handshake means non-TLS data was received - // before the handshake completed (e.g., HTTP request to TLS server) - if matches!(e, rustls::Error::InvalidMessage(_)) { - return Err(SslError::PreauthData); - } - - // Certificate verification errors are already handled by from_rustls - - return Err(SslError::from_rustls(e)); - } - - // Check if handshake is complete and handle post-handshake processing - // CRITICAL: We do NOT check wants_read() - this matches CPython/OpenSSL behavior! - if handle_handshake_complete(conn, socket, is_server, vm)? { + // Both transports drain writes first and feed complete/partial records + // through the same path. An empty BIO naturally returns WantRead. + handshake_write_loop(conn, socket, vm)?; + if !conn.is_handshaking() { return Ok(()); } - - // In BIO mode, stop after one iteration - if is_bio { - // Before returning WANT error, write any pending TLS data to BIO - // This is critical: if wants_write is true after process_new_packets, - // we need to write that data to the outgoing BIO before returning - if conn.wants_write() { - // Write all pending TLS data to outgoing BIO - loop { - let mut buf = vec![0u8; SSL3_RT_MAX_PACKET_SIZE]; - let n = match conn.write_tls(&mut buf.as_mut_slice()) { - Ok(n) => n, - Err(_) => break, - }; - if n == 0 { - break; - } - // Send to outgoing BIO - send_all_bytes(socket, buf[..n].to_vec(), vm, None)?; - // Check if there's more to write - if !conn.wants_write() { - break; - } - } - // After writing, check if we still want more - // If all data was written, wants_write may now be false - if conn.wants_write() { - // Still need more - return WANT_WRITE - return Err(SslError::WantWrite); + if !conn.wants_read() { + return Err(SslError::WantRead); + } + handshake_read_data(conn, socket, vm)?; + if let Err(error) = conn.process_new_packets() { + // Preserve the existing socket error path until the caller owns + // fatal-alert draining as a lifecycle state. + if !socket.is_bio_mode() { + if let Ok(bytes) = ssl_write_tls_records(conn) { + let _ = send_all_bytes(socket, bytes, vm, None); } - // Otherwise fall through to check wants_read } - - // Check if we need to read - if conn.wants_read() { - return Err(SslError::WantRead); - } - break; - } - - // Mark that we've completed the first iteration - first_iteration = false; - - // Improved loop termination logic: - // Continue looping if: - // 1. Rustls wants more I/O (wants_read or wants_write), OR - // 2. We made progress in this iteration - // - // This is more robust than just checking made_progress, because: - // - Rustls may need multiple iterations to process TLS state machine - // - Network delays may cause temporary "no progress" situations - // - wants_read/wants_write accurately reflect Rustls internal state - let should_continue = conn.wants_read() || conn.wants_write() || made_progress; - - if !should_continue { - break; - } - } - - // If we exit the loop without completing handshake, return appropriate error - if conn.is_handshaking() { - // For non-blocking sockets, return WantRead/WantWrite to signal caller - // should retry when socket is ready. This matches OpenSSL behavior. - if conn.wants_write() { - return Err(SslError::WantWrite); - } - if conn.wants_read() { - return Err(SslError::WantRead); + return Err(if matches!(error, rustls::Error::InvalidMessage(_)) { + SslError::PreauthData + } else { + SslError::from_rustls(error) + }); } - // Neither wants_read nor wants_write - this is a real error - Err(SslError::Syscall( - "SSL handshake failed: incomplete handshake".to_string(), - )) - } else { - // Handshake completed successfully (shouldn't reach here normally) - Ok(()) } } From 1201e22d1a4432aa0a40bd526499dc833645086d Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sat, 5 Sep 2026 21:05:16 +0900 Subject: [PATCH 3/4] ssl: unify shutdown and fatal alert lifecycle Track close-notify progress in TlsState, preserve pending alerts across retries, and share record-boundary shutdown for sockets and MemoryBIO. Retain rustls errors until Python exception conversion. Assisted-by: Codex:GPT-5 --- crates/stdlib/src/ssl.rs | 579 +++++++---------------------- crates/stdlib/src/ssl/compat.rs | 36 +- crates/stdlib/src/ssl/handshake.rs | 33 +- 3 files changed, 179 insertions(+), 469 deletions(-) diff --git a/crates/stdlib/src/ssl.rs b/crates/stdlib/src/ssl.rs index 5f2c74f2592..8e2f98a2a8f 100644 --- a/crates/stdlib/src/ssl.rs +++ b/crates/stdlib/src/ssl.rs @@ -67,8 +67,8 @@ mod _ssl { // Import error types used in this module (others are exposed via pymodule(with(...))) use super::error::{ - PySSLError, create_ssl_eof_error, create_ssl_want_read_error, create_ssl_want_write_error, - create_ssl_zero_return_error, + PySSLError, PySSLWantReadError, PySSLWantWriteError, create_ssl_eof_error, + create_ssl_want_read_error, create_ssl_want_write_error, create_ssl_zero_return_error, }; use alloc::sync::Arc; use core::{ @@ -114,7 +114,7 @@ mod _ssl { ssl_do_handshake, }; - use super::handshake::HandshakeState; + use super::handshake::TlsState; use super::providers::CryptoExt; // Type aliases for better readability @@ -1910,7 +1910,7 @@ mod _ssl { server_side: args.server_side, server_hostname: PyRwLock::new(hostname), connection: PyMutex::new(None), - handshake_state: PyMutex::new(HandshakeState::new(args.server_side)), + state: PyMutex::new(TlsState::new(args.server_side)), session_was_reused: PyMutex::new(false), owner: PyRwLock::new( args.owner @@ -1923,7 +1923,6 @@ mod _ssl { chain_builder: PyRwLock::new(None), verified_chain: PyRwLock::new(None), client_session_store: PyRwLock::new(None), - shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), deferred_cert_error: Arc::new(ParkingRwLock::new(None)), @@ -1999,7 +1998,7 @@ mod _ssl { server_side, server_hostname: PyRwLock::new(hostname), connection: PyMutex::new(None), - handshake_state: PyMutex::new(HandshakeState::new(server_side)), + state: PyMutex::new(TlsState::new(server_side)), session_was_reused: PyMutex::new(false), owner: PyRwLock::new( args.owner @@ -2012,7 +2011,6 @@ mod _ssl { chain_builder: PyRwLock::new(None), verified_chain: PyRwLock::new(None), client_session_store: PyRwLock::new(None), - shutdown_state: PyMutex::new(ShutdownState::NotStarted), pending_tls_output: PyMutex::new(Vec::new()), write_buffered_len: PyMutex::new(0), deferred_cert_error: Arc::new(ParkingRwLock::new(None)), @@ -2413,7 +2411,7 @@ mod _ssl { #[pytraverse(skip)] connection: PyMutex>, // Includes the saved exception while a rejected handshake sends its alert. - handshake_state: PyMutex, + state: PyMutex, // Session was reused (for session resumption tracking) #[pytraverse(skip)] session_was_reused: PyMutex, @@ -2432,9 +2430,6 @@ mod _ssl { // Per-connection store containing the session selected by this connection. #[pytraverse(skip)] client_session_store: PyRwLock>>, - // Shutdown state for tracking close-notify exchange - #[pytraverse(skip)] - shutdown_state: PyMutex, // Pending TLS output buffer for non-blocking sockets // Stores unsent TLS bytes when sock_send() would block // This prevents data loss when write_tls() drains rustls' internal buffer @@ -2453,14 +2448,6 @@ mod _ssl { deferred_cert_error: Arc>>, } - // Shutdown state for tracking close-notify exchange - #[derive(Debug, Clone, Copy, PartialEq)] - enum ShutdownState { - NotStarted, // unwrap() not called yet - SentCloseNotify, // close-notify sent, waiting for peer's response - Completed, // unwrap() completed successfully - } - /// TLS record header size (content_type + version + length). const TLS_RECORD_HEADER_SIZE: usize = 5; @@ -2659,7 +2646,7 @@ mod _ssl { } fn complete_handshake(&self, vm: &VirtualMachine) { - *self.handshake_state.lock() = HandshakeState::Connected; + *self.state.lock() = TlsState::Connected; // Check if session was resumed - get value and release lock immediately let was_resumed = self @@ -2745,15 +2732,15 @@ mod _ssl { } fn handshake_completed(&self) -> bool { - matches!(*self.handshake_state.lock(), HandshakeState::Connected) + matches!( + *self.state.lock(), + TlsState::Connected | TlsState::ShuttingDown | TlsState::ShutDown + ) } - fn reject_handshake(&self, error: PyBaseExceptionRef, bytes: Vec) { - *self.handshake_state.lock() = HandshakeState::SendingAlert { - error, - bytes, - sent: 0, - }; + fn reject_connection(&self, error: PyBaseExceptionRef, bytes: Vec) { + self.pending_tls_output.lock().extend_from_slice(&bytes); + *self.state.lock() = TlsState::SendingAlert { error }; } /// Receive ClientHello and select the context before rustls can generate @@ -2762,19 +2749,16 @@ mod _ssl { loop { // Never hold the state lock across Python callbacks or I/O. // Reentrant handshake calls see an in-progress operation. - let state = core::mem::replace( - &mut *self.handshake_state.lock(), - HandshakeState::CallingSni, - ); + let state = core::mem::replace(&mut *self.state.lock(), TlsState::InProgress); match state { - HandshakeState::WaitingForClientHello(mut acceptor) => { + TlsState::WaitingForClientHello(mut acceptor) => { match acceptor.accept() { Ok(Some(accepted)) => { let name = accepted.client_hello().server_name().map(str::to_owned); if let Err((description, error)) = self.invoke_sni_callback(name.as_deref(), vm) { - self.reject_handshake( + self.reject_connection( error, super::handshake::sni_alert(description), ); @@ -2791,7 +2775,7 @@ mod _ssl { } else { SslError::from_rustls(error) }; - self.reject_handshake(error.into_py_err(vm), bytes); + self.reject_connection(error.into_py_err(vm), bytes); continue; } Ok(None) => {} @@ -2826,56 +2810,25 @@ mod _ssl { .map_err(|e| e.into_pyexception(vm))?; Ok(()) })(); - *self.handshake_state.lock() = - HandshakeState::WaitingForClientHello(acceptor); + *self.state.lock() = TlsState::WaitingForClientHello(acceptor); read?; } - HandshakeState::SendingAlert { - error, - bytes, - mut sent, - } => { - let write = (|| { - if sent == bytes.len() { - return Ok(()); - } - if self.sock_wait_for_io_impl(SockWaitKind::Write, vm)? { - return Err(timeout_error_msg( - vm, - "The handshake operation timed out".to_owned(), - ) - .upcast()); - } - let result = self.sock_send(&bytes[sent..], vm).map_err(|e| { - if is_blocking_io_error(&e, vm) { - create_ssl_want_write_error(vm).upcast() - } else { - e - } - })?; - let count = result.try_to_value::(vm)?; - if count == 0 { - return Err(create_ssl_want_write_error(vm).upcast()); - } - sent += count; - Ok(()) - })(); - let done = sent == bytes.len(); - *self.handshake_state.lock() = HandshakeState::SendingAlert { + TlsState::SendingAlert { error } => { + let result = self.flush_pending_tls_output(vm, None); + *self.state.lock() = TlsState::SendingAlert { error: error.clone(), - bytes, - sent, }; - write?; - if done { - return Err(error); - } + result?; + return Err(error); } - state @ (HandshakeState::Handshaking | HandshakeState::Connected) => { - *self.handshake_state.lock() = state; + state @ (TlsState::Handshaking + | TlsState::Connected + | TlsState::ShuttingDown + | TlsState::ShutDown) => { + *self.state.lock() = state; return Ok(None); } - HandshakeState::CallingSni => { + TlsState::InProgress => { return Err(vm.new_value_error("handshake already in progress")); } } @@ -3053,15 +3006,6 @@ mod _ssl { Ok(()) } - /// Queue TLS output through the same partial-write path used by data I/O. - fn send_tls_output(&self, buf: Vec, vm: &VirtualMachine) -> PyResult<()> { - super::compat::send_all_bytes(self, buf, vm, None).map_err(|e| e.into_py_err(vm)) - } - - pub(crate) fn blocking_flush_all_pending(&self, vm: &VirtualMachine) -> PyResult<()> { - self.flush_pending_tls_output(vm, None) - } - // Helper function to convert Python PROTO_* constants to rustls versions fn get_rustls_versions( minimum: i32, @@ -3296,14 +3240,14 @@ mod _ssl { match accepted.into_connection(config_arc) { Ok(conn) => { *conn_guard = Some(Connection::Server(conn)); - *self.handshake_state.lock() = HandshakeState::Handshaking; + *self.state.lock() = TlsState::Handshaking; } Err((error, mut alert)) => { let mut bytes = Vec::new(); alert .write_all(&mut bytes) .map_err(|e| e.into_pyexception(vm))?; - self.reject_handshake(SslError::from_rustls(error).into_py_err(vm), bytes); + self.reject_connection(SslError::from_rustls(error).into_py_err(vm), bytes); self.accept_client_hello(vm)?; } } @@ -3318,11 +3262,7 @@ mod _ssl { return Ok(()); } - let accepted = if self.server_side { - self.accept_client_hello(vm)? - } else { - None - }; + let accepted = self.accept_client_hello(vm)?; let mut conn_guard = self.connection.lock(); // Initialize connection if not already done @@ -3335,13 +3275,10 @@ mod _ssl { vm, ) { drop(conn_guard); - if matches!( - *self.handshake_state.lock(), - HandshakeState::SendingAlert { .. } - ) { + if matches!(*self.state.lock(), TlsState::SendingAlert { .. }) { return Err(error); } - self.reject_handshake(error, super::handshake::sni_alert(40)); + self.reject_connection(error, super::handshake::sni_alert(40)); return self.accept_client_hello(vm).map(|_| ()); } } else { @@ -3454,8 +3391,19 @@ mod _ssl { // Perform the actual handshake by exchanging data with the socket/BIO - let conn = conn_guard.as_mut().expect("unreachable"); + let conn = conn_guard + .as_mut() + .ok_or_else(|| vm.new_value_error("TLS connection is not available"))?; let handshake_result = ssl_do_handshake(conn, self, vm); + if let Err(error @ (SslError::Rustls(_) | SslError::PreauthData)) = handshake_result { + // rustls queued the fatal alert while processing the failing + // record. Retain it and the original error across output retries. + let mut bytes = Vec::new(); + let _ = conn.write_tls(&mut bytes); + drop(conn_guard); + self.reject_connection(error.into_py_err(vm), bytes); + return self.accept_client_hello(vm).map(|_| ()); + } drop(conn_guard); handshake_result.map_err(|e| e.into_py_err(vm))?; self.complete_handshake(vm); @@ -3510,8 +3458,7 @@ mod _ssl { // Check if connection has been shut down // Only block after shutdown is COMPLETED, not during shutdown process - let shutdown_state = *self.shutdown_state.lock(); - if shutdown_state == ShutdownState::Completed { + if matches!(*self.state.lock(), TlsState::ShutDown) { return Err(vm .new_os_subtype_error( PySSLError::class(&vm.ctx).to_owned(), @@ -3563,7 +3510,7 @@ mod _ssl { buf.truncate(n); return_data(buf, &buffer, vm) } - Err(crate::ssl::compat::SslError::Eof) => { + Err(error) if error.is_eof() => { // If plaintext is still buffered, return it before EOF. let pending = { let mut conn_guard = self.connection.lock(); @@ -3592,7 +3539,7 @@ mod _ssl { // EOF occurred in violation of protocol (unexpected closure) Err(create_ssl_eof_error(vm).upcast()) } - Err(crate::ssl::compat::SslError::ZeroReturn) => { + Err(error) if error.is_zero_return() => { // If plaintext is still buffered, return it before clean EOF. let pending = { let mut conn_guard = self.connection.lock(); @@ -3623,10 +3570,10 @@ mod _ssl { // raise SSLZeroReturnError (bidirectional shutdown). // Otherwise return empty bytes, which callers (asyncore, // asyncio sslproto) interpret as EOF. - let our_shutdown_state = *self.shutdown_state.lock(); - if our_shutdown_state == ShutdownState::SentCloseNotify - || our_shutdown_state == ShutdownState::Completed - { + if matches!( + *self.state.lock(), + TlsState::ShuttingDown | TlsState::ShutDown + ) { Err(create_ssl_zero_return_error(vm).upcast()) } else { return_data(vec![], &buffer, vm) @@ -3692,9 +3639,11 @@ mod _ssl { self.do_handshake(vm)?; } - // Check shutdown state - // Only block after shutdown is COMPLETED, not during shutdown process - if *self.shutdown_state.lock() == ShutdownState::Completed { + // Application data cannot follow our close_notify. + if matches!( + *self.state.lock(), + TlsState::ShuttingDown | TlsState::ShutDown + ) { return Err(vm .new_os_subtype_error( PySSLError::class(&vm.ctx).to_owned(), @@ -3997,6 +3946,15 @@ mod _ssl { #[pymethod] fn shutdown(&self, vm: &VirtualMachine) -> PyResult { + // A failed shutdown may still owe the peer a fatal alert. Resume + // that output on unwrap retries before reporting the saved error. + if matches!(*self.state.lock(), TlsState::SendingAlert { .. }) + && !self.pending_tls_output.lock().is_empty() + { + return self + .accept_client_hello(vm) + .map(|_| self.io.socket_object(vm)); + } if !self.handshake_completed() { return Err(SslError::create_ssl_error_with_reason( vm, @@ -4005,356 +3963,89 @@ mod _ssl { "[SSL: SHUTDOWN_WHILE_IN_INIT] shutdown while in init", )); } - // Check current shutdown state - let current_state = *self.shutdown_state.lock(); - - // If already completed, return immediately - if current_state == ShutdownState::Completed { - if self.is_bio_mode() { - return Ok(vm.ctx.none()); - } + if matches!(*self.state.lock(), TlsState::ShutDown) { return Ok(self.io.socket_object(vm)); } - // Get connection + let timeout = self.get_socket_timeout(vm)?; + let deadline = timeout + .filter(|timeout| !timeout.is_zero()) + .map(|timeout| std::time::Instant::now() + timeout); let mut conn_guard = self.connection.lock(); let conn = conn_guard .as_mut() .ok_or_else(|| vm.new_value_error("Connection not established"))?; - let is_bio = self.is_bio_mode(); - - // Step 1: Send our close_notify if not already sent - if current_state == ShutdownState::NotStarted { - // First, flush ALL pending TLS data BEFORE sending close_notify - // This is CRITICAL - close_notify must come AFTER all application data - // Otherwise data loss occurs when peer receives close_notify first - - // Step 1a: Flush any pending TLS records from rustls internal buffer - // This ensures all application data is converted to TLS records - while conn.wants_write() { - let mut buf = Vec::new(); - conn.write_tls(&mut buf) - .map_err(|e| vm.new_os_error(format!("TLS write failed: {e}")))?; - if !buf.is_empty() { - self.send_tls_output(buf, vm)?; - } - } - - // Step 1b: Flush pending_tls_output buffer to socket - if !is_bio { - // Socket mode: blocking flush to ensure data order - // Must complete before sending close_notify - self.blocking_flush_all_pending(vm)?; - } else { - // BIO mode: non-blocking flush (caller handles pending data) - let _ = self.flush_pending_tls_output(vm, None); - } - + if matches!(*self.state.lock(), TlsState::Connected) { + // rustls queues close_notify after previously buffered data. + // Record this before any fallible write so retries never queue + // another alert, even if the transport cannot accept output. conn.send_close_notify(); - - // Write close_notify to outgoing buffer/BIO - self.write_pending_tls(conn, vm)?; - // Ensure close_notify and any pending TLS data are flushed - if !is_bio { - self.flush_pending_tls_output(vm, None)?; - } - - // Update state - *self.shutdown_state.lock() = ShutdownState::SentCloseNotify; - } - - // Step 2: Try to read and process peer's close_notify - - // First check if we already have peer's close_notify - // This can happen if it was received during a previous read() call - let mut peer_closed = self.check_peer_closed(conn, vm)?; - - // If peer hasn't closed yet, try to read from socket - if !peer_closed { - // Check socket timeout mode - let timeout_mode = if !is_bio { - // Get socket timeout - match self.io.socket_object(vm).get_attr("gettimeout", vm) { - Ok(method) => match method.call((), vm) { - Ok(timeout) => { - if vm.is_none(&timeout) { - // timeout=None means blocking - Some(None) - } else if let Ok(t) = timeout.try_float(vm).map(|f| f.to_f64()) { - if t == 0.0 { - // timeout=0 means non-blocking - Some(Some(0.0)) - } else { - // timeout>0 means timeout mode - Some(Some(t)) - } - } else { - None - } - } - Err(_) => None, - }, - Err(_) => None, - } - } else { - None // BIO mode - }; - - if is_bio { - // In BIO mode: non-blocking read attempt - if self.try_read_close_notify(conn, vm)? { - peer_closed = true; - } - } else if let Some(timeout) = timeout_mode { - match timeout { - Some(0.0) => { - // Non-blocking: return immediately after sending close_notify. - // Don't wait for peer's close_notify to avoid blocking. - drop(conn_guard); - // Best-effort flush; WouldBlock is expected in non-blocking mode. - // Other errors indicate close_notify may not have been sent, - // but we still complete shutdown to avoid inconsistent state. - let _ = self.flush_pending_tls_output(vm, None); - *self.shutdown_state.lock() = ShutdownState::Completed; - *self.connection.lock() = None; - return Ok(self.io.socket_object(vm)); - } - _ => { - // Blocking or timeout mode: wait for peer's close_notify. - // This is proper TLS shutdown - we should receive peer's - // close_notify before closing the connection. - drop(conn_guard); - - // Flush our close_notify first - if timeout.is_none() { - self.blocking_flush_all_pending(vm)?; - } else { - self.flush_pending_tls_output(vm, None)?; - } - - // Calculate deadline for timeout mode - let deadline = timeout.map(|t| { - std::time::Instant::now() + core::time::Duration::from_secs_f64(t) - }); - - // Wait for peer's close_notify - loop { - // Re-acquire connection lock for each iteration - let mut conn_guard = self.connection.lock(); - let conn = match conn_guard.as_mut() { - Some(c) => c, - None => break, // Connection already closed - }; - - // Check if peer already sent close_notify - if self.check_peer_closed(conn, vm)? { - break; - } - - drop(conn_guard); - - // Check timeout - let remaining_timeout = if let Some(dl) = deadline { - let now = std::time::Instant::now(); - if now >= dl { - // Timeout reached - raise TimeoutError - return Err(timeout_error_msg( - vm, - "The read operation timed out".to_string(), - ) - .upcast()); - } - Some(dl - now) - } else { - None // Blocking mode: no timeout - }; - - // Wait for socket to be readable - let timed_out = self.sock_wait_for_io_with_timeout( - SockWaitKind::Read, - remaining_timeout, - vm, - )?; - - if timed_out { - // Timeout waiting for peer's close_notify - return Err(timeout_error_msg( - vm, - "The read operation timed out".to_string(), - ) - .upcast()); - } - - // Try to read data from socket - let mut conn_guard = self.connection.lock(); - let conn = match conn_guard.as_mut() { - Some(c) => c, - None => break, - }; - - // Read and process any incoming TLS data - match self.try_read_close_notify(conn, vm) { - Ok(closed) => { - if closed { - break; - } - // Check again after processing - if self.check_peer_closed(conn, vm)? { - break; - } - } - Err(_) => { - // Socket error - peer likely closed connection - break; - } - } - } - - // Shutdown complete - *self.shutdown_state.lock() = ShutdownState::Completed; - *self.connection.lock() = None; - return Ok(self.io.socket_object(vm)); - } - } - } - - // Step 3: Check again if peer has sent close_notify (non-blocking/BIO mode only) - if !peer_closed { - peer_closed = self.check_peer_closed(conn, vm)?; - } - } - - drop(conn_guard); // Release lock before returning - - if !peer_closed { - // Still waiting for peer's close-notify - // Raise SSLWantReadError to signal app needs to transfer data - // This is correct for non-blocking sockets and BIO mode - return Err(create_ssl_want_read_error(vm).upcast()); - } - // Both close-notify exchanged, shutdown complete - *self.shutdown_state.lock() = ShutdownState::Completed; - - if is_bio { - return Ok(vm.ctx.none()); - } - Ok(self.io.socket_object(vm)) - } - - // Helper: Write all pending TLS data (including close_notify) to outgoing buffer/BIO - fn write_pending_tls(&self, conn: &mut Connection, vm: &VirtualMachine) -> PyResult<()> { - // First, flush any previously pending TLS output - // Must succeed before sending new data to maintain order - self.flush_pending_tls_output(vm, None)?; - - loop { - if !conn.wants_write() { - break; - } - - let mut buf = vec![0u8; SSL3_RT_MAX_PACKET_SIZE]; - let written = conn - .write_tls(&mut buf.as_mut_slice()) - .map_err(|e| vm.new_os_error(format!("TLS write failed: {e}")))?; - - if written == 0 { - break; - } - - // Send TLS data, saving unsent bytes to pending buffer if needed - self.send_tls_output(buf[..written].to_vec(), vm)?; - } - - Ok(()) - } - - // Helper: Try to read incoming data from socket/BIO - // Returns true if peer closed connection (with or without close_notify) - fn try_read_close_notify( - &self, - conn: &mut Connection, - vm: &VirtualMachine, - ) -> PyResult { - // In socket mode, peek first to avoid consuming post-TLS cleartext - // data. During STARTTLS, after close_notify exchange, the socket - // transitions to cleartext. Without peeking, sock_recv may consume - // cleartext data meant for the application after unwrap(). - if self.io.incoming().is_none() { - return Ok(self.try_read_close_notify_socket(conn, vm)); + *self.state.lock() = TlsState::ShuttingDown; } - // BIO mode: stop at the TLS record boundary so cleartext queued - // after close_notify remains available to the BIO's owner. - match self.sock_recv_at_most_one_tls_record(vm) { - Ok(bytes) => { - let data = bytes.as_bytes(); - - if data.is_empty() { - if let Some(ref bio) = self.io.incoming() { - // BIO mode: check if EOF was signaled via write_eof() - let bio_obj: PyObjectRef = bio.clone().into(); - let eof_attr = bio_obj.get_attr("eof", vm)?; - let is_eof = eof_attr.try_to_bool(vm)?; - if !is_eof { - return Ok(false); - } - } - return Ok(true); - } + let result = (|| loop { + let mut bytes = Vec::new(); + conn.write_tls(&mut bytes) + .map_err(|e| e.into_pyexception(vm))?; + super::compat::send_all_bytes(self, bytes, vm, deadline) + .map_err(|e| e.into_py_err(vm))?; - let data_slice: &[u8] = data; - let mut cursor = std::io::Cursor::new(data_slice); - let _ = conn.read_tls(&mut cursor); - let _ = conn.process_new_packets(); - Ok(false) + let io_state = conn + .process_new_packets() + .map_err(|e| SslError::from_rustls(e).into_py_err(vm))?; + if io_state.plaintext_bytes_to_read() != 0 { + return Err(SslError::create_ssl_error_with_reason( + vm, + Some("SSL"), + "APPLICATION_DATA_AFTER_CLOSE_NOTIFY", + "[SSL: APPLICATION_DATA_AFTER_CLOSE_NOTIFY] application data after close notify", + )); } - Err(e) => { - if is_blocking_io_error(&e, vm) { - return Ok(false); - } - Ok(true) + if io_state.peer_has_closed() { + *self.state.lock() = TlsState::ShutDown; + return Ok(self.io.socket_object(vm)); } - } - } - /// Socket-mode close_notify reader that respects TLS record boundaries. - /// Uses MSG_PEEK to inspect data before consuming, preventing accidental - /// consumption of post-TLS cleartext data during STARTTLS transitions. - /// - /// Equivalent to OpenSSL's `SSL_set_read_ahead(ssl, 0)` — rustls has no - /// such knob, so we enforce record-level reads manually via peek. - fn try_read_close_notify_socket(&self, conn: &mut Connection, vm: &VirtualMachine) -> bool { - // Consume at most one TLS record from the socket - match self.sock_recv_at_most_one_tls_record(vm) { - Ok(data) => { - if data.is_empty() { - return true; - } - - let data_slice: &[u8] = data.as_ref(); - let mut cursor = std::io::Cursor::new(data_slice); - let _ = conn.read_tls(&mut cursor); - let _ = conn.process_new_packets(); - false - } - Err(e) => { - if is_blocking_io_error(&e, vm) { - return false; - } - true + let remaining = match deadline { + Some(deadline) => Some( + deadline + .checked_duration_since(std::time::Instant::now()) + .ok_or_else(|| { + timeout_error_msg(vm, "The read operation timed out".to_owned()) + .upcast() + })?, + ), + None => timeout, + }; + if self.sock_wait_for_io_with_timeout(SockWaitKind::Read, remaining, vm)? { + return Err( + timeout_error_msg(vm, "The read operation timed out".to_owned()).upcast(), + ); } - } - } - - // Helper: Check if peer has sent close_notify - fn check_peer_closed(&self, conn: &mut Connection, vm: &VirtualMachine) -> PyResult { - // Process any remaining packets and check peer_has_closed - let io_state = conn - .process_new_packets() - .map_err(|e| vm.new_os_error(format!("Failed to process packets: {e}")))?; - - Ok(io_state.peer_has_closed()) + // Stop at each TLS record boundary, leaving an unencrypted + // trailer after close_notify in the socket or incoming BIO. + let data = super::compat::recv_at_most_one_tls_record(self, vm) + .map_err(|e| e.into_py_err(vm))?; + super::compat::ssl_read_tls_records(conn, data, self.is_bio_mode(), vm) + .map_err(|e| e.into_py_err(vm))?; + })(); + if let Err(error) = &result + && error.fast_isinstance(PySSLError::class(&vm.ctx)) + && !error.fast_isinstance(PySSLWantReadError::class(&vm.ctx)) + && !error.fast_isinstance(PySSLWantWriteError::class(&vm.ctx)) + { + // A protocol failure cannot resume as a normal shutdown. Keep + // any fatal alert rustls queued and preserve the original error. + let mut bytes = Vec::new(); + let _ = conn.write_tls(&mut bytes); + drop(conn_guard); + self.reject_connection(error.clone(), bytes); + return self + .accept_client_hello(vm) + .map(|_| self.io.socket_object(vm)); + } + result } #[pymethod] diff --git a/crates/stdlib/src/ssl/compat.rs b/crates/stdlib/src/ssl/compat.rs index b514be2cf49..d12416ef794 100644 --- a/crates/stdlib/src/ssl/compat.rs +++ b/crates/stdlib/src/ssl/compat.rs @@ -201,6 +201,8 @@ pub(super) enum SslError { Timeout(String), /// Python exception (pass through directly) Py(PyBaseExceptionRef), + /// Preserve the TLS error until the Python exception boundary. + Rustls(rustls::Error), /// TLS alert received with OpenSSL-compatible error code AlertReceived { lib: i32, reason: i32 }, /// NO_SHARED_CIPHER error (OpenSSL SSL_R_NO_SHARED_CIPHER) @@ -217,6 +219,28 @@ impl SslError { /// Convert rustls error to SslError pub(super) fn from_rustls(err: rustls::Error) -> Self { + Self::Rustls(err) + } + + pub(super) fn is_eof(&self) -> bool { + match self { + Self::Rustls(error) => matches!(Self::map_rustls(error.clone()), Self::Eof), + Self::Eof => true, + _ => false, + } + } + + pub(super) fn is_zero_return(&self) -> bool { + matches!( + self, + Self::ZeroReturn + | Self::Rustls(rustls::Error::AlertReceived( + rustls::AlertDescription::CloseNotify + )) + ) + } + + fn map_rustls(err: rustls::Error) -> Self { match err { rustls::Error::InvalidCertificate(cert_err) => Self::CertVerification(cert_err), rustls::Error::AlertReceived(alert_desc) => { @@ -346,6 +370,7 @@ impl SslError { /// Convert to Python exception pub(super) fn into_py_err(self, vm: &VirtualMachine) -> PyBaseExceptionRef { match self { + Self::Rustls(error) => Self::map_rustls(error).into_py_err(vm), Self::WantRead => create_ssl_want_read_error(vm).upcast(), Self::WantWrite => create_ssl_want_write_error(vm).upcast(), Self::Timeout(msg) => timeout_error_msg(vm, msg).upcast(), @@ -931,7 +956,7 @@ fn handshake_write_loop( /// Fix: peek at the TCP buffer to find the first complete TLS record boundary /// and recv() only that many bytes. Any remaining data stays in the kernel /// buffer and remains visible to select(). -fn recv_at_most_one_tls_record( +pub(super) fn recv_at_most_one_tls_record( socket: &PySSLSocket, vm: &VirtualMachine, ) -> SslResult { @@ -1047,13 +1072,6 @@ pub(super) fn ssl_do_handshake( } handshake_read_data(conn, socket, vm)?; if let Err(error) = conn.process_new_packets() { - // Preserve the existing socket error path until the caller owns - // fatal-alert draining as a lifecycle state. - if !socket.is_bio_mode() { - if let Ok(bytes) = ssl_write_tls_records(conn) { - let _ = send_all_bytes(socket, bytes, vm, None); - } - } return Err(if matches!(error, rustls::Error::InvalidMessage(_)) { SslError::PreauthData } else { @@ -1454,7 +1472,7 @@ fn ssl_write_tls_records(conn: &mut Connection) -> SslResult> { } /// Read TLS records from socket to rustls -fn ssl_read_tls_records( +pub(super) fn ssl_read_tls_records( conn: &mut Connection, data: PyObjectRef, is_bio: bool, diff --git a/crates/stdlib/src/ssl/handshake.rs b/crates/stdlib/src/ssl/handshake.rs index 414fa5e1e55..f07a20f7119 100644 --- a/crates/stdlib/src/ssl/handshake.rs +++ b/crates/stdlib/src/ssl/handshake.rs @@ -6,44 +6,45 @@ use rustpython_vm::{ }; /// Server configuration is selected only after receiving ClientHello and -/// invoking SNI. A rejected handshake stays terminal while its alert drains. -pub(super) enum HandshakeState { - WaitingForClientHello(Acceptor), - CallingSni, +/// invoking SNI. A failed connection stays terminal while its alert drains. +/// ShuttingDown means our close_notify is queued and must not be sent twice. +pub(super) enum TlsState { + WaitingForClientHello(Box), + InProgress, Handshaking, Connected, - SendingAlert { - error: PyBaseExceptionRef, - bytes: Vec, - sent: usize, - }, + ShuttingDown, + ShutDown, + SendingAlert { error: PyBaseExceptionRef }, } -impl HandshakeState { +impl TlsState { pub(super) fn new(server_side: bool) -> Self { if server_side { - Self::WaitingForClientHello(Acceptor::default()) + Self::WaitingForClientHello(Box::default()) } else { Self::Handshaking } } } -impl fmt::Debug for HandshakeState { +impl fmt::Debug for TlsState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(match self { Self::WaitingForClientHello(_) => "WaitingForClientHello", - Self::CallingSni => "CallingSni", + Self::InProgress => "InProgress", Self::Handshaking => "Handshaking", Self::Connected => "Connected", + Self::ShuttingDown => "ShuttingDown", + Self::ShutDown => "ShutDown", Self::SendingAlert { .. } => "SendingAlert", }) } } -// Only the saved Python exception can contain GC references. Acceptor and the -// encoded TLS alert contain Rust-owned protocol data. -unsafe impl Traverse for HandshakeState { +// Only the saved Python exception can contain GC references. Acceptor contains +// Rust-owned protocol data; encoded alerts use the shared pending output queue. +unsafe impl Traverse for TlsState { fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) { if let Self::SendingAlert { error, .. } = self { error.traverse(tracer_fn); From e939d92e4b38ff2d2cd30ce7bc5cd742f3b116a1 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sat, 5 Sep 2026 22:26:32 +0900 Subject: [PATCH 4/4] ssl: emit the full peer certificate dict Use the same certificate dict helper for getpeercert() and Certificate.get_info() so OCSP, caIssuers, and CRL URLs are present. Clear CHANNEL_BINDING_TYPES on the rustls backend, and add a MemoryBIO leftover-plaintext snippet. Assisted-by: Claude --- Lib/ssl.py | 3 + Lib/test/test_asyncio/test_events.py | 1 - Lib/test/test_ssl.py | 1 - crates/stdlib/src/ssl.rs | 10 +-- crates/stdlib/src/ssl/cert.rs | 122 +-------------------------- extra_tests/snippets/stdlib_ssl.py | 46 ++++++++++ 6 files changed, 54 insertions(+), 129 deletions(-) create mode 100644 extra_tests/snippets/stdlib_ssl.py diff --git a/Lib/ssl.py b/Lib/ssl.py index 8889aff92fa..a78d2a696bc 100644 --- a/Lib/ssl.py +++ b/Lib/ssl.py @@ -268,6 +268,9 @@ class _TLSMessageType: socket_error = OSError # keep that public name in module namespace CHANNEL_BINDING_TYPES = ['tls-unique'] +# TODO: RUSTPYTHON; rustls does not expose TLS Finished messages +if "rustls" in OPENSSL_VERSION: + CHANNEL_BINDING_TYPES = [] HAS_NEVER_CHECK_COMMON_NAME = hasattr(_ssl, 'HOSTFLAG_NEVER_CHECK_SUBJECT') diff --git a/Lib/test/test_asyncio/test_events.py b/Lib/test/test_asyncio/test_events.py index f373483066a..4b4927f6c5d 100644 --- a/Lib/test/test_asyncio/test_events.py +++ b/Lib/test/test_asyncio/test_events.py @@ -1289,7 +1289,6 @@ def test_create_unix_server_ssl_verified(self): server.close() self.loop.run_until_complete(proto.done) - @unittest.expectedFailure # TODO: RUSTPYTHON; - SSL peer certificate format differs @unittest.skipIf(ssl is None, 'No ssl module') def test_create_server_ssl_verified(self): proto = MyProto(loop=self.loop) diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py index f7ecabd853c..3b1a3c62728 100644 --- a/Lib/test/test_ssl.py +++ b/Lib/test/test_ssl.py @@ -4192,7 +4192,6 @@ def test_default_ecdh_curve(self): s.connect((HOST, server.port)) self.assertIn("ECDH", s.cipher()[0]) - @unittest.expectedFailure # TODO: RUSTPYTHON @unittest.skipUnless("tls-unique" in ssl.CHANNEL_BINDING_TYPES, "'tls-unique' channel binding not available") def test_tls_unique_channel_binding(self): diff --git a/crates/stdlib/src/ssl.rs b/crates/stdlib/src/ssl.rs index 8e2f98a2a8f..bd072d45d77 100644 --- a/crates/stdlib/src/ssl.rs +++ b/crates/stdlib/src/ssl.rs @@ -3723,11 +3723,7 @@ mod _ssl { return Ok(Some(vm.ctx.new_dict().into())); } - // Parse DER certificate and convert to dict (outside lock) - let (_, cert) = x509_parser::parse_x509_certificate(&der_bytes) - .map_err(|e| vm.new_value_error(format!("Failed to parse certificate: {e}")))?; - - cert::cert_to_dict(vm, &cert).map(Some) + cert::cert_der_to_dict_helper(vm, &der_bytes).map(Some) } #[pymethod] @@ -4643,7 +4639,6 @@ mod _ssl { } impl PySSLCertificate { - // Parse the certificate lazily fn parse(&self) -> Result, String> { match x509_parser::parse_x509_certificate(&self.der_bytes) { Ok((_, cert)) => Ok(cert), @@ -4682,8 +4677,7 @@ mod _ssl { #[pymethod] fn get_info(&self, vm: &VirtualMachine) -> PyResult { - let cert = self.parse().map_err(|e| vm.new_value_error(e))?; - cert::cert_to_dict(vm, &cert) + cert::cert_der_to_dict_helper(vm, &self.der_bytes) } } diff --git a/crates/stdlib/src/ssl/cert.rs b/crates/stdlib/src/ssl/cert.rs index abe6d244af3..d10a652c67c 100644 --- a/crates/stdlib/src/ssl/cert.rs +++ b/crates/stdlib/src/ssl/cert.rs @@ -228,45 +228,6 @@ fn normalize_wildcard_hostname(hostname: &str) -> &str { hostname.strip_prefix("*.").unwrap_or(hostname) } -/// Process Subject Alternative Name (SAN) general names into Python tuples -/// -/// Converts X.509 GeneralName entries into Python tuple format. -/// Returns a vector of PyObjectRef tuples in the format: (type, value) -fn process_san_general_names( - vm: &VirtualMachine, - general_names: &[GeneralName<'_>], -) -> Vec { - general_names - .iter() - .filter_map(|name| match name { - GeneralName::DNSName(dns) => Some(vm.new_tuple(("DNS", *dns)).into()), - GeneralName::IPAddress(ip) => { - let ip_str = format_ip_address(ip); - Some(vm.new_tuple(("IP Address", ip_str)).into()) - } - GeneralName::RFC822Name(email) => Some(vm.new_tuple(("email", *email)).into()), - GeneralName::URI(uri) => Some(vm.new_tuple(("URI", *uri)).into()), - GeneralName::DirectoryName(dn) => { - let dn_str = format!("{dn}"); - Some(vm.new_tuple(("DirName", dn_str)).into()) - } - GeneralName::RegisteredID(oid) => { - let oid_str = oid.to_string(); - Some(vm.new_tuple(("Registered ID", oid_str)).into()) - } - GeneralName::OtherName(oid, value) => { - let oid_str = oid.to_string(); - let value_str = format!("{value:?}"); - Some( - vm.new_tuple(("othername", format!("{oid_str}:{value_str}"))) - .into(), - ) - } - _ => None, - }) - .collect() -} - // Certificate Validation and Parsing /// Check if a certificate is a CA certificate by examining the Basic Constraints extension @@ -294,87 +255,10 @@ pub(super) fn is_ca_certificate(cert_der: &[u8]) -> bool { cert.version().0 == 0 && cert.subject() == cert.issuer() } -/// Convert an X509Name to Python nested tuple format for SSL certificate dicts -/// -/// Format: ((('CN', 'example.com'),), (('O', 'Example Org'),), ...) -fn name_to_py(vm: &VirtualMachine, name: &x509_parser::x509::X509Name<'_>) -> PyObjectRef { - let list = name - .iter() - .flat_map(|rdn| { - // Each RDN can have multiple attributes - rdn.iter() - .map(|attr| { - let oid_str = attr.attr_type().to_id_string(); - let value_str = attr.attr_value().as_str().unwrap_or("").to_string(); - let key = oid_to_attribute_name(&oid_str); - - vm.new_tuple((vm.new_tuple((vm.ctx.new_str(key), vm.ctx.new_str(value_str))),)) - .into() - }) - .collect::>() - }) - .collect::>(); - - vm.ctx.new_tuple(list).into() -} - -/// Convert DER-encoded certificate to Python dict (for getpeercert with binary_form=False) -/// -/// Returns a dict with fields: subject, issuer, version, serialNumber, -/// notBefore, notAfter, subjectAltName (if present) -pub(super) fn cert_to_dict( - vm: &VirtualMachine, - cert: &x509_parser::certificate::X509Certificate<'_>, -) -> PyResult { - let dict = vm.ctx.new_dict(); - - // Subject and Issuer - dict.set_item("subject", name_to_py(vm, cert.subject()), vm)?; - dict.set_item("issuer", name_to_py(vm, cert.issuer()), vm)?; - - // Version (X.509 v3 = version 2 in the cert, but Python uses 3) - dict.set_item( - "version", - vm.ctx.new_int(cert.version().0 as i32 + 1).into(), - vm, - )?; - - // Serial number - hex format with even length - let serial = format_serial_number(&cert.serial); - dict.set_item("serialNumber", vm.ctx.new_str(serial).into(), vm)?; - - // Validity dates - format with GMT using jiff - dict.set_item( - "notBefore", - vm.ctx - .new_str(format_asn1_time(&cert.validity().not_before)) - .into(), - vm, - )?; - dict.set_item( - "notAfter", - vm.ctx - .new_str(format_asn1_time(&cert.validity().not_after)) - .into(), - vm, - )?; - - // Subject Alternative Names (if present) - if let Ok(Some(san_ext)) = cert.subject_alternative_name() { - let san_list = process_san_general_names(vm, &san_ext.value.general_names); - - if !san_list.is_empty() { - dict.set_item("subjectAltName", vm.ctx.new_tuple(san_list).into(), vm)?; - } - } - - Ok(dict.into()) -} - -/// Convert DER-encoded certificate to Python dict (for get_ca_certs) +/// Convert DER-encoded certificate to Python dict. /// -/// Similar to cert_to_dict but includes additional fields like crlDistributionPoints -/// and uses CPython's specific ordering: issuer, notAfter, notBefore, serialNumber, subject, version +/// Includes OCSP, caIssuers, and crlDistributionPoints when present, and uses +/// the field order: issuer, notAfter, notBefore, serialNumber, subject, version pub(super) fn cert_der_to_dict_helper( vm: &VirtualMachine, cert_der: &[u8], diff --git a/extra_tests/snippets/stdlib_ssl.py b/extra_tests/snippets/stdlib_ssl.py new file mode 100644 index 00000000000..cad505eebc1 --- /dev/null +++ b/extra_tests/snippets/stdlib_ssl.py @@ -0,0 +1,46 @@ +import pathlib +import ssl + +CERT = ( + pathlib.Path(__file__).resolve().parent.parent.parent + / "Lib" + / "test" + / "certdata" + / "keycert.pem" +) + +TAIL = b"tail" + +server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) +server_context.load_cert_chain(str(CERT)) +client_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) +client_context.check_hostname = False +client_context.verify_mode = ssl.CERT_NONE +server_context.maximum_version = client_context.maximum_version = ssl.TLSVersion.TLSv1_2 + +client_in, client_out = ssl.MemoryBIO(), ssl.MemoryBIO() +server_in, server_out = ssl.MemoryBIO(), ssl.MemoryBIO() +client = client_context.wrap_bio(client_in, client_out) +server = server_context.wrap_bio(server_in, server_out, server_side=True) + +for _ in range(5): + try: + client.do_handshake() + except ssl.SSLWantReadError: + pass + server_in.write(client_out.read()) + try: + server.do_handshake() + except ssl.SSLWantReadError: + pass + client_in.write(server_out.read()) +client.do_handshake() +server.do_handshake() + +try: + server.unwrap() +except ssl.SSLWantReadError: + pass +client_in.write(server_out.read() + TAIL) +client.unwrap() +assert client_in.read() == TAIL