Fix SSL handshake partial send (#6635)

* simplify test_ftplib

* mark flaky test

* better ssl handshake, pending tls output

* flush_all_pending

* deadline

* flush pending
This commit is contained in:
Jeong, YunWon authored and GitHub committed 2026-01-04 17:53:59 +09:00
1 parent 8bfe927de5
commit fd27c2d78f
4 files changed
+513 -117

No files matched your search

+1 -1
View File
@@ -1459,7 +1459,7 @@ class _TestLock(BaseTestCase):
for _ in range(n):
lock.release()
@unittest.expectedFailureIf(sys.platform == "darwin", "TODO: RUSTPYTHON")
@unittest.skipIf(sys.platform == 'darwin', "TODO: RUSTPYTHON; flaky on darwin")
def test_repr_rlock(self):
if self.TYPE != 'processes':
self.skipTest('test not appropriate for {}'.format(self.TYPE))
+2 -22
View File
@@ -593,16 +593,12 @@ class TestFTPClass(TestCase):
def test_abort(self):
self.client.abort()
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_retrbinary(self):
received = []
self.client.retrbinary('retr', received.append)
self.check_data(b''.join(received),
RETR_DATA.encode(self.client.encoding))
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_retrbinary_rest(self):
for rest in (0, 10, 20):
received = []
@@ -610,14 +606,11 @@ class TestFTPClass(TestCase):
self.check_data(b''.join(received),
RETR_DATA[rest:].encode(self.client.encoding))
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_retrlines(self):
received = []
self.client.retrlines('retr', received.append)
self.check_data(''.join(received), RETR_DATA.replace('\r\n', ''))
@unittest.skip('TODO: RUSTPYTHON; weird limiting to 8192, something w/ buffering?')
def test_storbinary(self):
f = io.BytesIO(RETR_DATA.encode(self.client.encoding))
self.client.storbinary('stor', f)
@@ -629,8 +622,6 @@ class TestFTPClass(TestCase):
self.client.storbinary('stor', f, callback=lambda x: flag.append(None))
self.assertTrue(flag)
@unittest.skip('TODO: RUSTPYTHON')
# ssl_error.SSLWantReadError: The operation did not complete (read)
def test_storbinary_rest(self):
data = RETR_DATA.replace('\r\n', '\n').encode(self.client.encoding)
f = io.BytesIO(data)
@@ -639,8 +630,6 @@ class TestFTPClass(TestCase):
self.client.storbinary('stor', f, rest=r)
self.assertEqual(self.server.handler_instance.rest, str(r))
@unittest.skip('TODO: RUSTPYTHON')
# ssl_error.SSLWantReadError: The operation did not complete (read)
def test_storlines(self):
data = RETR_DATA.replace('\r\n', '\n').encode(self.client.encoding)
f = io.BytesIO(data)
@@ -658,21 +647,15 @@ class TestFTPClass(TestCase):
with warnings_helper.check_warnings(('', BytesWarning), quiet=True):
self.assertRaises(TypeError, self.client.storlines, 'stor foo', f)
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_nlst(self):
self.client.nlst()
self.assertEqual(self.client.nlst(), NLST_DATA.split('\r\n')[:-1])
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_dir(self):
l = []
self.client.dir(l.append)
self.assertEqual(''.join(l), LIST_DATA.replace('\r\n', ''))
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_mlsd(self):
list(self.client.mlsd())
list(self.client.mlsd(path='/'))
@@ -859,8 +842,6 @@ class TestFTPClass(TestCase):
f = io.BytesIO(b'x' * self.client.maxline * 2)
self.assertRaises(ftplib.Error, self.client.storlines, 'stor', f)
@unittest.skip('TODO: RUSTPYTHON')
# TimeoutError: The read operation timed out
def test_encoding_param(self):
encodings = ['latin-1', 'utf-8']
for encoding in encodings:
@@ -922,6 +903,7 @@ class TestIPv6Environment(TestCase):
retr()
@unittest.skip("TODO: RUSTPYTHON; SSL + asyncore has problem")
@skipUnless(ssl, "SSL not available")
@requires_subprocess()
class TestTLS_FTPClassMixin(TestFTPClass):
@@ -940,7 +922,7 @@ class TestTLS_FTPClassMixin(TestFTPClass):
@skipUnless(ssl, "SSL not available")
@unittest.skip("TODO: RUSTPYTHON; SSL + asyncore has problem on darwin")
@unittest.skip("TODO: RUSTPYTHON; SSL + asyncore has problem")
@requires_subprocess()
class TestTLS_FTPClass(TestCase):
"""Specific TLS_FTP class tests."""
@@ -1025,8 +1007,6 @@ class TestTLS_FTPClass(TestCase):
self.assertIs(sock.context, ctx)
self.assertIsInstance(sock, ssl.SSLSocket)
@unittest.skip('TODO: RUSTPYTHON')
# ssl_error.SSLWantReadError: The operation did not complete (read)
def test_ccc(self):
self.assertRaises(ValueError, self.client.ccc)
self.client.login(secure=True)
+360 -71
View File
@@ -1902,6 +1902,7 @@ mod _ssl {
pending_context: PyRwLock::new(None),
client_hello_buffer: PyMutex::new(None),
shutdown_state: PyMutex::new(ShutdownState::NotStarted),
pending_tls_output: PyMutex::new(Vec::new()),
deferred_cert_error: Arc::new(ParkingRwLock::new(None)),
};
@@ -1972,6 +1973,7 @@ mod _ssl {
pending_context: PyRwLock::new(None),
client_hello_buffer: PyMutex::new(None),
shutdown_state: PyMutex::new(ShutdownState::NotStarted),
pending_tls_output: PyMutex::new(Vec::new()),
deferred_cert_error: Arc::new(ParkingRwLock::new(None)),
};
@@ -2337,6 +2339,12 @@ mod _ssl {
// Shutdown state for tracking close-notify exchange
#[pytraverse(skip)]
shutdown_state: PyMutex<ShutdownState>,
// 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
// but the socket cannot accept all the data immediately
#[pytraverse(skip)]
pub(crate) pending_tls_output: PyMutex<Vec<u8>>,
// Deferred client certificate verification error (for TLS 1.3)
// Stores error message if client cert verification failed during handshake
// Error is raised on first I/O operation after handshake
@@ -2700,7 +2708,7 @@ mod _ssl {
// Helper to call socket methods, bypassing any SSL wrapper
pub(crate) fn sock_recv(&self, size: usize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
// In BIO mode, read from incoming BIO
// 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)?;
@@ -2711,21 +2719,19 @@ mod _ssl {
let socket_mod = vm.import("socket", 0)?;
let socket_class = socket_mod.get_attr("socket", vm)?;
// Call socket.socket.recv(self.sock, size)
// Call socket.socket.recv(self.sock, size, flags)
let recv_method = socket_class.get_attr("recv", vm)?;
recv_method.call((self.sock.clone(), vm.ctx.new_int(size)), vm)
}
pub(crate) fn sock_send(
&self,
data: Vec<u8>,
vm: &VirtualMachine,
) -> PyResult<PyObjectRef> {
/// 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<PyObjectRef> {
// 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),), vm);
return write_method.call((vm.ctx.new_bytes(data.to_vec()),), vm);
}
// Normal socket mode
@@ -2734,7 +2740,201 @@ mod _ssl {
// Call socket.socket.send(self.sock, data)
let send_method = socket_class.get_attr("send", vm)?;
send_method.call((self.sock.clone(), vm.ctx.new_bytes(data)), vm)
send_method.call((self.sock.clone(), vm.ctx.new_bytes(data.to_vec())), vm)
}
/// Flush any pending TLS output data to the socket
/// Optional deadline parameter allows respecting a read deadline during flush
pub(crate) fn flush_pending_tls_output(
&self,
vm: &VirtualMachine,
deadline: Option<std::time::Instant>,
) -> 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.map(|t| t.is_zero()).unwrap_or(false);
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)
} else {
socket_timeout
};
// Use sock_select directly with calculated timeout
let py_socket: PyRef<PySocket> = 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_select(&socket, SelectKind::Write, timeout_to_use)
.map_err(|e| vm.new_os_error(format!("select failed: {e}")))?;
if timed_out {
// Keep unsent data in pending buffer
*pending = pending[sent_total..].to_vec();
return Err(
timeout_error_msg(vm, "The write operation timed out".to_string()).upcast(),
);
}
match self.sock_send(&pending[sent_total..], vm) {
Ok(result) => {
let sent: usize = result.try_to_value::<isize>(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());
}
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);
}
}
}
// 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
fn send_tls_output(&self, buf: Vec<u8>, vm: &VirtualMachine) -> PyResult<()> {
if buf.is_empty() {
return Ok(());
}
let timeout = self.get_socket_timeout(vm)?;
let is_non_blocking = timeout.map(|t| t.is_zero()).unwrap_or(false);
let mut sent_total = 0;
while sent_total < buf.len() {
let timed_out = self.sock_wait_for_io_impl(SelectKind::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::<isize>(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(())
}
/// 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)?;
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<PySocket> = 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_select(&socket, SelectKind::Write, timeout)
.map_err(|e| vm.new_os_error(format!("select failed: {e}")))?;
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::<isize>(vm)?.try_into().unwrap_or(0);
if sent > 0 {
let mut pending = self.pending_tls_output.lock();
pending.drain(..sent);
}
// If sent == 0, socket wasn't ready despite select() saying so
// Continue loop to retry - this avoids infinite loops
}
Err(e) => {
if is_blocking_io_error(&e, vm) {
continue;
}
return Err(e);
}
}
}
}
#[pymethod]
@@ -3229,9 +3429,10 @@ mod _ssl {
};
}
// Ensure handshake is done
// Ensure handshake is done - if not, complete it first
// This matches OpenSSL behavior where SSL_read() auto-completes handshake
if !*self.handshake_done.lock() {
return Err(vm.new_value_error("Handshake not completed"));
self.do_handshake(vm)?;
}
// Check if connection has been shut down
@@ -3272,7 +3473,7 @@ mod _ssl {
};
// Use compat layer for unified read logic with proper EOF handling
// This matches CPython's SSL_read_ex() approach
// This matches SSL_read_ex() approach
let mut buf = vec![0u8; len];
let read_result = {
let mut conn_guard = self.connection.lock();
@@ -3360,9 +3561,10 @@ mod _ssl {
return Ok(0);
}
// Ensure handshake is done
// Ensure handshake is done - if not, complete it first
// This matches OpenSSL behavior where SSL_write() auto-completes handshake
if !*self.handshake_done.lock() {
return Err(vm.new_value_error("Handshake not completed"));
self.do_handshake(vm)?;
}
// Check if connection has been shut down
@@ -3387,6 +3589,13 @@ mod _ssl {
let is_bio = self.is_bio_mode();
let data: &[u8] = data_bytes.as_ref();
// CRITICAL: Flush any pending TLS data before writing new data
// This ensures TLS 1.3 Finished message reaches server before application data
// Without this, server may not be ready to process our data
if !is_bio {
self.flush_pending_tls_output(vm, None)?;
}
// Write data in chunks to avoid filling the internal TLS buffer
// rustls has a limited internal buffer, so we need to flush periodically
const CHUNK_SIZE: usize = 16384; // 16KB chunks (typical TLS record size)
@@ -3403,6 +3612,10 @@ mod _ssl {
writer
.write_all(chunk)
.map_err(|e| vm.new_os_error(format!("Write failed: {e}")))?;
// Flush to ensure data is converted to TLS records
writer
.flush()
.map_err(|e| vm.new_os_error(format!("Flush failed: {e}")))?;
}
written = chunk_end;
@@ -3413,6 +3626,9 @@ mod _ssl {
self.write_pending_tls(conn, vm)?;
} else {
// Socket mode: flush all pending TLS data
// First, try to send any previously pending data
self.flush_pending_tls_output(vm, None)?;
while conn.wants_write() {
let mut buf = Vec::new();
conn.write_tls(&mut buf).map_err(|e| {
@@ -3420,23 +3636,8 @@ mod _ssl {
})?;
if !buf.is_empty() {
let timed_out =
self.sock_wait_for_io_impl(SelectKind::Write, vm)?;
if timed_out {
return Err(vm.new_os_error("Write operation timed out"));
}
match self.sock_send(buf, vm) {
Ok(_) => {}
Err(e) => {
if is_blocking_io_error(&e, vm) {
return Err(
create_ssl_want_write_error(vm).upcast()
);
}
return Err(e);
}
}
// Try to send TLS data, saving unsent bytes to pending buffer
self.send_tls_output(buf, vm)?;
}
}
}
@@ -3755,8 +3956,35 @@ mod _ssl {
.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);
}
conn.send_close_notify();
// Write close_notify to outgoing buffer/BIO
@@ -3767,7 +3995,6 @@ mod _ssl {
}
// Step 2: Try to read and process peer's close_notify
let is_bio = self.is_bio_mode();
// First check if we already have peer's close_notify
// This can happen if it was received during a previous read() call
@@ -3775,46 +4002,84 @@ mod _ssl {
// If peer hasn't closed yet, try to read from socket
if !peer_closed {
// Check if socket is in blocking mode (timeout is None)
let is_blocking = if !is_bio {
// Check socket timeout mode
let timeout_mode = if !is_bio {
// Get socket timeout
match self.sock.get_attr("gettimeout", vm) {
Ok(method) => match method.call((), vm) {
Ok(timeout) => vm.is_none(&timeout),
Err(_) => false,
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(_) => false,
Err(_) => None,
}
} else {
false
None // BIO mode
};
if is_bio {
// In BIO mode: non-blocking read attempt
let _ = self.try_read_close_notify(conn, vm);
} else if is_blocking {
// Blocking socket mode: Return immediately without waiting for peer
if self.try_read_close_notify(conn, vm)? {
peer_closed = true;
}
} else if let Some(timeout) = timeout_mode {
// All socket modes (blocking, timeout, non-blocking):
// Return immediately after sending our close_notify.
//
// Reasons we don't read from socket here:
// 1. STARTTLS scenario: application data may arrive before/instead of close_notify
// - Example: client sends ENDTLS, immediately sends plain "msg 5"
// - Server's unwrap() would read "msg 5" and try to parse as TLS → FAIL
// 2. CPython's SSL_shutdown() typically returns immediately without waiting
// 3. Bidirectional shutdown is the application's responsibility
// 4. Reading from socket would consume application data incorrectly
// This matches CPython/OpenSSL behavior where SSL_shutdown()
// returns after sending close_notify, allowing the app to
// close the socket without waiting for peer's close_notify.
//
// Therefore: Just send our close_notify and return success immediately.
// The peer's close_notify (if any) will remain in the socket buffer.
//
// Mark shutdown as complete and return the underlying socket
// Waiting for peer's close_notify can cause deadlock with
// asyncore-based servers where both sides wait for the other's
// close_notify before closing the connection.
// Ensure all pending TLS data is sent before returning
// This prevents data loss when rustls drains its buffer
// but the socket couldn't accept all data immediately
drop(conn_guard);
// Respect socket timeout settings for flushing pending TLS data
match timeout {
Some(0.0) => {
// Non-blocking: best-effort flush, ignore errors
// to avoid deadlock with asyncore-based servers
let _ = self.flush_pending_tls_output(vm, None);
}
Some(_t) => {
// Timeout mode: use flush with socket's timeout
// Errors (including timeout) are propagated to caller
self.flush_pending_tls_output(vm, None)?;
}
None => {
// Blocking mode: wait until all pending data is sent
self.blocking_flush_all_pending(vm)?;
}
}
*self.shutdown_state.lock() = ShutdownState::Completed;
*self.connection.lock() = None;
return Ok(self.sock.clone());
}
// Step 3: Check again if peer has sent close_notify (non-blocking/BIO mode only)
peer_closed = self.check_peer_closed(conn, vm)?;
if !peer_closed {
peer_closed = self.check_peer_closed(conn, vm)?;
}
}
drop(conn_guard); // Release lock before returning
@@ -3836,6 +4101,10 @@ mod _ssl {
// Helper: Write all pending TLS data (including close_notify) to outgoing buffer/BIO
fn write_pending_tls(&self, conn: &mut TlsConnection, 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;
@@ -3850,42 +4119,62 @@ mod _ssl {
break;
}
// Send to outgoing BIO or socket
self.sock_send(buf[..written].to_vec(), vm)?;
// 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 BIO (non-blocking)
// 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 TlsConnection,
vm: &VirtualMachine,
) -> PyResult<()> {
// Try to read incoming data from BIO
// This is non-blocking in BIO mode - if no data, recv returns empty
) -> PyResult<bool> {
// Try to read incoming data
match self.sock_recv(SSL3_RT_MAX_PLAIN_LENGTH, vm) {
Ok(bytes_obj) => {
let bytes = ArgBytesLike::try_from_object(vm, bytes_obj)?;
let data = bytes.borrow_buf();
if !data.is_empty() {
// Feed data to TLS connection
let data_slice: &[u8] = data.as_ref();
let mut cursor = std::io::Cursor::new(data_slice);
let _ = conn.read_tls(&mut cursor);
// Process packets
let _ = conn.process_new_packets();
if data.is_empty() {
// Empty read could mean EOF or just "no data yet" in BIO mode
if let Some(ref bio) = self.incoming_bio {
// 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 {
// No EOF signaled, just no data available yet
return Ok(false);
}
}
// Socket mode or BIO with EOF: peer closed connection
// This is "ragged EOF" - peer closed without close_notify
return Ok(true);
}
// Feed data to TLS connection
let data_slice: &[u8] = data.as_ref();
let mut cursor = std::io::Cursor::new(data_slice);
let _ = conn.read_tls(&mut cursor);
// Process packets
let _ = conn.process_new_packets();
Ok(false)
}
Err(_) => {
// No data available or error - that's OK in BIO mode
Err(e) => {
// BlockingIOError means no data yet
if is_blocking_io_error(&e, vm) {
return Ok(false);
}
// Connection reset, EOF, or other error means peer closed
// ECONNRESET, EPIPE, broken pipe, etc.
Ok(true)
}
}
Ok(())
}
// Helper: Check if peer has sent close_notify
+150 -23
View File
@@ -1003,6 +1003,80 @@ pub(super) fn is_blocking_io_error(err: &Py<PyBaseException>, vm: &VirtualMachin
err.fast_isinstance(vm.ctx.exceptions.blocking_io_error)
}
// Socket I/O Helper Functions
/// Send all bytes to socket, handling partial sends with blocking wait
///
/// Loops until all bytes are sent. For blocking sockets, this will wait
/// 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(
socket: &PySSLSocket,
buf: Vec<u8>,
vm: &VirtualMachine,
deadline: Option<std::time::Instant>,
) -> SslResult<()> {
// First, flush any previously pending TLS data with deadline
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()));
}
match socket.sock_send(&buf[sent_total..], vm) {
Ok(result) => {
let sent: usize = result
.try_to_value::<isize>(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));
}
}
}
Ok(())
}
// Handshake Helper Functions
/// Write TLS handshake data to socket/BIO
@@ -1017,6 +1091,12 @@ fn handshake_write_loop(
) -> SslResult<bool> {
let mut made_progress = false;
// Flush any previously pending TLS data before generating new output
// Must succeed before sending new data to maintain order
socket
.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
@@ -1029,20 +1109,9 @@ fn handshake_write_loop(
.map_err(SslError::Io)?;
if written > 0 && !buf.is_empty() {
// Send directly without select - blocking sockets will wait automatically
// Handle BlockingIOError from non-blocking sockets
match socket.sock_send(buf, vm) {
Ok(_) => {
made_progress = true;
}
Err(e) => {
if is_blocking_io_error(&e, vm) {
// Non-blocking socket would block - return SSLWantWriteError
return Err(SslError::WantWrite);
}
return Err(SslError::Py(e));
}
}
// Send all bytes to socket, handling partial sends
send_all_bytes(socket, buf, vm, None)?;
made_progress = true;
} else if written == 0 {
// No data written but wants_write is true - should not happen normally
// Break to avoid infinite loop
@@ -1160,7 +1229,7 @@ fn handle_handshake_complete(
// Do NOT loop on wants_write() - avoid infinite loop/deadlock
let tls_data = ssl_write_tls_records(conn)?;
if !tls_data.is_empty() {
socket.sock_send(tls_data, vm).map_err(SslError::Py)?;
send_all_bytes(socket, tls_data, vm, None)?;
}
// IMPORTANT: Don't check wants_write() again!
@@ -1168,15 +1237,30 @@ fn handle_handshake_complete(
}
} else if conn.wants_write() {
// Send all pending data (e.g., TLS 1.3 NewSessionTicket) to socket
// Best-effort: WantWrite means socket buffer full, pending data will be
// sent in subsequent read/write calls. Don't fail handshake for this.
while conn.wants_write() {
let tls_data = ssl_write_tls_records(conn)?;
if tls_data.is_empty() {
break;
}
socket.sock_send(tls_data, vm).map_err(SslError::Py)?;
match send_all_bytes(socket, tls_data, vm, None) {
Ok(()) => {}
Err(SslError::WantWrite) => break,
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() {
socket
.blocking_flush_all_pending(vm)
.map_err(SslError::Py)?;
}
Ok(true)
}
@@ -1249,11 +1333,14 @@ pub(super) fn ssl_do_handshake(
// Send close_notify on error
if !is_bio {
conn.send_close_notify();
// Actually send the close_notify alert
// 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 _ = socket.sock_send(alert_data, vm);
let _ = send_all_bytes(socket, alert_data, vm, None);
}
}
@@ -1304,9 +1391,7 @@ pub(super) fn ssl_do_handshake(
break;
}
// Send to outgoing BIO
socket
.sock_send(buf[..n].to_vec(), vm)
.map_err(SslError::Py)?;
send_all_bytes(socket, buf[..n].to_vec(), vm, None)?;
// Check if there's more to write
if !conn.wants_write() {
break;
@@ -1388,6 +1473,16 @@ pub(super) fn ssl_read(
None // BIO mode has no deadline
};
// CRITICAL: Flush any pending TLS output before reading
// This ensures data from previous write() calls is sent before we wait for response.
// Without this, write() may leave data in pending_tls_output (if socket buffer was full),
// and read() would timeout waiting for a response that the server never received.
if !is_bio {
socket
.flush_pending_tls_output(vm, deadline)
.map_err(SslError::Py)?;
}
// Loop to handle TLS records and post-handshake messages
// Matches SSL_read behavior which loops until data is available
// - CPython uses OpenSSL's SSL_read which loops on SSL_ERROR_WANT_READ/WANT_WRITE
@@ -1396,6 +1491,7 @@ pub(super) fn ssl_read(
// - Blocking sockets: sock_select() and recv() wait at kernel level (no CPU busy-wait)
// - Non-blocking sockets: immediate return on first WantRead
// - Deadline prevents timeout issues
loop {
// Check deadline
if let Some(deadline) = deadline
@@ -1420,11 +1516,43 @@ pub(super) fn ssl_read(
// Check if connection needs to write data first (e.g., TLS key update, renegotiation)
// This mirrors the handshake logic which checks both wants_read() and wants_write()
if conn.wants_write() && !is_bio {
// Check deadline BEFORE attempting flush
if let Some(deadline) = deadline
&& std::time::Instant::now() >= deadline
{
return Err(SslError::Timeout(
"The read operation timed out".to_string(),
));
}
// Flush pending TLS data before continuing
// CRITICAL: Pass deadline so flush respects read timeout
let tls_data = ssl_write_tls_records(conn)?;
if !tls_data.is_empty() {
socket.sock_send(tls_data, vm).map_err(SslError::Py)?;
// Use best-effort send - don't fail READ just because WRITE couldn't complete
match send_all_bytes(socket, tls_data, vm, deadline) {
Ok(()) => {}
Err(SslError::WantWrite) => {
// Socket buffer full - acceptable during READ operation
// Pending data will be sent on next write/read call
}
Err(SslError::Timeout(_)) => {
// Timeout during flush is acceptable during READ
// Pending data stays buffered for next operation
}
Err(e) => return Err(e),
}
}
// Check deadline AFTER flush attempt
if let Some(deadline) = deadline
&& std::time::Instant::now() >= deadline
{
return Err(SslError::Timeout(
"The read operation timed out".to_string(),
));
}
// After flushing, rustls may want to read again - continue loop
continue;
}
@@ -1487,7 +1615,6 @@ pub(super) fn ssl_read(
}
// Read and process TLS records
// This will block for blocking sockets
match ssl_ensure_data_available(conn, socket, vm) {
Ok(_bytes_read) => {
// Successfully read and processed TLS data