diff --git a/src/socketengine.rs b/src/socketengine.rs index 22440fa..9007cb0 100644 --- a/src/socketengine.rs +++ b/src/socketengine.rs @@ -204,6 +204,21 @@ impl Sock { Sock::Tls(t) => t.write(buf), } } + /// Whether a TLS session still has outbound bytes buffered internally (rustls + /// after a WouldBlock); a plaintext socket never buffers in-process. + fn wants_write(&self) -> bool { + match self { + Sock::Plain(_) => false, + Sock::Tls(t) => t.wants_write(), + } + } + /// Push a TLS session's buffered ciphertext to the socket; no-op for plaintext. + fn flush(&mut self) -> io::Result<()> { + match self { + Sock::Plain(_) => Ok(()), + Sock::Tls(t) => t.flush(), + } + } /// Best-effort graceful close. TLS sends a close_notify alert; plaintext relies on /// the socket's own FIN when the stream drops. fn shutdown(&mut self) { @@ -245,8 +260,9 @@ impl Conn { fn set_interest(poll: &mut Poll, c: &mut Conn, t: usize) { let want_read = !c.paused; // a TLS handshake may need to write (its flight) as well as read, so keep both - // until it completes; after that, write only when there's a backlog to drain. - let want_write = c.handshaking || !c.wbuf.is_empty() || c.paused; + // until it completes; after that, write when there's a backlog to drain — either + // our own queued plaintext, or ciphertext still buffered inside a TLS session. + let want_write = c.handshaking || !c.wbuf.is_empty() || c.paused || c.sock.wants_write(); if want_read == c.want_read && want_write == c.want_write { return; } @@ -979,12 +995,20 @@ fn flush_conn(poll: &mut Poll, conns: &mut HashMap, t: usize, core: c.wbuf.clear(); c.wpos = 0; } + // push any ciphertext a TLS session still holds buffered (rustls keeps it when + // the socket filled mid-write); our plaintext queue draining doesn't mean the + // socket has it all. WouldBlock leaves the rest for the next writable event. + match c.sock.flush() { + Ok(()) => {} + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} + Err(_) => close = true, + } if c.paused && c.pending() <= c.softsendq { c.paused = false; unpaused = true; } set_interest(poll, c, t); - if c.closing && c.wbuf.is_empty() { + if c.closing && c.wbuf.is_empty() && !c.sock.wants_write() { close = true; } } diff --git a/src/tls.rs b/src/tls.rs index 75cfacf..7c0ae1f 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -55,6 +55,17 @@ pub trait TlsSession: Send { fn read(&mut self, buf: &mut [u8]) -> io::Result; /// Encrypt+queue application data; returns the plaintext bytes accepted. fn write(&mut self, buf: &[u8]) -> io::Result; + /// Whether the session still holds outbound TLS bytes not yet pushed to the + /// socket. rustls buffers ciphertext internally when the socket is full; + /// openssl surfaces backpressure through `write`, so it never buffers. + fn wants_write(&self) -> bool { + false + } + /// Push any buffered outbound TLS bytes to the socket. `WouldBlock` leaves the + /// remainder for the next writable event; a no-op when nothing is buffered. + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } /// The underlying mio socket, for the reactor's poll (re)registration. fn source(&mut self) -> &mut MioStream; /// SHA-256 fingerprint of the peer certificate (CertFP / SASL EXTERNAL), if any. diff --git a/src/tls_rustls.rs b/src/tls_rustls.rs index 836bc12..5698924 100644 --- a/src/tls_rustls.rs +++ b/src/tls_rustls.rs @@ -219,7 +219,11 @@ impl TlsBackend for RustlsBackend { } fn start(&self, sock: MioStream) -> io::Result> { - let conn = ServerConnection::new(self.cfg()).map_err(err)?; + let mut conn = ServerConnection::new(self.cfg()).map_err(err)?; + // bound the buffered plaintext so a slow-reading client makes writer().write() + // return short (backpressure) instead of growing without limit; the reactor's + // sendq caps then govern it, matching the openssl backend. + conn.set_buffer_limit(Some(256 * 1024)); Ok(Box::new(RustlsSession { conn, sock })) } } @@ -307,6 +311,14 @@ impl TlsSession for RustlsSession { self.pump_write()?; Ok(n) } + fn wants_write(&self) -> bool { + // rustls holds encrypted bytes when the socket filled mid-flush; the reactor + // must keep WRITABLE interest and drain them, or a burst strands here. + self.conn.wants_write() + } + fn flush(&mut self) -> io::Result<()> { + self.pump_write() + } fn source(&mut self) -> &mut MioStream { &mut self.sock }