Skip to content

Commit aaa6fd6

Browse files
youknowoneyouknowone
committed
fix ssl
1 parent 5a07e8e commit aaa6fd6

2 files changed

Lines changed: 219 additions & 80 deletions

File tree

‎crates/stdlib/src/ssl.rs‎

Lines changed: 166 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -2700,7 +2700,7 @@ mod _ssl {
27002700

27012701
// Helper to call socket methods, bypassing any SSL wrapper
27022702
pub(crate) fn sock_recv(&self, size: usize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
2703-
// In BIO mode, read from incoming BIO
2703+
// In BIO mode, read from incoming BIO (flags not supported)
27042704
if let Some(ref bio) = self.incoming_bio {
27052705
let bio_obj: PyObjectRef = bio.clone().into();
27062706
let read_method = bio_obj.get_attr("read", vm)?;
@@ -2711,7 +2711,7 @@ mod _ssl {
27112711
let socket_mod = vm.import("socket", 0)?;
27122712
let socket_class = socket_mod.get_attr("socket", vm)?;
27132713

2714-
// Call socket.socket.recv(self.sock, size)
2714+
// Call socket.socket.recv(self.sock, size, flags)
27152715
let recv_method = socket_class.get_attr("recv", vm)?;
27162716
recv_method.call((self.sock.clone(), vm.ctx.new_int(size)), vm)
27172717
}
@@ -3229,9 +3229,10 @@ mod _ssl {
32293229
};
32303230
}
32313231

3232-
// Ensure handshake is done
3232+
// Ensure handshake is done - if not, complete it first
3233+
// This matches OpenSSL behavior where SSL_read() auto-completes handshake
32333234
if !*self.handshake_done.lock() {
3234-
return Err(vm.new_value_error("Handshake not completed"));
3235+
self.do_handshake(vm)?;
32353236
}
32363237

32373238
// Check if connection has been shut down
@@ -3272,7 +3273,7 @@ mod _ssl {
32723273
};
32733274

32743275
// Use compat layer for unified read logic with proper EOF handling
3275-
// This matches CPython's SSL_read_ex() approach
3276+
// This matches SSL_read_ex() approach
32763277
let mut buf = vec![0u8; len];
32773278
let read_result = {
32783279
let mut conn_guard = self.connection.lock();
@@ -3360,9 +3361,10 @@ mod _ssl {
33603361
return Ok(0);
33613362
}
33623363

3363-
// Ensure handshake is done
3364+
// Ensure handshake is done - if not, complete it first
3365+
// This matches OpenSSL behavior where SSL_write() auto-completes handshake
33643366
if !*self.handshake_done.lock() {
3365-
return Err(vm.new_value_error("Handshake not completed"));
3367+
self.do_handshake(vm)?;
33663368
}
33673369

33683370
// Check if connection has been shut down
@@ -3420,21 +3422,57 @@ mod _ssl {
34203422
})?;
34213423

34223424
if !buf.is_empty() {
3423-
let timed_out =
3424-
self.sock_wait_for_io_impl(SelectKind::Write, vm)?;
3425-
if timed_out {
3426-
return Err(vm.new_os_error("Write operation timed out"));
3427-
}
3425+
// Check if this is a non-blocking socket
3426+
let timeout = self.get_socket_timeout(vm)?;
3427+
let is_non_blocking =
3428+
timeout.map(|t| t.is_zero()).unwrap_or(false);
3429+
3430+
// Send all bytes, handling partial sends
3431+
let mut sent_total = 0;
3432+
while sent_total < buf.len() {
3433+
let timed_out =
3434+
self.sock_wait_for_io_impl(SelectKind::Write, vm)?;
3435+
if timed_out {
3436+
return Err(
3437+
vm.new_os_error("Write operation timed out")
3438+
);
3439+
}
34283440

3429-
match self.sock_send(buf, vm) {
3430-
Ok(_) => {}
3431-
Err(e) => {
3432-
if is_blocking_io_error(&e, vm) {
3433-
return Err(
3434-
create_ssl_want_write_error(vm).upcast()
3435-
);
3441+
let to_send = buf[sent_total..].to_vec();
3442+
match self.sock_send(to_send, vm) {
3443+
Ok(result) => {
3444+
let sent: usize = result
3445+
.try_to_value::<isize>(vm)?
3446+
.try_into()
3447+
.unwrap_or(0);
3448+
if sent == 0 {
3449+
// No progress - socket buffer full
3450+
if is_non_blocking {
3451+
// Non-blocking: must return to avoid infinite loop
3452+
return Err(create_ssl_want_write_error(
3453+
vm,
3454+
)
3455+
.upcast());
3456+
}
3457+
// Blocking: retry
3458+
continue;
3459+
}
3460+
sent_total += sent;
3461+
}
3462+
Err(e) => {
3463+
if is_blocking_io_error(&e, vm) {
3464+
if is_non_blocking {
3465+
// Non-blocking: must return to avoid infinite loop
3466+
return Err(create_ssl_want_write_error(
3467+
vm,
3468+
)
3469+
.upcast());
3470+
}
3471+
// Blocking: retry
3472+
continue;
3473+
}
3474+
return Err(e);
34363475
}
3437-
return Err(e);
34383476
}
34393477
}
34403478
}
@@ -3775,46 +3813,61 @@ mod _ssl {
37753813

37763814
// If peer hasn't closed yet, try to read from socket
37773815
if !peer_closed {
3778-
// Check if socket is in blocking mode (timeout is None)
3779-
let is_blocking = if !is_bio {
3816+
// Check socket timeout mode
3817+
let timeout_mode = if !is_bio {
37803818
// Get socket timeout
37813819
match self.sock.get_attr("gettimeout", vm) {
37823820
Ok(method) => match method.call((), vm) {
3783-
Ok(timeout) => vm.is_none(&timeout),
3784-
Err(_) => false,
3821+
Ok(timeout) => {
3822+
if vm.is_none(&timeout) {
3823+
// timeout=None means blocking
3824+
Some(None)
3825+
} else if let Ok(t) = timeout.try_float(vm).map(|f| f.to_f64()) {
3826+
if t == 0.0 {
3827+
// timeout=0 means non-blocking
3828+
Some(Some(0.0))
3829+
} else {
3830+
// timeout>0 means timeout mode
3831+
Some(Some(t))
3832+
}
3833+
} else {
3834+
None
3835+
}
3836+
}
3837+
Err(_) => None,
37853838
},
3786-
Err(_) => false,
3839+
Err(_) => None,
37873840
}
37883841
} else {
3789-
false
3842+
None // BIO mode
37903843
};
37913844

37923845
if is_bio {
37933846
// In BIO mode: non-blocking read attempt
3794-
let _ = self.try_read_close_notify(conn, vm);
3795-
} else if is_blocking {
3796-
// Blocking socket mode: Return immediately without waiting for peer
3797-
//
3798-
// Reasons we don't read from socket here:
3799-
// 1. STARTTLS scenario: application data may arrive before/instead of close_notify
3800-
// - Example: client sends ENDTLS, immediately sends plain "msg 5"
3801-
// - Server's unwrap() would read "msg 5" and try to parse as TLS → FAIL
3802-
// 2. CPython's SSL_shutdown() typically returns immediately without waiting
3803-
// 3. Bidirectional shutdown is the application's responsibility
3804-
// 4. Reading from socket would consume application data incorrectly
3847+
if self.try_read_close_notify(conn, vm)? {
3848+
peer_closed = true;
3849+
}
3850+
} else if let Some(_timeout) = timeout_mode {
3851+
// All socket modes (blocking, timeout, non-blocking):
3852+
// Return immediately after sending our close_notify.
38053853
//
3806-
// Therefore: Just send our close_notify and return success immediately.
3807-
// The peer's close_notify (if any) will remain in the socket buffer.
3854+
// This matches CPython/OpenSSL behavior where SSL_shutdown()
3855+
// returns after sending close_notify, allowing the app to
3856+
// close the socket without waiting for peer's close_notify.
38083857
//
3809-
// Mark shutdown as complete and return the underlying socket
3858+
// Waiting for peer's close_notify can cause deadlock with
3859+
// asyncore-based servers where both sides wait for the other's
3860+
// close_notify before closing the connection.
38103861
drop(conn_guard);
38113862
*self.shutdown_state.lock() = ShutdownState::Completed;
38123863
*self.connection.lock() = None;
38133864
return Ok(self.sock.clone());
38143865
}
38153866

38163867
// Step 3: Check again if peer has sent close_notify (non-blocking/BIO mode only)
3817-
peer_closed = self.check_peer_closed(conn, vm)?;
3868+
if !peer_closed {
3869+
peer_closed = self.check_peer_closed(conn, vm)?;
3870+
}
38183871
}
38193872

38203873
drop(conn_guard); // Release lock before returning
@@ -3850,42 +3903,96 @@ mod _ssl {
38503903
break;
38513904
}
38523905

3853-
// Send to outgoing BIO or socket
3854-
self.sock_send(buf[..written].to_vec(), vm)?;
3906+
// Check if this is a non-blocking socket
3907+
let timeout = self.get_socket_timeout(vm)?;
3908+
let is_non_blocking = timeout.map(|t| t.is_zero()).unwrap_or(false);
3909+
3910+
// Send to outgoing BIO or socket, handling partial sends
3911+
let data = buf[..written].to_vec();
3912+
let mut sent_total = 0;
3913+
while sent_total < data.len() {
3914+
let to_send = data[sent_total..].to_vec();
3915+
match self.sock_send(to_send, vm) {
3916+
Ok(result) => {
3917+
let sent: usize =
3918+
result.try_to_value::<isize>(vm)?.try_into().unwrap_or(0);
3919+
if sent == 0 {
3920+
// Would block
3921+
if is_non_blocking {
3922+
// Non-blocking socket: return WantWrite to avoid infinite loop
3923+
return Err(create_ssl_want_write_error(vm).upcast());
3924+
}
3925+
// Wait for socket to be ready
3926+
let timed_out =
3927+
self.sock_wait_for_io_impl(SelectKind::Write, vm)?;
3928+
if timed_out {
3929+
return Err(vm.new_os_error("Write operation timed out"));
3930+
}
3931+
continue;
3932+
}
3933+
sent_total += sent;
3934+
}
3935+
Err(e) => {
3936+
if is_blocking_io_error(&e, vm) {
3937+
if is_non_blocking {
3938+
// Non-blocking socket: return WantWrite
3939+
return Err(create_ssl_want_write_error(vm).upcast());
3940+
}
3941+
// Wait for socket to be ready
3942+
let timed_out =
3943+
self.sock_wait_for_io_impl(SelectKind::Write, vm)?;
3944+
if timed_out {
3945+
return Err(vm.new_os_error("Write operation timed out"));
3946+
}
3947+
continue;
3948+
}
3949+
return Err(e);
3950+
}
3951+
}
3952+
}
38553953
}
38563954

38573955
Ok(())
38583956
}
38593957

3860-
// Helper: Try to read incoming data from BIO (non-blocking)
3958+
// Helper: Try to read incoming data from socket/BIO
3959+
// Returns true if peer closed connection (with or without close_notify)
38613960
fn try_read_close_notify(
38623961
&self,
38633962
conn: &mut TlsConnection,
38643963
vm: &VirtualMachine,
3865-
) -> PyResult<()> {
3866-
// Try to read incoming data from BIO
3867-
// This is non-blocking in BIO mode - if no data, recv returns empty
3964+
) -> PyResult<bool> {
3965+
// Try to read incoming data
38683966
match self.sock_recv(SSL3_RT_MAX_PLAIN_LENGTH, vm) {
38693967
Ok(bytes_obj) => {
38703968
let bytes = ArgBytesLike::try_from_object(vm, bytes_obj)?;
38713969
let data = bytes.borrow_buf();
38723970

3873-
if !data.is_empty() {
3874-
// Feed data to TLS connection
3875-
let data_slice: &[u8] = data.as_ref();
3876-
let mut cursor = std::io::Cursor::new(data_slice);
3877-
let _ = conn.read_tls(&mut cursor);
3878-
3879-
// Process packets
3880-
let _ = conn.process_new_packets();
3971+
if data.is_empty() {
3972+
// Empty read = EOF = peer closed connection
3973+
// This is "ragged EOF" - peer closed without close_notify
3974+
return Ok(true);
38813975
}
3976+
3977+
// Feed data to TLS connection
3978+
let data_slice: &[u8] = data.as_ref();
3979+
let mut cursor = std::io::Cursor::new(data_slice);
3980+
let _ = conn.read_tls(&mut cursor);
3981+
3982+
// Process packets
3983+
let _ = conn.process_new_packets();
3984+
Ok(false)
38823985
}
3883-
Err(_) => {
3884-
// No data available or error - that's OK in BIO mode
3986+
Err(e) => {
3987+
// BlockingIOError means no data yet
3988+
if is_blocking_io_error(&e, vm) {
3989+
return Ok(false);
3990+
}
3991+
// Connection reset, EOF, or other error means peer closed
3992+
// ECONNRESET, EPIPE, broken pipe, etc.
3993+
Ok(true)
38853994
}
38863995
}
3887-
3888-
Ok(())
38893996
}
38903997

38913998
// Helper: Check if peer has sent close_notify

0 commit comments

Comments
 (0)