mirror of
https://github.com/RustPython/RustPython.git
synced 2026-10-29 03:04:38 +08:00
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:
1 parent
8bfe927de5
commit
fd27c2d78f
4 files changed
+513
-117
No files matched your search
Vendored
+1
-1
@@ -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))
|
||||
|
||||
Vendored
+2
-22
@@ -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
@@ -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
@@ -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
|
||||
|
||||
Reference in new issue
Block a user