Add net module

Author: khvzak

Cargo.toml

@@ -22,6 +22,7 @@ regex = ["dep:regex", "dep:ouroboros", "dep:quick_cache"]
 yaml = ["mlua/serde", "dep:ouroboros", "dep:serde", "dep:serde_yaml"]
 http = ["dep:http", "dep:http-serde-ext"]
 task = ["async"]
+net = ["task"]
 
 [dependencies]
 bytes = "1"

src/lib.rs

@@ -27,3 +27,6 @@ pub mod http;
 
 #[cfg(feature = "task")]
 pub mod task;
+
+#[cfg(feature = "net")]
+pub mod net;

src/net/mod.rs

@@ -0,0 +1,17 @@
+use mlua::{Lua, Result, Table};
+
+/// A loader for the `net` module.
+fn loader(lua: &Lua) -> Result<Table> {
+    let t = lua.create_table()?;
+    Ok(t)
+}
+
+/// Registers the `net` module in the given Lua state.
+pub fn register(lua: &Lua, name: Option<&str>) -> Result<Table> {
+    let name = name.unwrap_or("@net");
+    let value = loader(lua)?;
+    lua.register_module(name, &value)?;
+    Ok(value)
+}
+
+pub mod tcp;

src/net/tcp/listener.rs

@@ -0,0 +1,57 @@
+use std::io;
+use std::net::SocketAddr;
+use std::result::Result as StdResult;
+
+use mlua::{Lua, Result, Table, UserData, UserDataMethods, UserDataRegistry};
+use tokio::net::lookup_host;
+
+use super::{SocketOptions, TcpSocket, TcpStream};
+
+pub struct TcpListener(tokio::net::TcpListener);
+
+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_async_function("listen", listen);
+
+        registry.add_async_method("accept", |_, this, ()| async move {
+            let (stream, _) = lua_try!(this.0.accept().await);
+            Ok(Ok(TcpStream::from(stream)))
+        });
+    }
+}
+
+pub async fn listen(
+    _: Lua,
+    (addr, params): (String, Option<Table>),
+) -> Result<StdResult<TcpListener, String>> {
+    let addrs = lua_try!(lookup_host(addr).await);
+
+    let sock_options = SocketOptions::from_table(&params)?;
+    let backlog = opt_param!(params, "backlog")?;
+
+    let try_listen = |addr: SocketAddr| {
+        let sock = TcpSocket::new_for_addr(addr)?;
+        sock.set_options(sock_options)?;
+        sock.0.set_reuseaddr(true)?;
+        sock.0.bind(addr)?;
+        let listener = TcpListener(sock.0.listen(backlog.unwrap_or(1024))?);
+        io::Result::Ok(listener)
+    };
+
+    let mut last_err = None;
+    for addr in addrs {
+        match try_listen(addr) {
+            Ok(sock) => return Ok(Ok(sock)),
+            Err(e) => {
+                last_err = Some(e);
+                continue;
+            }
+        }
+    }
+
+    Ok(Err(last_err.map(|err| err.to_string()).unwrap_or_else(|| {
+        "could not resolve to any address".to_string()
+    })))
+}

src/net/tcp/mod.rs

@@ -0,0 +1,28 @@
+use mlua::{Lua, Result, Table};
+
+pub use listener::{TcpListener, listen};
+pub use stream::{TcpStream, connect};
+
+use socket::{SocketOptions, TcpSocket};
+
+/// A loader for the `net/tcp` module.
+fn loader(lua: &Lua) -> Result<Table> {
+    let t = lua.create_table()?;
+    t.set("TcpListener", lua.create_proxy::<TcpListener>()?)?;
+    t.set("TcpStream", lua.create_proxy::<TcpStream>()?)?;
+    t.set("listen", lua.create_async_function(listen)?)?;
+    t.set("connect", lua.create_async_function(connect)?)?;
+    Ok(t)
+}
+
+/// Registers the `net/tcp` module in the given Lua state.
+pub fn register(lua: &Lua, name: Option<&str>) -> Result<Table> {
+    let name = name.unwrap_or("@net/tcp");
+    let value = loader(lua)?;
+    lua.register_module(name, &value)?;
+    Ok(value)
+}
+
+mod listener;
+mod socket;
+mod stream;

src/net/tcp/socket.rs

@@ -0,0 +1,72 @@
+use std::io::Result as IoResult;
+use std::net::SocketAddr;
+use std::ops::Deref;
+
+use mlua::{Result, Table};
+use tokio::net::TcpSocket as TokioTcpSocket;
+
+pub(crate) struct TcpSocket(pub(crate) TokioTcpSocket);
+
+impl Deref for TcpSocket {
+    type Target = TokioTcpSocket;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+pub(super) struct SocketOptions {
+    keepalive: Option<bool>,
+    nodelay: Option<bool>,
+    recv_buffer_size: Option<u32>,
+    send_buffer_size: Option<u32>,
+    reuseraddr: Option<bool>,
+    reuseport: Option<bool>,
+}
+
+impl TcpSocket {
+    pub(crate) fn new_for_addr(addr: SocketAddr) -> IoResult<Self> {
+        let sock = match addr {
+            SocketAddr::V4(_) => TokioTcpSocket::new_v4()?,
+            SocketAddr::V6(_) => TokioTcpSocket::new_v6()?,
+        };
+        Ok(TcpSocket(sock))
+    }
+
+    pub(crate) fn set_options(&self, options: SocketOptions) -> IoResult<()> {
+        if let Some(keepalive) = options.keepalive {
+            self.set_keepalive(keepalive)?;
+        }
+        if let Some(nodelay) = options.nodelay {
+            self.set_nodelay(nodelay)?;
+        }
+        if let Some(reuseraddr) = options.reuseraddr {
+            self.set_reuseaddr(reuseraddr)?;
+        }
+        if let Some(reuseport) = options.reuseport {
+            self.set_reuseport(reuseport)?;
+        }
+        if let Some(recv_buffer_size) = options.recv_buffer_size {
+            self.set_recv_buffer_size(recv_buffer_size)?;
+        }
+        if let Some(send_buffer_size) = options.send_buffer_size {
+            self.set_send_buffer_size(send_buffer_size)?;
+        }
+        Ok(())
+    }
+}
+
+impl SocketOptions {
+    pub(crate) fn from_table(params: &Option<Table>) -> Result<Self> {
+        Ok(SocketOptions {
+            keepalive: opt_param!(params, "keepalive")?,
+            nodelay: opt_param!(params, "nodelay")?,
+            recv_buffer_size: opt_param!(params, "recv_buffer_size")?,
+            send_buffer_size: opt_param!(params, "send_buffer_size")?,
+            reuseraddr: opt_param!(params, "reuseaddr")?,
+            reuseport: opt_param!(params, "reuseport")?,
+        })
+    }
+}

src/net/tcp/stream.rs

@@ -0,0 +1,152 @@
+use std::net::SocketAddr;
+use std::ops::{Deref, DerefMut};
+use std::result::Result as StdResult;
+
+use mlua::{Lua, Result, String as LuaString, Table, UserData, UserDataMethods, UserDataRegistry};
+use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
+use tokio::net::lookup_host;
+use tokio::time::timeout;
+
+use super::{SocketOptions, TcpSocket};
+use crate::time::Duration;
+
+pub struct TcpStream {
+    stream: tokio::net::TcpStream,
+    read_timeout: Option<Duration>,
+    write_timeout: Option<Duration>,
+}
+
+impl Deref for TcpStream {
+    type Target = tokio::net::TcpStream;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.stream
+    }
+}
+
+impl DerefMut for TcpStream {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.stream
+    }
+}
+
+macro_rules! with_io_timeout {
+    ($timeout:expr, $fut:expr) => {
+        match $timeout {
+            Some(dur) => (timeout(dur.0, $fut).await)
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::TimedOut, e))
+                .flatten(),
+            None => $fut.await,
+        }
+    };
+}
+
+impl From<tokio::net::TcpStream> for TcpStream {
+    fn from(stream: tokio::net::TcpStream) -> Self {
+        TcpStream {
+            stream,
+            read_timeout: None,
+            write_timeout: None,
+        }
+    }
+}
+
+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_mut("set_read_timeout", |_, this, dur: Option<Duration>| {
+            this.read_timeout = dur;
+            Ok(())
+        });
+
+        registry.add_method_mut("set_write_timeout", |_, this, dur: Option<Duration>| {
+            this.write_timeout = dur;
+            Ok(())
+        });
+
+        registry.add_async_method_mut("read", |lua, mut this, size: usize| async move {
+            let mut buf = vec![0; size];
+            let n = with_io_timeout!(this.read_timeout, this.read(&mut buf));
+            let n = lua_try!(n);
+            buf.truncate(n);
+            Ok(Ok(lua.create_string(buf)?))
+        });
+
+        registry.add_async_method_mut("read_to_end", |lua, mut this, ()| async move {
+            let mut buf = Vec::new();
+            let n = with_io_timeout!(this.read_timeout, this.read_to_end(&mut buf));
+            let _n = lua_try!(n);
+            Ok(Ok(lua.create_string(buf)?))
+        });
+
+        registry.add_async_method_mut("write", |_, mut this, data: LuaString| async move {
+            let n = with_io_timeout!(this.write_timeout, this.write(&data.as_bytes()));
+            let n = lua_try!(n);
+            Ok(Ok(n))
+        });
+
+        registry.add_async_method_mut("write_all", |_, mut this, data: LuaString| async move {
+            let r = with_io_timeout!(this.write_timeout, this.write_all(&data.as_bytes()));
+            lua_try!(r);
+            Ok(Ok(true))
+        });
+
+        registry.add_async_method_mut("flush", |_, mut this, ()| async move {
+            let r = with_io_timeout!(this.write_timeout, this.flush());
+            lua_try!(r);
+            Ok(Ok(true))
+        });
+
+        registry.add_async_method_mut("shutdown", |_, mut this, ()| async move {
+            lua_try!(this.shutdown().await);
+            Ok(Ok(true))
+        });
+    }
+}
+
+pub async fn connect(
+    _: Lua,
+    (addr, params): (String, Option<Table>),
+) -> Result<StdResult<TcpStream, String>> {
+    let addrs = lua_try!(lookup_host(addr).await);
+    let options = SocketOptions::from_table(&params)?;
+
+    let timeout = opt_param!(Duration, params, "timeout")?; // A single timeout for any operation
+    let connect_timeout = opt_param!(Duration, params, "connect_timeout")?.or(timeout);
+    let read_timeout = opt_param!(Duration, params, "read_timeout")?.or(timeout);
+    let write_timeout = opt_param!(Duration, params, "write_timeout")?.or(timeout);
+
+    let try_connect = |addr: SocketAddr| async move {
+        let sock = TcpSocket::new_for_addr(addr)?;
+        sock.set_options(options)?;
+        with_io_timeout!(connect_timeout, sock.0.connect(addr))
+    };
+
+    let mut last_err = None;
+    for addr in addrs {
+        match try_connect(addr).await {
+            Ok(stream) => {
+                return Ok(Ok(TcpStream {
+                    stream,
+                    read_timeout,
+                    write_timeout,
+                }));
+            }
+            Err(e) => {
+                last_err = Some(e);
+                continue;
+            }
+        }
+    }
+
+    Ok(Err(last_err.map(|err| err.to_string()).unwrap_or_else(|| {
+        "could not resolve to any address".to_string()
+    })))
+}

tests/lib.rs

@@ -21,6 +21,11 @@ async fn run_file(modname: &str) -> Result<()> {
     mlua_stdlib::regex::register(&lua, None)?;
     #[cfg(feature = "http")]
     mlua_stdlib::http::register(&lua, None)?;
+    #[cfg(feature = "net")]
+    {
+        mlua_stdlib::net::register(&lua, None)?;
+        mlua_stdlib::net::tcp::register(&lua, None)?;
+    }
     #[cfg(feature = "task")]
     mlua_stdlib::task::register(&lua, None)?;
 
@@ -91,6 +96,11 @@ include_tests! {
         headers,
     },
 
+    #[cfg(feature = "net")]
+    net {
+        tcp,
+    },
+
     #[cfg(feature = "task")]
     task,
 }