Add AnyStream/AnyListener types to combine TCP and TLS streams/listeners.

Add `Accept` trait to abstract over different listener types.

Author: khvzak

src/net/common.rs

@@ -1,9 +1,19 @@
 //! Common types and traits shared between client and server TLS implementations.
 
+use std::any::TypeId;
 use std::borrow::Cow;
+use std::pin::Pin;
+use std::task::{Context, Poll};
 use std::{fmt, io};
 
-use mlua::{IntoLua, Lua, Result, Value};
+use mlua::{AnyUserData, Error, FromLua, IntoLua, Lua, MaybeSend, Result, Value};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+use super::tcp::{TcpListener, TcpStream};
+#[cfg(feature = "tls")]
+use super::tls::{TlsListener, TlsStream};
+#[cfg(unix)]
+use super::unix::{UnixListener, UnixStream};
 
 /// Socket address that can be either TCP or Unix domain socket.
 pub enum AnySocketAddr {
@@ -60,3 +70,215 @@ impl AddressProvider for tokio::net::UnixStream {
         Ok(AnySocketAddr::Unix(self.peer_addr()?))
     }
 }
+
+/// A stream that can be either TCP or Unix domain socket, possibly wrapped in TLS.
+pub enum AnyStream {
+    Tcp(TcpStream),
+    #[cfg(unix)]
+    Unix(UnixStream),
+    #[cfg(feature = "tls")]
+    TcpTls(TlsStream<TcpStream>),
+    #[cfg(all(unix, feature = "tls"))]
+    UnixTls(TlsStream<UnixStream>),
+}
+
+impl AsyncRead for AnyStream {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            AnyStream::Tcp(s) => Pin::new(s).poll_read(cx, buf),
+            #[cfg(unix)]
+            AnyStream::Unix(s) => Pin::new(s).poll_read(cx, buf),
+            #[cfg(feature = "tls")]
+            AnyStream::TcpTls(s) => Pin::new(s).poll_read(cx, buf),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyStream::UnixTls(s) => Pin::new(s).poll_read(cx, buf),
+        }
+    }
+}
+
+impl AsyncWrite for AnyStream {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        match self.get_mut() {
+            AnyStream::Tcp(s) => Pin::new(s).poll_write(cx, buf),
+            #[cfg(unix)]
+            AnyStream::Unix(s) => Pin::new(s).poll_write(cx, buf),
+            #[cfg(feature = "tls")]
+            AnyStream::TcpTls(s) => Pin::new(s).poll_write(cx, buf),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyStream::UnixTls(s) => Pin::new(s).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            AnyStream::Tcp(s) => Pin::new(s).poll_flush(cx),
+            #[cfg(unix)]
+            AnyStream::Unix(s) => Pin::new(s).poll_flush(cx),
+            #[cfg(feature = "tls")]
+            AnyStream::TcpTls(s) => Pin::new(s).poll_flush(cx),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyStream::UnixTls(s) => Pin::new(s).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            AnyStream::Tcp(s) => Pin::new(s).poll_shutdown(cx),
+            #[cfg(unix)]
+            AnyStream::Unix(s) => Pin::new(s).poll_shutdown(cx),
+            #[cfg(feature = "tls")]
+            AnyStream::TcpTls(s) => Pin::new(s).poll_shutdown(cx),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyStream::UnixTls(s) => Pin::new(s).poll_shutdown(cx),
+        }
+    }
+}
+
+impl IntoLua for AnyStream {
+    fn into_lua(self, lua: &Lua) -> Result<Value> {
+        match self {
+            AnyStream::Tcp(s) => lua.create_userdata(s).map(Value::UserData),
+            #[cfg(unix)]
+            AnyStream::Unix(s) => lua.create_userdata(s).map(Value::UserData),
+            #[cfg(feature = "tls")]
+            AnyStream::TcpTls(s) => lua.create_userdata(s).map(Value::UserData),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyStream::UnixTls(s) => lua.create_userdata(s).map(Value::UserData),
+        }
+    }
+}
+
+impl FromLua for AnyStream {
+    fn from_lua(value: Value, lua: &Lua) -> Result<Self> {
+        let value = lua.unpack::<AnyUserData>(value)?;
+        match value.type_id() {
+            Some(id) if id == TypeId::of::<TcpStream>() => {
+                let stream = value.take::<TcpStream>()?;
+                Ok(AnyStream::Tcp(stream))
+            }
+            #[cfg(unix)]
+            Some(id) if id == TypeId::of::<UnixStream>() => {
+                let stream = value.take::<UnixStream>()?;
+                Ok(AnyStream::Unix(stream))
+            }
+            #[cfg(feature = "tls")]
+            Some(id) if id == TypeId::of::<TlsStream<TcpStream>>() => {
+                let stream = value.take::<TlsStream<TcpStream>>()?;
+                Ok(AnyStream::TcpTls(stream))
+            }
+            #[cfg(all(unix, feature = "tls"))]
+            Some(id) if id == TypeId::of::<TlsStream<UnixStream>>() => {
+                let stream = value.take::<TlsStream<UnixStream>>()?;
+                Ok(AnyStream::UnixTls(stream))
+            }
+            _ => {
+                let type_name = value.type_name().ok().flatten();
+                let type_name = type_name.as_deref().unwrap_or("unknown");
+                Err(Error::FromLuaConversionError {
+                    from: "UserData",
+                    to: "AnyStream".to_string(),
+                    message: Some(format!("expected TcpStream or UnixStream, got {type_name}",)),
+                })
+            }
+        }
+    }
+}
+
+/// A listener that can be either TCP or Unix domain socket, possibly wrapped in TLS.
+pub enum AnyListener {
+    Tcp(TcpListener),
+    #[cfg(unix)]
+    Unix(UnixListener),
+    #[cfg(feature = "tls")]
+    TcpTls(TlsListener<TcpListener>),
+    #[cfg(all(unix, feature = "tls"))]
+    UnixTls(TlsListener<UnixListener>),
+}
+
+/// Trait for accepting incoming connections from various listener types.
+pub trait Accept {
+    type Stream: AsyncRead + AsyncWrite + Unpin + MaybeSend + 'static;
+
+    /// Get the local address that the listener is bound to.
+    fn local_addr(&self) -> io::Result<AnySocketAddr>;
+
+    /// Accept an incoming connection.
+    #[allow(async_fn_in_trait)]
+    async fn accept(&self) -> io::Result<(Self::Stream, AnySocketAddr)>;
+}
+
+impl Accept for AnyListener {
+    type Stream = AnyStream;
+
+    fn local_addr(&self) -> io::Result<AnySocketAddr> {
+        match self {
+            AnyListener::Tcp(l) => l.local_addr(),
+            #[cfg(unix)]
+            AnyListener::Unix(l) => l.local_addr(),
+            #[cfg(feature = "tls")]
+            AnyListener::TcpTls(l) => l.local_addr(),
+            #[cfg(all(unix, feature = "tls"))]
+            AnyListener::UnixTls(l) => l.local_addr(),
+        }
+    }
+
+    async fn accept(&self) -> io::Result<(AnyStream, AnySocketAddr)> {
+        match self {
+            AnyListener::Tcp(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok((AnyStream::Tcp(stream), addr))
+            }
+            #[cfg(unix)]
+            AnyListener::Unix(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok((AnyStream::Unix(stream), addr))
+            }
+            #[cfg(feature = "tls")]
+            AnyListener::TcpTls(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok((AnyStream::TcpTls(stream), addr))
+            }
+            #[cfg(all(unix, feature = "tls"))]
+            AnyListener::UnixTls(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok((AnyStream::UnixTls(stream), addr))
+            }
+        }
+    }
+}
+
+impl FromLua for AnyListener {
+    fn from_lua(value: Value, lua: &Lua) -> Result<Self> {
+        let value = lua.unpack::<AnyUserData>(value)?;
+        match value.type_id() {
+            Some(id) if id == TypeId::of::<TcpListener>() => {
+                let listener = value.take::<TcpListener>()?;
+                Ok(AnyListener::Tcp(listener))
+            }
+            #[cfg(unix)]
+            Some(id) if id == TypeId::of::<UnixListener>() => {
+                let listener = value.take::<UnixListener>()?;
+                Ok(AnyListener::Unix(listener))
+            }
+            #[cfg(feature = "tls")]
+            Some(id) if id == TypeId::of::<TlsListener<TcpListener>>() => {
+                let listener = value.take::<TlsListener<TcpListener>>()?;
+                Ok(AnyListener::TcpTls(listener))
+            }
+            #[cfg(all(unix, feature = "tls"))]
+            Some(id) if id == TypeId::of::<TlsListener<UnixListener>>() => {
+                let listener = value.take::<TlsListener<UnixListener>>()?;
+                Ok(AnyListener::UnixTls(listener))
+            }
+            _ => {
+                let type_name = value.type_name().ok().flatten();
+                let type_name = type_name.as_deref().unwrap_or("unknown");
+                Err(Error::FromLuaConversionError {
+                    from: "UserData",
+                    to: "AnyListener".to_string(),
+                    message: Some(format!("expected TcpListener or UnixListener, got {type_name}")),
+                })
+            }
+        }
+    }
+}

src/net/mod.rs

@@ -1,6 +1,7 @@
 use mlua::{Lua, Result, Table};
 
-pub use common::{AddressProvider, AnySocketAddr};
+pub use common::{AddressProvider, AnyListener, AnySocketAddr, AnyStream};
+pub use tcp::{TcpListener, TcpStream};
 
 /// A loader for the `net` module.
 fn loader(lua: &Lua) -> Result<Table> {
@@ -27,7 +28,7 @@ macro_rules! with_io_timeout {
     };
 }
 
-mod common;
+pub(crate) mod common;
 
 pub mod tcp;
 #[cfg(feature = "tls")]

src/net/tcp/listener.rs

@@ -6,14 +6,23 @@ use mlua::{Lua, Result, Table, UserData, UserDataMethods, UserDataRegistry};
 use tokio::net::lookup_host;
 
 use super::{SocketOptions, TcpSocket, TcpStream};
-use crate::net::common::AnySocketAddr;
+use crate::net::common::{Accept, AnySocketAddr};
 
 pub struct TcpListener(pub(crate) tokio::net::TcpListener);
 
-impl TcpListener {
-    pub(crate) fn local_addr(&self) -> io::Result<AnySocketAddr> {
+impl Accept for TcpListener {
+    type Stream = TcpStream;
+
+    fn local_addr(&self) -> io::Result<AnySocketAddr> {
         self.0.local_addr().map(AnySocketAddr::IP)
     }
+
+    async fn accept(&self) -> io::Result<(Self::Stream, AnySocketAddr)> {
+        let (stream, addr) = self.0.accept().await?;
+        let io = TcpStream::from(stream);
+        let addr = AnySocketAddr::IP(addr);
+        Ok((io, addr))
+    }
 }
 
 impl UserData for TcpListener {

src/net/tcp/stream.rs

@@ -1,10 +1,12 @@
 use std::io;
 use std::net::SocketAddr;
 use std::ops::{Deref, DerefMut};
+use std::pin::Pin;
 use std::result::Result as StdResult;
+use std::task::{Context, Poll};
 
 use mlua::{Lua, Result, String as LuaString, Table, UserData, UserDataMethods, UserDataRegistry};
-use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
+use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _, ReadBuf};
 use tokio::net::lookup_host;
 
 use super::{SocketOptions, TcpSocket};
@@ -55,6 +57,30 @@ impl AddressProvider for TcpStream {
     }
 }
 
+impl AsyncRead for TcpStream {
+    #[inline]
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.get_mut().stream).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for TcpStream {
+    #[inline]
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.get_mut().stream).poll_write(cx, buf)
+    }
+
+    #[inline]
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.get_mut().stream).poll_flush(cx)
+    }
+
+    #[inline]
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.get_mut().stream).poll_shutdown(cx)
+    }
+}
+
 impl UserData for TcpStream {
     fn register(registry: &mut UserDataRegistry<Self>) {
         registry.add_async_function("connect", connect);

src/net/tls/client.rs

@@ -12,9 +12,9 @@ use rustls::pki_types::{CertificateDer, DnsName, ServerName, UnixTime};
 use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme};
 
 use super::stream::TlsStream;
-use crate::net::tcp::TcpStream;
 #[cfg(unix)]
 use crate::net::unix::UnixStream;
+use crate::net::{AnyStream, TcpStream};
 
 /// TLS configuration options for client connections
 #[derive(Debug, Clone)]
@@ -222,9 +222,9 @@ static DEFAULT_TLS_CLIENT_CONFIG: LazyLock<TlsClientConfig> =
 
 /// Wrap a stream with TLS (client-side).
 pub async fn wrap_stream(
-    lua: Lua,
+    _lua: Lua,
     (stream, server_name, config): (AnyUserData, Option<String>, Option<TlsClientConfig>),
-) -> LuaResult<StdResult<AnyUserData, String>> {
+) -> LuaResult<StdResult<AnyStream, String>> {
     let server_name = server_name.and_then(|name| ServerName::try_from(name).ok());
     let config = config.unwrap_or_else(|| DEFAULT_TLS_CLIENT_CONFIG.clone());
 
@@ -236,11 +236,11 @@ pub async fn wrap_stream(
                 .or_else(|| host.and_then(|host| ServerName::try_from(host).ok()))
                 .or_else(|| stream.peer_addr().map(|addr| ServerName::from(addr.ip())).ok())
                 .unwrap_or_else(default_server_name);
-            match TlsStream::new_client(stream, server_name, config).await {
+            match TlsStream::new_client(stream.into(), server_name, config).await {
                 Ok(mut tls_stream) => {
                     tls_stream.set_read_timeout(read_timeout);
                     tls_stream.set_write_timeout(write_timeout);
-                    Ok(Ok(lua.create_userdata(tls_stream)?))
+                    Ok(Ok(AnyStream::TcpTls(tls_stream)))
                 }
                 Err(e) => Ok(Err(e.to_string())),
             }
@@ -250,11 +250,11 @@ pub async fn wrap_stream(
             #[rustfmt::skip]
             let UnixStream { stream, read_timeout, write_timeout } = stream.take::<UnixStream>()?;
             let server_name = server_name.unwrap_or_else(default_server_name);
-            match TlsStream::new_client(stream, server_name, config).await {
+            match TlsStream::new_client(stream.into(), server_name, config).await {
                 Ok(mut tls_stream) => {
                     tls_stream.set_read_timeout(read_timeout);
                     tls_stream.set_write_timeout(write_timeout);
-                    Ok(Ok(lua.create_userdata(tls_stream)?))
+                    Ok(Ok(AnyStream::UnixTls(tls_stream)))
                 }
                 Err(e) => Ok(Err(e.to_string())),
             }

src/net/tls/mod.rs

@@ -2,6 +2,9 @@
 
 use mlua::{Result, Table};
 
+pub use server::TlsListener;
+pub use stream::TlsStream;
+
 /// Registers the `tls` module in the given Lua state.
 pub fn register(lua: &mlua::Lua, name: Option<&str>) -> Result<Table> {
     let name = name.unwrap_or("@tls");