rust: bssl-tls: Application facing I/O

Bug: 479599893

Signed-off-by: Xiangfei Ding <xfding@google.com>
Change-Id: I71dcfa3b634abce34b1fb9904517846b6a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/91467
Reviewed-by: Adam Langley <agl@google.com>
Presubmit-BoringSSL-Verified: boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com <boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/rust/bssl-tls/src/config.rs b/rust/bssl-tls/src/config.rs
index 4de75ce..baee676 100644
--- a/rust/bssl-tls/src/config.rs
+++ b/rust/bssl-tls/src/config.rs
@@ -54,6 +54,9 @@
     pub(crate) struct ConnectionMode: u32 {
         /// Deny session creation.
         const MODE_NO_SESSION_CREATION = bssl_sys::SSL_MODE_NO_SESSION_CREATION as u32;
+        /// Allow moving write buffer.
+        /// This is indispensable for async I/O because the future could be freely cancelled.
+        const ACCEPT_MOVING_WRITE_BUFFER = bssl_sys::SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER as u32;
     }
 }
 
diff --git a/rust/bssl-tls/src/connection.rs b/rust/bssl-tls/src/connection.rs
index 94a1aeb..ded375d 100644
--- a/rust/bssl-tls/src/connection.rs
+++ b/rust/bssl-tls/src/connection.rs
@@ -45,6 +45,7 @@
 };
 
 mod credentials;
+pub mod io;
 pub mod lifecycle;
 pub(crate) mod methods;
 pub mod transport;
@@ -102,9 +103,15 @@
         let data = Box::into_raw(Box::new(methods::RustConnectionMethods::<M>::new())) as _;
         unsafe {
             // Safety:
+            // - the validity of the handle `ptr` is witnessed by `self`.
             // - `M::registration` will return a valid ex-data index.
             // - `data` should be valid by non-null invariant.
             bssl_sys::SSL_set_ex_data(ptr.as_ptr(), idx, data);
+            // Safety: the validity of the handle `ptr` is witnessed by `self`.
+            bssl_sys::SSL_set_mode(
+                ptr.as_ptr(),
+                ConnectionMode::ACCEPT_MOVING_WRITE_BUFFER.bits(),
+            );
         }
         Self {
             ptr,
@@ -178,7 +185,7 @@
 unsafe impl<R, M> Send for TlsConnectionRef<R, M> {}
 
 impl<R, M> TlsConnectionRef<R, M> {
-    #[allow(unused)]
+    /// Call this method whenever I/O is performed on the connection.
     pub(crate) fn categorise_error_for_io(&self, rc: c_int) -> Result<IoStatus, Error> {
         let reason = unsafe {
             // Safety: we only want to extract the last I/O error on an existing valid connection.
diff --git a/rust/bssl-tls/src/connection/io.rs b/rust/bssl-tls/src/connection/io.rs
new file mode 100644
index 0000000..80b61aa
--- /dev/null
+++ b/rust/bssl-tls/src/connection/io.rs
@@ -0,0 +1,331 @@
+// Copyright 2026 The BoringSSL Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! TLS I/O model
+
+use alloc::boxed::Box;
+use core::{
+    ffi::c_int,
+    future::poll_fn,
+    pin::Pin,
+    task::{
+        Context,
+        Poll, //
+    }, //
+};
+
+use crate::{
+    ReceiveBuffer,
+    connection::{
+        lifecycle::ShutdownStatus,
+        methods::HasTlsConnectionMethod, //
+    },
+    context::{
+        HasBasicIo,
+        TlsMode, //
+    },
+    errors::{
+        Error,
+        IoError,
+        TlsRetryReason, //
+    },
+    ffi::slice_into_ffi_raw_parts,
+    io::IoStatus, //
+};
+
+use super::TlsConnectionRef;
+
+impl<R, M> TlsConnectionRef<R, M>
+where
+    M: HasTlsConnectionMethod,
+{
+    /// Check if the connection has any buffered data pending reads.
+    pub fn has_pending_read(&self) -> bool {
+        unsafe {
+            // Safety: the validity of the handle `self.ptr()` is witnessed by `self`.
+            bssl_sys::SSL_has_pending(self.ptr()) == 1
+        }
+    }
+
+    fn take_io_err(&mut self) -> Option<Box<dyn core::error::Error + Send + Sync>> {
+        let bio = self.get_connection_methods().bio.as_mut()?;
+        bio.as_mut().take_io_err()
+    }
+
+    /// Translate I/O error into the right form.
+    ///
+    /// It is here we translate retry reason into a **soft** error [`IoStatus::Retry`].
+    fn translate_io_error(&mut self, rc: c_int) -> Result<IoStatus, Error> {
+        // Pre-emptively extract error and clear the error queue.
+        let ssl_err = self.categorise_error_for_io(rc);
+        if let Some(err) = self.take_io_err() {
+            Err(Error::Io(IoError::Transport(err)))
+        } else {
+            ssl_err
+        }
+    }
+
+    /// Read data from the socket.
+    ///
+    /// This method reads up to `buffer.len()` bytes from `buffer`.
+    pub fn sync_read(&mut self, buffer: &mut ReceiveBuffer<'_>) -> Result<IoStatus, Error> {
+        let buf = unsafe {
+            // Safety:
+            // - the use of this pointer is outlived by this function callframe.
+            // - the access to the buffer region is bounded by `buffer.remaining()` by `SSL_read`
+            //   contract.
+            // - there are no reads into the buffer region per `SSL_read` contract.
+            buffer.head()
+        };
+        let num = c_int::try_from(buffer.remaining()).unwrap_or(c_int::MAX);
+        let rc = unsafe {
+            // Safety: the validity of the handle `self.ptr()` is witnessed by `self`.
+            bssl_sys::SSL_read(self.ptr(), buf as _, num)
+        };
+        if rc > 0 {
+            let len = rc as usize;
+            unsafe {
+                // Safety: BoringSSL will ensure that `len` bytes have been written.
+                buffer.advance(len);
+            }
+            Ok(IoStatus::Ok(len))
+        } else {
+            self.translate_io_error(rc)
+        }
+    }
+
+    /// Peek `buffer.len()` bytes of application data into the `buffer`.
+    pub fn peek(&mut self, buffer: &mut ReceiveBuffer<'_>) -> Result<IoStatus, Error> {
+        let buf = unsafe {
+            // Safety:
+            // - the use of this pointer is outlived by this function callframe.
+            // - the access to the buffer region is bounded by `buffer.remaining()` by `SSL_peek`
+            //   contract.
+            // - there are no reads into the buffer region per `SSL_peek` contract.
+            buffer.head()
+        };
+        let num = c_int::try_from(buffer.remaining()).unwrap_or(c_int::MAX);
+        let rc = unsafe {
+            // Safety: the validity of the handle `self.ptr()` is witnessed by `self`
+            bssl_sys::SSL_peek(self.ptr(), buf as _, num)
+        };
+        if rc > 0 {
+            let len = rc as usize;
+            unsafe {
+                // Safety: BoringSSL will ensure that `len` bytes have been written.
+                buffer.advance(len);
+            }
+            Ok(IoStatus::Ok(len))
+        } else {
+            self.translate_io_error(rc)
+        }
+    }
+
+    /// Write data to the socket.
+    ///
+    /// This method writes up to `buffer.len()` bytes from `buffer`.
+    pub fn sync_write(&mut self, buffer: &[u8]) -> Result<IoStatus, Error> {
+        let (ptr, len) = slice_into_ffi_raw_parts(buffer);
+        let num = c_int::try_from(len).unwrap_or(c_int::MAX);
+        let rc = unsafe {
+            // Safety: the validity of the handle `self.ptr()` is witnessed by `self`
+            bssl_sys::SSL_write(self.ptr(), ptr as _, num)
+        };
+        if rc > 0 {
+            Ok(IoStatus::Ok(rc as usize))
+        } else {
+            self.translate_io_error(rc)
+        }
+    }
+
+    /// Flush the data on the **transport**.
+    ///
+    /// On success, this method always reports the number of bytes moved as `0`.
+    pub fn flush(&mut self) -> Result<IoStatus, Error> {
+        let bio = unsafe {
+            // Safety: the validity of the handle `self.ptr()` is witnessed by `self`.
+            bssl_sys::SSL_get_wbio(self.ptr())
+        };
+        if bio.is_null() {
+            return Ok(IoStatus::Empty);
+        }
+        let rc = unsafe {
+            // Safety: `bio` should still be valid by BoringSSL invariant.
+            bssl_sys::BIO_flush(bio)
+        };
+        if rc == 1 {
+            return Ok(IoStatus::Ok(0));
+        }
+        // We do not expect any SSL level error, but there could still be BIO level error.
+        let bio_retry = unsafe {
+            // Safety: `bio` should still be valid here.
+            bssl_sys::BIO_should_retry(bio)
+        };
+        if bio_retry != 1 {
+            return Ok(IoStatus::Retry(TlsRetryReason::WantWrite));
+        }
+        // Pre-emptively extract error and clear the error queue.
+        if let Some(err) = self.take_io_err() {
+            Err(Error::Io(IoError::Transport(err)))
+        } else {
+            Ok(IoStatus::Ok(0))
+        }
+    }
+}
+
+/// Async I/O
+impl<R, M> TlsConnectionRef<R, M>
+where
+    M: HasTlsConnectionMethod,
+{
+    /// For `async` operations, obtain a pinned mutable reference.
+    pub fn as_pin_mut(&mut self) -> Pin<&mut Self> {
+        Pin::new(self)
+    }
+
+    /// For `async` operations, obtain a pinned immutable reference.
+    pub fn as_pin(&self) -> Pin<&Self> {
+        Pin::new(self)
+    }
+}
+
+impl<R, M> TlsConnectionRef<R, M>
+where
+    M: HasTlsConnectionMethod + HasBasicIo,
+{
+    fn do_async_io(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        sync_op: impl FnOnce(&mut TlsConnectionRef<R, M>) -> Result<IoStatus, Error>,
+    ) -> Result<Option<IoStatus>, Error> {
+        self.set_waker(cx.waker());
+
+        let reason = match sync_op(&mut *self) {
+            Ok(
+                status @ (IoStatus::Ok(..)
+                | IoStatus::EndOfStream
+                | IoStatus::Empty
+                | IoStatus::Err),
+            ) => return Ok(Some(status)),
+            Err(e) => return Err(e),
+            Ok(IoStatus::Retry(reason)) => reason,
+        };
+        self.get_connection_methods().set_pending_reason(reason);
+        Ok(None)
+    }
+    fn aread_inner(
+        self: Pin<&mut Self>,
+        buffer: &mut [u8],
+        cx: &mut Context<'_>,
+    ) -> Result<Option<IoStatus>, Error> {
+        let mut buffer = ReceiveBuffer::new(buffer);
+        self.do_async_io(cx, move |this| this.sync_read(&mut buffer))
+    }
+
+    fn awrite_inner(
+        self: Pin<&mut Self>,
+        buffer: &[u8],
+        cx: &mut Context<'_>,
+    ) -> Result<Option<IoStatus>, Error> {
+        self.do_async_io(cx, move |this| this.sync_write(buffer))
+    }
+
+    fn aflush_inner(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Result<Option<IoStatus>, Error> {
+        self.do_async_io(cx, move |this| this.flush())
+    }
+
+    /// Asynchronously read application data from the TLS connection.
+    ///
+    /// This method will intercept [`IoStatus::Retry`] and suspend the future.
+    /// The reason can be inspected by invoking [`Self::take_pending_reason`].
+    pub fn async_read<'a>(
+        mut self: Pin<&'a mut Self>,
+        buffer: &'a mut [u8],
+    ) -> impl 'a + Send + Future<Output = Result<IoStatus, Error>> {
+        poll_fn(move |cx| match self.as_mut().aread_inner(buffer, cx) {
+            Ok(Some(status)) => Poll::Ready(Ok(status)),
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(e)),
+        })
+    }
+
+    /// Asynchronously write application data to the TLS connection.
+    ///
+    /// This method will intercept [`IoStatus::Retry`] and suspend the future.
+    /// The reason can be inspected by invoking [`Self::take_pending_reason`].
+    pub fn async_write<'a>(
+        mut self: Pin<&'a mut Self>,
+        buffer: &'a [u8],
+    ) -> impl 'a + Send + Future<Output = Result<IoStatus, Error>> {
+        poll_fn(move |cx| match self.as_mut().awrite_inner(buffer, cx) {
+            Ok(Some(status)) => Poll::Ready(Ok(status)),
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(e)),
+        })
+    }
+
+    /// Asynchronously flush the underlying transport attached to the TLS connection.
+    ///
+    /// This method will intercept [`IoStatus::Retry`] and suspend the future.
+    /// The reason can be inspected by invoking [`Self::take_pending_reason`].
+    pub fn async_flush<'a>(
+        mut self: Pin<&'a mut Self>,
+    ) -> impl 'a + Send + Future<Output = Result<IoStatus, Error>> {
+        poll_fn(move |cx| match self.as_mut().aflush_inner(cx) {
+            Ok(Some(status)) => Poll::Ready(Ok(status)),
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(e)),
+        })
+    }
+
+    fn ashutdown_inner(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Result<Option<ShutdownStatus>, Error> {
+        self.set_waker(cx.waker());
+        let Some(mut conn) = self.established() else {
+            return Err(Error::Io(IoError::EndOfStream));
+        };
+        loop {
+            match conn.sync_shutdown()? {
+                Some(ShutdownStatus::CloseNotifyPosted) => {}
+                status => return Ok(status),
+            }
+        }
+    }
+
+    /// Asynchronously shut down the connection.
+    pub fn async_shutdown<'a>(
+        mut self: Pin<&'a mut Self>,
+    ) -> impl 'a + Send + Future<Output = Result<(), Error>> {
+        poll_fn(move |cx| match self.as_mut().ashutdown_inner(cx) {
+            Ok(Some(ShutdownStatus::CloseNotifyReceived)) => Poll::Ready(Ok(())),
+            Ok(Some(ShutdownStatus::EndOfStream)) => {
+                Poll::Ready(Err(Error::Io(IoError::EndOfStream)))
+            }
+            Ok(Some(ShutdownStatus::RemainingApplicationData)) => Poll::Ready(Err(
+                Error::TlsReason(crate::errors::TlsErrorReason::ApplicationDataOnShutdown),
+            )),
+            Ok(Some(ShutdownStatus::CloseNotifyPosted)) => unreachable!(),
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(e)),
+        })
+    }
+}
+
+#[cfg(feature = "std")]
+mod stdio;
+
+#[cfg(feature = "tokio_io")]
+mod tokio_io_impl;
diff --git a/rust/bssl-tls/src/connection/io/stdio.rs b/rust/bssl-tls/src/connection/io/stdio.rs
new file mode 100644
index 0000000..2a67962
--- /dev/null
+++ b/rust/bssl-tls/src/connection/io/stdio.rs
@@ -0,0 +1,107 @@
+// Copyright 2026 The BoringSSL Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::io;
+
+use super::{
+    Error,
+    TlsConnectionRef,
+    TlsMode, //
+};
+use crate::{
+    ReceiveBuffer,
+    context::DtlsMode,
+    errors::{
+        IoError,
+        TlsRetryReason, //
+    },
+    io::{
+        AbstractSocketResult,
+        IoStatus,
+        stdio::DatagramSocket, //
+    }, //
+};
+
+fn translate_res_for_stdio(res: Result<IoStatus, Error>) -> Result<usize, io::Error> {
+    match res {
+        Ok(IoStatus::Ok(bytes)) => Ok(bytes),
+        Ok(IoStatus::EndOfStream) | Err(Error::Io(IoError::EndOfStream)) => Ok(0),
+        Ok(IoStatus::Retry(TlsRetryReason::WantRead | TlsRetryReason::WantWrite)) => {
+            Err(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
+        }
+        Ok(IoStatus::Retry(reason)) => Err(io::Error::new(io::ErrorKind::Other, reason)),
+        Ok(IoStatus::Err) => Err(io::Error::new(
+            io::ErrorKind::Other,
+            "The transport has failed the I/O operation",
+        )),
+        Ok(IoStatus::Empty) => Err(io::Error::new(
+            io::ErrorKind::ConnectionReset,
+            "connection reset or panicked",
+        )),
+        Err(
+            e @ (Error::Library(..)
+            | Error::Configuration(..)
+            | Error::TlsReason(..)
+            | Error::PemReason(..)
+            | Error::Quic(..)
+            | Error::Pki(..)
+            | Error::Io(..)
+            | Error::Unknown(..)),
+        ) => Err(io::Error::new(io::ErrorKind::Other, e)),
+    }
+}
+
+impl<R> io::Read for TlsConnectionRef<R, TlsMode> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let mut buf = ReceiveBuffer::new(buf);
+        let res = self.sync_read(&mut buf);
+        translate_res_for_stdio(res)
+    }
+}
+
+impl<R> io::Write for TlsConnectionRef<R, TlsMode> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        translate_res_for_stdio(self.sync_write(buf))
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+fn translate_result_for_datagram(res: Result<IoStatus, Error>) -> AbstractSocketResult {
+    match res {
+        Ok(IoStatus::Ok(bytes)) => AbstractSocketResult::Ok(bytes),
+        Ok(IoStatus::EndOfStream) | Err(Error::Io(IoError::EndOfStream)) => {
+            AbstractSocketResult::EndOfStream
+        }
+        Ok(IoStatus::Retry(_)) => AbstractSocketResult::Retry,
+        Ok(IoStatus::Empty | IoStatus::Err) => AbstractSocketResult::Err(Box::new(io::Error::new(
+            io::ErrorKind::Other,
+            "transport failed or empty",
+        ))),
+        Err(e) => AbstractSocketResult::Err(Box::new(io::Error::new(io::ErrorKind::Other, e))),
+    }
+}
+
+impl<R> DatagramSocket for TlsConnectionRef<R, DtlsMode> {
+    fn send(&mut self, datagram: &[u8]) -> AbstractSocketResult {
+        translate_result_for_datagram(self.sync_write(datagram))
+    }
+
+    fn recv(&mut self, datagram: &mut [u8]) -> AbstractSocketResult {
+        let mut datagram = ReceiveBuffer::new(datagram);
+        translate_result_for_datagram(self.sync_read(&mut datagram))
+    }
+}
diff --git a/rust/bssl-tls/src/connection/io/tokio_io_impl.rs b/rust/bssl-tls/src/connection/io/tokio_io_impl.rs
new file mode 100644
index 0000000..40a3c34
--- /dev/null
+++ b/rust/bssl-tls/src/connection/io/tokio_io_impl.rs
@@ -0,0 +1,142 @@
+// Copyright 2026 The BoringSSL Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{
+    fmt::Display,
+    io,
+    pin::Pin,
+    task::{
+        Context,
+        Poll, //
+    }, //
+};
+
+use tokio::io::{
+    AsyncRead,
+    AsyncWrite,
+    ReadBuf, //
+};
+
+use super::{
+    IoStatus,
+    TlsConnectionRef,
+    TlsMode, //
+};
+use crate::connection::lifecycle::ShutdownStatus;
+
+#[inline]
+fn handle_io_status<T>(status: IoStatus) -> Poll<io::Result<T>> {
+    match status {
+        IoStatus::EndOfStream | IoStatus::Ok(_) => unreachable!(),
+        IoStatus::Retry(_) => unreachable!("we should have handled retry earlier"),
+        IoStatus::Err => Poll::Ready(Err(io::Error::new(
+            io::ErrorKind::Other,
+            "The transport has failed the I/O operation",
+        ))),
+        IoStatus::Empty => Poll::Ready(Err(io::Error::new(
+            io::ErrorKind::ConnectionReset,
+            "TlsConnection has no backing transport or the transport has panicked",
+        ))),
+    }
+}
+
+impl<R> AsyncRead for TlsConnectionRef<R, TlsMode> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let status = match self.as_mut().aread_inner(buf.initialize_unfilled(), cx) {
+            Ok(status) => status,
+            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+        };
+        match status {
+            None => Poll::Pending,
+            Some(IoStatus::Ok(bytes)) => {
+                buf.advance(bytes);
+                Poll::Ready(Ok(()))
+            }
+            Some(IoStatus::EndOfStream) => Poll::Ready(Ok(())),
+            Some(status) => handle_io_status(status),
+        }
+    }
+}
+
+impl<R> AsyncWrite for TlsConnectionRef<R, TlsMode> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        let status = match self.as_mut().awrite_inner(buf, cx) {
+            Ok(status) => status,
+            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+        };
+        match status {
+            None => Poll::Pending,
+            Some(IoStatus::Ok(bytes)) => Poll::Ready(Ok(bytes)),
+            Some(IoStatus::EndOfStream) => Poll::Ready(Ok(0)),
+            Some(status) => handle_io_status(status),
+        }
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        let status = match self.as_mut().aflush_inner(cx) {
+            Ok(status) => status,
+            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+        };
+        match status {
+            None => Poll::Pending,
+            Some(IoStatus::Ok(_)) => Poll::Ready(Ok(())),
+            Some(IoStatus::EndOfStream) => Poll::Ready(Ok(())),
+            Some(status) => handle_io_status(status),
+        }
+    }
+
+    /// # Warning ⚠️
+    ///
+    /// Calling this may fail with error [`NeedToDrainAppData`].
+    /// This is not a hard error.
+    /// Caller should continue reading from the connection until the end of the stream.
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.as_mut().ashutdown_inner(cx) {
+            Ok(Some(ShutdownStatus::CloseNotifyReceived)) => Poll::Ready(Ok(())),
+            Ok(Some(ShutdownStatus::RemainingApplicationData)) => Poll::Ready(Err(io::Error::new(
+                io::ErrorKind::Other,
+                NeedToDrainAppData,
+            ))),
+            Ok(Some(ShutdownStatus::EndOfStream)) => Poll::Ready(Err(io::Error::new(
+                io::ErrorKind::UnexpectedEof,
+                "unexpected eof while waiting for peek close_notify",
+            ))),
+            Ok(Some(ShutdownStatus::CloseNotifyPosted)) => unreachable!(),
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct NeedToDrainAppData;
+
+impl Display for NeedToDrainAppData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "caller needs to drain application data before polling on shutdown again"
+        )
+    }
+}
+
+impl std::error::Error for NeedToDrainAppData {}
diff --git a/rust/bssl-tls/src/connection/lifecycle.rs b/rust/bssl-tls/src/connection/lifecycle.rs
index 50552ad..6d921e3 100644
--- a/rust/bssl-tls/src/connection/lifecycle.rs
+++ b/rust/bssl-tls/src/connection/lifecycle.rs
@@ -30,10 +30,10 @@
 use crate::{
     check_tls_error,
     connection::{
+        methods::HasTlsConnectionMethod, //
         Client,
         Server,
         TlsConnectionRef,
-        methods::HasTlsConnectionMethod, //
     },
     context::{
         HasBasicIo,
@@ -156,14 +156,17 @@
 where
     M: HasTlsConnectionMethod,
 {
-    /// Continue the handshake.
+    /// Drive the handshake.
     ///
     /// Call this method after the initial [`Self::accept`] or [`Self::connect`],
     /// should the handshake be suspended.
-    pub fn do_handshake(&mut self) -> Result<&mut Self, Error> {
+    ///
+    /// This method returns `Ok(None)` to signal handshake completion;
+    /// otherwise, `Ok(Some(reason))` is returned and the suspension `reason` must be resolved first
+    /// before this method can make progress again.
+    pub fn do_handshake(&mut self) -> Result<Option<TlsRetryReason>, Error> {
         let conn = self.ptr();
-        check_tls_error!(conn, bssl_sys::SSL_do_handshake(conn));
-        Ok(self)
+        Ok(check_tls_error!(conn, bssl_sys::SSL_do_handshake(conn)))
     }
 }
 
@@ -172,10 +175,12 @@
     M: HasTlsConnectionMethod,
 {
     /// Accept a connection by responding to `ClientHello` with `ServerHello`.
-    pub fn accept(&mut self) -> Result<&mut Self, Error> {
-        let conn = self.ptr();
-        check_tls_error!(conn, bssl_sys::SSL_accept(conn));
-        Ok(self)
+    ///
+    /// This method returns `Ok(None)` to signal handshake completion;
+    /// otherwise, given `Ok(Some(reason))` the suspension `reason` must be resolved first
+    /// before calling [`Self::do_handshake`] can make progress again.
+    pub fn accept(&mut self) -> Result<Option<TlsRetryReason>, Error> {
+        self.do_handshake()
     }
 }
 
@@ -184,10 +189,12 @@
     M: HasTlsConnectionMethod,
 {
     /// Initiate a connection by sending a `ClientHello`.
-    pub fn connect(&mut self) -> Result<&mut Self, Error> {
-        let conn = self.ptr();
-        check_tls_error!(conn, bssl_sys::SSL_connect(conn));
-        Ok(self)
+    ///
+    /// This method returns `Ok(None)` to signal handshake completion;
+    /// otherwise, given `Ok(Some(reason))` the suspension `reason` must be resolved first
+    /// before calling [`Self::do_handshake`] can make progress again.
+    pub fn connect(&mut self) -> Result<Option<TlsRetryReason>, Error> {
+        self.do_handshake()
     }
 }
 
@@ -215,6 +222,8 @@
 {
     /// Perform synchronising shutdown.
     ///
+    /// If the method returns `Ok(None)`, the shutdown will not progress until I/O makes progress.
+    ///
     /// # Shutdown protocol
     /// A live connection can be actively shut down by calling this method at most two times.
     /// The first call will send `close_notify` down the transport.
@@ -231,25 +240,30 @@
     ///   The connection is then in terminal state.
     /// To process the remaining application data, normal reading should continue until the end of
     /// stream, at which [`Self::sync_shutdown`] can be called again to set the connection to the terminal state.
-    pub fn sync_shutdown(&mut self) -> Result<ShutdownStatus, Error> {
+    pub fn sync_shutdown(&mut self) -> Result<Option<ShutdownStatus>, Error> {
         let rc = unsafe {
             // Safety: we have exclusive access to the connection state.
             bssl_sys::SSL_shutdown(self.ptr())
         };
         if self.is_write_closed() {
-            return Ok(ShutdownStatus::EndOfStream);
+            return Ok(Some(ShutdownStatus::EndOfStream));
         }
         match rc {
-            0 => Ok(ShutdownStatus::CloseNotifyPosted),
-            1 => Ok(ShutdownStatus::CloseNotifyReceived),
+            0 => Ok(Some(ShutdownStatus::CloseNotifyPosted)),
+            1 => Ok(Some(ShutdownStatus::CloseNotifyReceived)),
             _ => match self.categorise_error_for_io(rc) {
                 Ok(IoStatus::Ok(_)) => unreachable!(),
-                Ok(IoStatus::Empty | IoStatus::EndOfStream) => Ok(ShutdownStatus::EndOfStream),
-                Ok(IoStatus::Retry(reason)) => Err(Error::TlsRetry(reason)),
-                Err(Error::TlsReason(TlsErrorReason::ApplicationDataOnShutdown)) => {
-                    Ok(ShutdownStatus::RemainingApplicationData)
+                Ok(IoStatus::Empty | IoStatus::EndOfStream) => {
+                    Ok(Some(ShutdownStatus::EndOfStream))
                 }
-                Err(Error::Library(0, _, _)) => Ok(ShutdownStatus::CloseNotifyReceived),
+                Ok(IoStatus::Retry(TlsRetryReason::WantRead | TlsRetryReason::WantWrite)) => {
+                    Ok(None)
+                }
+                Ok(IoStatus::Retry(reason)) => panic!("unexpected retry reason {reason:?}"),
+                Err(Error::TlsReason(TlsErrorReason::ApplicationDataOnShutdown)) => {
+                    Ok(Some(ShutdownStatus::RemainingApplicationData))
+                }
+                Err(Error::Library(0, _, _)) => Ok(Some(ShutdownStatus::CloseNotifyReceived)),
                 Ok(IoStatus::Err) => Err(Error::Unknown(Box::new("transport error".to_string()))),
                 Err(e) => Err(e),
             },
@@ -265,18 +279,15 @@
     ///
     /// The caller needs to ensure that any pending operations during the handshake are resolved,
     /// before polling [`async_handshake`] again.
-    pub fn async_handshake(&mut self) -> impl Send + Future<Output = Result<(), Error>> + '_ {
+    pub fn async_handshake(
+        &mut self,
+    ) -> impl Send + Future<Output = Result<Option<TlsRetryReason>, Error>> + '_ {
         poll_fn(move |cx| {
             self.set_waker(cx.waker());
             match self.do_handshake() {
-                Ok(_) => Poll::Ready(Ok(())),
-                Err(Error::TlsRetry(r)) => {
-                    if matches!(r, TlsRetryReason::WantRead | TlsRetryReason::WantWrite) {
-                        Poll::Pending
-                    } else {
-                        Poll::Ready(Err(Error::TlsRetry(r)))
-                    }
-                }
+                Ok(Some(TlsRetryReason::WantRead | TlsRetryReason::WantWrite)) => Poll::Pending,
+                Ok(Some(reason)) => Poll::Ready(Ok(Some(reason))),
+                Ok(None) => Poll::Ready(Ok(None)),
                 Err(e) => Poll::Ready(Err(e)),
             }
         })
@@ -295,8 +306,8 @@
         poll_fn(move |cx| {
             self.set_waker(cx.waker());
             match self.do_handshake() {
-                Ok(_) => Poll::Ready(Ok(())),
-                Err(Error::TlsRetry(_)) => Poll::Pending,
+                Ok(Some(_)) => Poll::Pending,
+                Ok(None) => Poll::Ready(Ok(())),
                 Err(e) => Poll::Ready(Err(e)),
             }
         })
diff --git a/rust/bssl-tls/src/connection/methods.rs b/rust/bssl-tls/src/connection/methods.rs
index 640ecf9..7a05231 100644
--- a/rust/bssl-tls/src/connection/methods.rs
+++ b/rust/bssl-tls/src/connection/methods.rs
@@ -60,6 +60,10 @@
         }
     }
 
+    pub fn set_pending_reason(&mut self, reason: TlsRetryReason) {
+        self.pending_reason = Some(reason);
+    }
+
     pub fn take_pending_reason(&mut self) -> Option<TlsRetryReason> {
         self.pending_reason.take()
     }
diff --git a/rust/bssl-tls/src/errors.rs b/rust/bssl-tls/src/errors.rs
index 66d0c0d..f5aeebc 100644
--- a/rust/bssl-tls/src/errors.rs
+++ b/rust/bssl-tls/src/errors.rs
@@ -47,8 +47,6 @@
     Library(u32, Option<LibCode>, Option<i32>),
     /// Configuration errors
     Configuration(ConfigurationError),
-    /// TLS retryable errors
-    TlsRetry(TlsRetryReason),
     /// Other TLS errors with reason
     TlsReason(TlsErrorReason),
     /// PEM encoding failures
@@ -107,6 +105,10 @@
         }
     }
 
+    fn is_trivial(&self) -> bool {
+        matches!(self, Self::Library(0, _, _))
+    }
+
     pub(crate) fn extract_lib_err() -> Self {
         let packed_error = unsafe {
             // Safety: extracting error code does not have side-effect
@@ -120,14 +122,21 @@
         error
     }
 
-    pub(crate) fn extract_tls_err(code: c_int) -> Self {
+    pub(crate) fn extract_tls_err(code: c_int) -> Result<TlsRetryReason, Self> {
+        let lib_err = Self::extract_lib_err();
         if code == bssl_sys::SSL_ERROR_SSL {
-            return Self::extract_lib_err();
+            return Err(lib_err);
         }
-        if let Ok(err) = TlsRetryReason::try_from(code) {
-            return Self::TlsRetry(err);
+        if let Ok(reason) = TlsRetryReason::try_from(code) {
+            return Ok(reason);
         }
-        Self::Unknown(Box::new(alloc::format!("unknown tls error ({code})")))
+        if lib_err.is_trivial() {
+            Err(Self::Unknown(Box::new(alloc::format!(
+                "unknown tls error ({code})"
+            ))))
+        } else {
+            Err(lib_err)
+        }
     }
 }
 
@@ -167,7 +176,6 @@
                 f.write_str(&err_str.to_string_lossy())
             }
             Error::Configuration(err) => Display::fmt(err, f),
-            Error::TlsRetry(err) => Display::fmt(err, f),
             Error::TlsReason(err) => Display::fmt(err, f),
             Error::PemReason(err) => Debug::fmt(err, f),
             Error::Quic(err) => Display::fmt(err, f),
@@ -718,6 +726,7 @@
 bssl_enum! {
     /// TLS errors
     #[derive(Debug, Clone, Copy)]
+    #[must_use]
     pub enum TlsRetryReason: i32 {
         /// TLS is blocked on read.
         WantRead = bssl_sys::SSL_ERROR_WANT_READ as i32,
@@ -754,6 +763,8 @@
     }
 }
 
+impl core::error::Error for TlsRetryReason {}
+
 impl Display for TlsRetryReason {
     fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
         match self {
diff --git a/rust/bssl-tls/src/ffi.rs b/rust/bssl-tls/src/ffi.rs
index d141417..35f5e03 100644
--- a/rust/bssl-tls/src/ffi.rs
+++ b/rust/bssl-tls/src/ffi.rs
@@ -16,9 +16,9 @@
     marker::PhantomData,
     mem::MaybeUninit,
     ptr::{
-        NonNull,
         null,
         null_mut, //
+        NonNull,
     },
     slice::{
         from_raw_parts,
@@ -224,7 +224,6 @@
     /// - all reads into the byte under the returned pointer and those [`Self::remaining`]
     ///   bytes following it must be proceeded by at least one write; otherwise, it is **undefined
     ///   behaviour**.
-    #[allow(unused)]
     pub(crate) unsafe fn head(&mut self) -> *mut u8 {
         debug_assert!(self.cursor <= self.capacity && self.cursor <= isize::MAX as usize);
         unsafe {
@@ -240,7 +239,6 @@
     /// # Safety
     /// The bytes in between [`Self::head`] and `Self::head() + bytes` must have been filled
     /// by the caller before calling this method.
-    #[allow(unused)]
     pub(crate) unsafe fn advance(&mut self, bytes: usize) {
         self.cursor += bytes;
         debug_assert!(self.cursor <= self.capacity);
diff --git a/rust/bssl-tls/src/io.rs b/rust/bssl-tls/src/io.rs
index 590b6df..b1e23ea 100644
--- a/rust/bssl-tls/src/io.rs
+++ b/rust/bssl-tls/src/io.rs
@@ -45,6 +45,9 @@
     },
 };
 
+#[cfg(feature = "std")]
+pub mod stdio;
+
 /// A wrapper around a `dyn AbstractSocket`, delegating BIO methods to the
 /// underlying `AbstractSocket` implementations.
 ///
@@ -63,7 +66,7 @@
     io_err: Option<Box<dyn core::error::Error + Send + Sync>>,
 }
 
-/// Safety: `socket` field is a exclusively owned `Box<dyn AbstractSocket>` pointer,
+/// Safety: `socket` field is an exclusively owned `Box<dyn AbstractSocket>` pointer,
 /// and `AbstractSocket: Send + Sync`.
 unsafe impl Send for RustBio {}
 unsafe impl Sync for RustBio {}
@@ -168,6 +171,10 @@
         Ok(RustBioHandle(bio))
     }
 
+    pub fn take_io_err(&mut self) -> Option<Box<dyn core::error::Error + Send + Sync>> {
+        self.io_err.take()
+    }
+
     fn transform_result(
         &mut self,
         res: AbstractSocketResult,
@@ -185,7 +192,7 @@
     }
 }
 
-/// A exclusively owned handle to a BIO constructed by this crate.
+/// An exclusively owned handle to a BIO constructed by this crate.
 pub(crate) struct RustBioHandle(NonNull<bssl_sys::BIO>);
 
 impl RustBioHandle {
diff --git a/rust/bssl-tls/src/io/stdio.rs b/rust/bssl-tls/src/io/stdio.rs
new file mode 100644
index 0000000..8da2071
--- /dev/null
+++ b/rust/bssl-tls/src/io/stdio.rs
@@ -0,0 +1,50 @@
+// Copyright 2026 The BoringSSL Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! TLS I/O protocols under `std`
+
+use std::{
+    io,
+    task::{Context, Poll},
+};
+
+use super::AbstractSocketResult;
+
+/// A datagram socket protocol
+pub trait DatagramSocket: Send {
+    /// Send a complete datagram through the socket.
+    ///
+    /// By returning [`AbstractSocketResult::Retry`] the socket signals that the datagram
+    /// has not been sent down the transport.
+    fn send(&mut self, datagram: &[u8]) -> AbstractSocketResult;
+    /// Receive a complete datagram through the socket.
+    ///
+    /// By returning [`AbstractSocketResult::Retry`] the socket signals that a datagram
+    /// has not been received from the transport.
+    ///
+    /// If the `datagram` is not large enough to receive the whole datagram,
+    /// the datagram will be truncated while the actual size of the consumed datagram is reported
+    /// as [`AbstractSocketResult::Ok`].
+    ///
+    /// The datagram will be consumed on successful reception, even with `datagram.is_empty()`.
+    fn recv(&mut self, datagram: &mut [u8]) -> AbstractSocketResult;
+}
+
+/// Protocol and mechanism to register interest in an `async` runtime.
+pub trait PollFor<Io> {
+    /// Register an interest in reading from the `io` object.
+    fn poll_read(&mut self, async_ctx: &mut Context<'_>) -> Poll<Result<(), io::Error>>;
+    /// Register an interest in writing to the `io` object.
+    fn poll_write(&mut self, async_ctx: &mut Context<'_>) -> Poll<Result<(), io::Error>>;
+}
diff --git a/rust/bssl-tls/src/macros.rs b/rust/bssl-tls/src/macros.rs
index c1258ec..eea95e3 100644
--- a/rust/bssl-tls/src/macros.rs
+++ b/rust/bssl-tls/src/macros.rs
@@ -20,8 +20,8 @@
         unsafe {
             // Safety: we have exclusive access to the connection state.
             match ::bssl_sys::SSL_get_error($tls, $e) {
-                0 => {}
-                rc => return Err($crate::errors::Error::extract_tls_err(rc)),
+                0 => None,
+                rc => Some($crate::errors::Error::extract_tls_err(rc)?),
             }
         }
     };