Add net/tls module

Author: khvzak

Cargo.toml

@@ -23,10 +23,11 @@ yaml = ["mlua/serde", "dep:ouroboros", "dep:serde", "dep:serde_yaml"]
 http = ["dep:http", "dep:http-serde-ext"]
 task = ["async"]
 net = ["task"]
+tls = ["net", "dep:rustls", "dep:tokio-rustls", "dep:rustls-pemfile", "dep:rustls-native-certs", "dep:webpki-roots"]
 
 [dependencies]
 bytes = "1"
-mlua = { version = "0.11", features = ["error-send"] }
+mlua = { version = "0.11", features = ["error-send", "macros"] }
 ouroboros = { version = "0.18", optional = true }
 serde = { version = "1.0", optional = true }
 serde_json = { version = "1.0", optional = true }
@@ -43,5 +44,12 @@ http-serde-ext = { version = "1.0", optional = true }
 tokio = { version = "1", features = ["full"], optional = true }
 tokio-util = { version = "0.7", features = ["time"], optional = true }
 
+# tls
+rustls = { version = "0.23", optional = true, default-features = false, features = ["aws-lc-rs"] }
+tokio-rustls = { version = "0.26", optional = true, default-features = false }
+rustls-pemfile = { version = "2", optional = true }
+rustls-native-certs = { version = "0.8", optional = true }
+webpki-roots = { version = "1", optional = true }
+
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }

src/macros.rs

@@ -25,7 +25,7 @@ macro_rules! opt_param {
             Ok(None) => Ok(None),
             Err(err) => {
                 use mlua::ErrorContext as _;
-                Err(err.with_context(|_| format!("invalid `{}`", $name)))
+                Err(err.with_context(|_| format!("invalid `{}` param", $name)))
             }
         }
     };
@@ -39,7 +39,20 @@ macro_rules! opt_param {
             Ok(None) => Ok(None),
             Err(err) => {
                 use mlua::ErrorContext as _;
-                Err(err.with_context(|_| format!("invalid `{}`", $name)))
+                Err(err.with_context(|_| format!("invalid `{}` param", $name)))
+            }
+        }
+    };
+}
+
+macro_rules! param {
+    ($table:expr, $name:expr) => {
+        match $table.raw_get::<Option<_>>($name) {
+            Ok(Some(v)) => Ok(v),
+            Ok(None) => Err(mlua::Error::runtime(format!("`{}` param is required", $name))),
+            Err(err) => {
+                use mlua::ErrorContext as _;
+                Err(err.with_context(|_| format!("invalid `{}` param", $name)))
             }
         }
     };

src/net/common.rs

@@ -0,0 +1,62 @@
+//! Common types and traits shared between client and server TLS implementations.
+
+use std::borrow::Cow;
+use std::{fmt, io};
+
+use mlua::{IntoLua, Lua, Result, Value};
+
+/// Socket address that can be either TCP or Unix domain socket.
+pub enum AnySocketAddr {
+    Tcp(std::net::SocketAddr),
+    #[cfg(unix)]
+    Unix(tokio::net::unix::SocketAddr),
+}
+
+impl fmt::Display for AnySocketAddr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            AnySocketAddr::Tcp(addr) => write!(f, "{addr}"),
+            #[cfg(unix)]
+            AnySocketAddr::Unix(addr) => {
+                let path = addr
+                    .as_pathname()
+                    .map(|p| p.to_string_lossy())
+                    .unwrap_or_else(|| Cow::Borrowed("(unnamed)"));
+                write!(f, "{path}")
+            }
+        }
+    }
+}
+
+impl IntoLua for AnySocketAddr {
+    fn into_lua(self, lua: &Lua) -> Result<Value> {
+        lua.create_string(self.to_string()).map(Value::String)
+    }
+}
+
+/// Trait for getting local and peer addresses from various stream types.
+pub trait AddressProvider {
+    fn local_addr(&self) -> io::Result<AnySocketAddr>;
+    fn peer_addr(&self) -> io::Result<AnySocketAddr>;
+}
+
+impl AddressProvider for tokio::net::TcpStream {
+    fn local_addr(&self) -> io::Result<AnySocketAddr> {
+        Ok(AnySocketAddr::Tcp(self.local_addr()?))
+    }
+
+    fn peer_addr(&self) -> io::Result<AnySocketAddr> {
+        Ok(AnySocketAddr::Tcp(self.peer_addr()?))
+    }
+}
+
+#[cfg(unix)]
+impl AddressProvider for tokio::net::UnixStream {
+    fn local_addr(&self) -> io::Result<AnySocketAddr> {
+        Ok(AnySocketAddr::Unix(self.local_addr()?))
+    }
+
+    fn peer_addr(&self) -> io::Result<AnySocketAddr> {
+        Ok(AnySocketAddr::Unix(self.peer_addr()?))
+    }
+}

src/net/mod.rs

@@ -1,5 +1,7 @@
 use mlua::{Lua, Result, Table};
 
+pub use common::{AddressProvider, AnySocketAddr};
+
 /// A loader for the `net` module.
 fn loader(lua: &Lua) -> Result<Table> {
     let t = lua.create_table()?;
@@ -25,7 +27,10 @@ macro_rules! with_io_timeout {
     };
 }
 
-pub mod tcp;
+mod common;
 
+pub mod tcp;
+#[cfg(feature = "tls")]
+pub mod tls;
 #[cfg(unix)]
 pub mod unix;

src/net/tcp/listener.rs

@@ -6,12 +6,19 @@ use mlua::{Lua, Result, Table, UserData, UserDataMethods, UserDataRegistry};
 use tokio::net::lookup_host;
 
 use super::{SocketOptions, TcpSocket, TcpStream};
+use crate::net::common::AnySocketAddr;
 
-pub struct TcpListener(tokio::net::TcpListener);
+pub struct TcpListener(pub(crate) tokio::net::TcpListener);
+
+impl TcpListener {
+    pub(crate) fn local_addr(&self) -> io::Result<AnySocketAddr> {
+        self.0.local_addr().map(AnySocketAddr::Tcp)
+    }
+}
 
 impl UserData for TcpListener {
     fn register(registry: &mut UserDataRegistry<Self>) {
-        registry.add_method("local_addr", |_, this, ()| Ok(this.0.local_addr()?.to_string()));
+        registry.add_method("local_addr", |_, this, ()| Ok(this.local_addr()?));
 
         registry.add_async_function("listen", listen);
 
@@ -24,9 +31,10 @@ impl UserData for TcpListener {
 
 pub async fn listen(
     _: Lua,
-    (addr, params): (String, Option<Table>),
+    (addr, port, params): (String, Option<u16>, Option<Table>),
 ) -> Result<StdResult<TcpListener, String>> {
-    let addrs = lua_try!(lookup_host(addr).await);
+    let port = port.unwrap_or(0);
+    let addrs = lua_try!(lookup_host((addr, port)).await);
 
     let sock_options = SocketOptions::from_table(&params)?;
     let backlog = opt_param!(params, "backlog")?;

src/net/tcp/socket.rs

@@ -1,4 +1,4 @@
-use std::io::Result as IoResult;
+use std::io;
 use std::net::SocketAddr;
 use std::ops::Deref;
 
@@ -26,15 +26,15 @@ pub(super) struct SocketOptions {
 }
 
 impl TcpSocket {
-    pub(crate) fn new_for_addr(addr: SocketAddr) -> IoResult<Self> {
+    pub(crate) fn new_for_addr(addr: SocketAddr) -> io::Result<Self> {
         let sock = match addr {
             SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
             SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
         };
         Ok(TcpSocket(sock))
     }
 
-    pub(crate) fn set_options(&self, options: SocketOptions) -> IoResult<()> {
+    pub(crate) fn set_options(&self, options: SocketOptions) -> io::Result<()> {
         if let Some(keepalive) = options.keepalive {
             self.set_keepalive(keepalive)?;
         }

src/net/tcp/stream.rs

@@ -1,3 +1,4 @@
+use std::io;
 use std::net::SocketAddr;
 use std::ops::{Deref, DerefMut};
 use std::result::Result as StdResult;
@@ -7,10 +8,12 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
 use tokio::net::lookup_host;
 
 use super::{SocketOptions, TcpSocket};
+use crate::net::{AddressProvider, AnySocketAddr};
 use crate::time::Duration;
 
 pub struct TcpStream {
     pub(crate) stream: tokio::net::TcpStream,
+    pub(crate) host: Option<String>,
     pub(crate) read_timeout: Option<Duration>,
     pub(crate) write_timeout: Option<Duration>,
 }
@@ -35,19 +38,29 @@ impl From<tokio::net::TcpStream> for TcpStream {
     fn from(stream: tokio::net::TcpStream) -> Self {
         TcpStream {
             stream,
+            host: None,
             read_timeout: None,
             write_timeout: None,
         }
     }
 }
 
+impl AddressProvider for TcpStream {
+    fn local_addr(&self) -> io::Result<AnySocketAddr> {
+        self.stream.local_addr().map(AnySocketAddr::Tcp)
+    }
+
+    fn peer_addr(&self) -> io::Result<AnySocketAddr> {
+        self.stream.peer_addr().map(AnySocketAddr::Tcp)
+    }
+}
+
 impl UserData for TcpStream {
     fn register(registry: &mut UserDataRegistry<Self>) {
         registry.add_async_function("connect", connect);
 
-        registry.add_method("local_addr", |_, this, ()| Ok(this.local_addr()?.to_string()));
-
-        registry.add_method("peer_addr", |_, this, ()| Ok(this.peer_addr()?.to_string()));
+        registry.add_method("local_addr", |_, this, ()| Ok(this.local_addr()?));
+        registry.add_method("peer_addr", |_, this, ()| Ok(this.peer_addr()?));
 
         registry.add_method_mut("set_read_timeout", |_, this, dur: Option<Duration>| {
             this.read_timeout = dur;
@@ -101,9 +114,9 @@ impl UserData for TcpStream {
 
 pub async fn connect(
     _: Lua,
-    (addr, params): (String, Option<Table>),
+    (host, port, params): (String, u16, Option<Table>),
 ) -> Result<StdResult<TcpStream, String>> {
-    let addrs = lua_try!(lookup_host(addr).await);
+    let addrs = lua_try!(lookup_host((&*host, port)).await);
     let options = SocketOptions::from_table(&params)?;
 
     let timeout = opt_param!(Duration, params, "timeout")?; // A single timeout for any operation
@@ -123,6 +136,7 @@ pub async fn connect(
             Ok(stream) => {
                 return Ok(Ok(TcpStream {
                     stream,
+                    host: Some(host.clone()),
                     read_timeout,
                     write_timeout,
                 }));