Add async support for DTLS

mbedtls/Cargo.toml

@@ -30,13 +30,14 @@ cbc = { version = "0.1.2", optional = true }
 rc2 = { version = "0.8.1", optional = true }
 cfg-if = "1.0.0"
 tokio = { version = "1.16.1", optional = true }
+async-trait = { version = "0.1", optional = true }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 rs-libc = "0.2.0"
 chrono = "0.4"
 
 [dependencies.mbedtls-sys-auto]
-version = "2.25.0"
+version = "2.28.0"
 default-features = false
 features = ["trusted_cert_callback", "threading"]
 path = "../mbedtls-sys"
@@ -78,7 +79,7 @@ pkcs12 = ["std", "yasna"]
 pkcs12_rc2 = ["pkcs12", "rc2", "cbc"]
 legacy_protocols = ["mbedtls-sys-auto/legacy_protocols"]
 async = ["std", "tokio","tokio/net","tokio/io-util", "tokio/macros"]
-async-rt = ["async", "tokio/rt", "tokio/sync", "tokio/rt-multi-thread"]
+async-rt = ["async","async-trait", "tokio/rt", "tokio/sync", "tokio/rt-multi-thread"]
 
 [[example]]
 name = "client"

mbedtls/src/ssl/async_io.rs

@@ -15,29 +15,63 @@ use crate::{
         io::{IoCallback, IoCallbackUnsafe},
     },
 };
+use async_trait::async_trait;
 use std::{
     future::Future,
     io::{Error as IoError, ErrorKind as IoErrorKind, Result as IoResult},
     pin::Pin,
+    result::Result as StdResult,
     task::{Context as TaskContext, Poll},
 };
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio::{
+    io::{AsyncRead, AsyncWrite, ReadBuf},
+    macros::support::poll_fn,
+    net::UdpSocket,
+};
 
 /// Marker type for an IO implementation that implements both
 /// `tokio::io::AsyncRead` and `tokio::io::AsyncWrite`.
 pub enum AsyncStream {}
 
-// TODO: Add enum `AnyAsyncIo` as marker type for an IO implementation that
-// doesn't implement `tokio::io::AsyncRead` and `tokio::io::AsyncWrite`.
+/// Marker type for an IO implementation that
+/// doesn't implement [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`].
+pub enum AnyAsyncIo {}
+
+/// Async read or write bytes or packets.
+///
+/// Implementors represent a socket or file descriptor that can be read from or
+/// written to.
+///
+/// By implementing [`AsyncIo`] you can wrap any type of
+/// [`AsyncIo`] with [`Context<T>::establish_async`] to protect that
+/// communication channel with (D)TLS. That [`Context`] then also implements
+/// [`AsyncIo`] so you can use it interchangeably.
+///
+/// If you are using byte streams and are using [`tokio::io`], you don't need
+/// this trait and can rely on [`tokio::io::AsyncRead`] and
+/// [`tokio::io::AsyncWrite`] instead.
+#[async_trait]
+pub trait AsyncIo {
+    async fn recv(&mut self, buf: &mut [u8]) -> IoResult<usize>;
+    async fn send(&mut self, buf: &[u8]) -> IoResult<usize>;
+}
+
+#[async_trait]
+impl<T: AsyncIo + Send> AsyncIo for Context<T> {
+    async fn recv(&mut self, buf: &mut [u8]) -> IoResult<usize> {
+        self.io_mut()
+            .ok_or(IoError::new(IoErrorKind::Other, "stream has been shutdown"))?
+            .recv(buf)
+            .await
+    }
 
-// TODO: Add `AsyncIo` trait for async IO that that doesn't implement
-// `tokio::io::AsyncRead` and `tokio::io::AsyncWrite`. For example:
-//     pub trait AsyncIo {
-//        async fn recv(&mut self, buf: &mut [u8]) -> Result<usize>;
-//         async fn send(&mut self, buf: &[u8]) -> Result<usize>;
-//     }
-// Could implement by using `async-trait` crate or
-// #![feature(async_fn_in_trait)] or Associated Types
+    async fn send(&mut self, buf: &[u8]) -> IoResult<usize> {
+        self.io_mut()
+            .ok_or(IoError::new(IoErrorKind::Other, "stream has been shutdown"))?
+            .send(buf)
+            .await
+    }
+}
 
 impl<'a, 'b, 'c, IO: AsyncRead + AsyncWrite + std::marker::Unpin + 'static> IoCallback<AsyncStream>
     for (&'a mut TaskContext<'b>, &'c mut IO)
@@ -62,6 +96,24 @@ impl<'a, 'b, 'c, IO: AsyncRead + AsyncWrite + std::marker::Unpin + 'static> IoCa
     }
 }
 
+impl<'a, 'b, 'c, IO: AsyncIo> IoCallback<AnyAsyncIo> for (&'a mut TaskContext<'b>, &'c mut IO) {
+    fn recv(&mut self, buf: &mut [u8]) -> Result<usize> {
+        match Pin::new(&mut self.1.recv(buf)).poll(self.0) {
+            Poll::Ready(Ok(n)) => Ok(n),
+            Poll::Ready(Err(_)) => Err(Error::NetRecvFailed),
+            Poll::Pending => Err(Error::SslWantRead),
+        }
+    }
+
+    fn send(&mut self, buf: &[u8]) -> Result<usize> {
+        match Pin::new(&mut self.1.send(buf)).poll(self.0) {
+            Poll::Ready(Ok(n)) => Ok(n),
+            Poll::Ready(Err(_)) => Err(Error::NetSendFailed),
+            Poll::Pending => Err(Error::SslWantWrite),
+        }
+    }
+}
+
 impl<T: Unpin + AsyncRead + AsyncWrite + 'static> Context<T> {
     pub async fn establish_async<IoType>(&mut self, io: T, hostname: Option<&str>) -> Result<()>
     where
@@ -121,14 +173,13 @@ where
             return Poll::Ready(Err(IoError::new(IoErrorKind::Other, "stream has been shutdown")));
         }
 
-        self
-            .with_bio_async(cx, |ssl_ctx| match ssl_ctx.async_write(buf) {
-                Err(Error::SslPeerCloseNotify) => Poll::Ready(Ok(0)),
-                Err(Error::SslWantWrite) => Poll::Pending,
-                Err(e) => Poll::Ready(Err(crate::private::error_to_io_error(e))),
-                Ok(i) => Poll::Ready(Ok(i)),
-            })
-            .unwrap_or_else(|| Poll::Ready(Err(crate::private::error_to_io_error(Error::NetSendFailed))))
+        self.with_bio_async(cx, |ssl_ctx| match ssl_ctx.async_write(buf) {
+            Err(Error::SslPeerCloseNotify) => Poll::Ready(Ok(0)),
+            Err(Error::SslWantWrite) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(crate::private::error_to_io_error(e))),
+            Ok(i) => Poll::Ready(Ok(i)),
+        })
+        .unwrap_or_else(|| Poll::Ready(Err(crate::private::error_to_io_error(Error::NetSendFailed))))
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<IoResult<()>> {
@@ -168,4 +219,42 @@ where
     }
 }
 
-// TODO: AsyncIo impl for tokio::net::UdpSocket
+/// A wrapper of [`tokio::net::UdpSocket`] on which
+/// [`tokio::net::UdpSocket::connect`] was successfully called.
+///
+/// Construct this type using [`ConnectedAsyncUdpSocket::connect`].
+pub struct ConnectedAsyncUdpSocket {
+    socket: UdpSocket,
+}
+
+impl ConnectedAsyncUdpSocket {
+    pub async fn connect<A: tokio::net::ToSocketAddrs>(socket: UdpSocket, addr: A) -> StdResult<Self, (IoError, UdpSocket)> {
+        match socket.connect(addr).await {
+            Ok(()) => Ok(ConnectedAsyncUdpSocket { socket }),
+            Err(e) => Err((e, socket)),
+        }
+    }
+
+    pub fn into_socket(self) -> UdpSocket {
+        self.socket
+    }
+}
+
+#[async_trait]
+impl AsyncIo for ConnectedAsyncUdpSocket {
+    async fn recv(&mut self, buf: &mut [u8]) -> IoResult<usize> {
+        poll_fn(|cx| {
+            let mut buf = ReadBuf::new(buf);
+            match self.socket.poll_recv(cx, &mut buf) {
+                Poll::Ready(res) => res,
+                Poll::Pending => return Poll::Pending,
+            }?;
+            Poll::Ready(Ok(buf.filled().len()))
+        })
+        .await
+    }
+
+    async fn send(&mut self, buf: &[u8]) -> IoResult<usize> {
+        poll_fn(|cx| self.socket.poll_send(cx, buf)).await
+    }
+}