111 lines
3.3 KiB
Rust
111 lines
3.3 KiB
Rust
use std::io;
|
|
use std::net::SocketAddr;
|
|
use std::rc::Rc;
|
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
use tokio::task;
|
|
|
|
use mlua::{chunk, Function, Lua, RegistryKey, String as LuaString, UserData, UserDataMethods};
|
|
|
|
struct LuaTcpStream(TcpStream);
|
|
|
|
impl UserData for LuaTcpStream {
|
|
fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
|
|
methods.add_method("peer_addr", |_, this, ()| {
|
|
Ok(this.0.peer_addr()?.to_string())
|
|
});
|
|
|
|
methods.add_async_method_mut("read", |lua, this, size| async move {
|
|
let mut buf = vec![0; size];
|
|
let n = this.0.read(&mut buf).await?;
|
|
buf.truncate(n);
|
|
lua.create_string(&buf)
|
|
});
|
|
|
|
methods.add_async_method_mut("write", |_, this, data: LuaString| async move {
|
|
let n = this.0.write(&data.as_bytes()).await?;
|
|
Ok(n)
|
|
});
|
|
|
|
methods.add_async_method_mut("close", |_, this, ()| async move {
|
|
this.0.shutdown().await?;
|
|
Ok(())
|
|
});
|
|
}
|
|
}
|
|
|
|
async fn run_server(lua: Lua, handler: RegistryKey) -> io::Result<()> {
|
|
let addr: SocketAddr = ([127, 0, 0, 1], 3000).into();
|
|
let listener = TcpListener::bind(addr).await.expect("cannot bind addr");
|
|
|
|
println!("Listening on {}", addr);
|
|
|
|
let lua = Rc::new(lua);
|
|
let handler = Rc::new(handler);
|
|
loop {
|
|
let (stream, _) = match listener.accept().await {
|
|
Ok(res) => res,
|
|
Err(err) if is_transient_error(&err) => continue,
|
|
Err(err) => return Err(err),
|
|
};
|
|
|
|
let lua = lua.clone();
|
|
let handler = handler.clone();
|
|
task::spawn_local(async move {
|
|
let handler: Function = lua
|
|
.registry_value(&handler)
|
|
.expect("cannot get Lua handler");
|
|
|
|
let stream = LuaTcpStream(stream);
|
|
if let Err(err) = handler.call_async::<_, ()>(stream).await {
|
|
eprintln!("{}", err);
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
#[tokio::main(flavor = "current_thread")]
|
|
async fn main() {
|
|
let lua = Lua::new();
|
|
|
|
// Create Lua handler function
|
|
let handler_fn = lua
|
|
.load(chunk! {
|
|
function(stream)
|
|
local peer_addr = stream:peer_addr()
|
|
print("connected from "..peer_addr)
|
|
|
|
while true do
|
|
local data = stream:read(100)
|
|
data = data:match("^%s*(.-)%s*$") // trim
|
|
print("["..peer_addr.."] "..data)
|
|
if data == "bye" then
|
|
stream:write("bye bye\n")
|
|
stream:close()
|
|
return
|
|
end
|
|
stream:write("echo: "..data.."\n")
|
|
end
|
|
end
|
|
})
|
|
.eval::<Function>()
|
|
.expect("cannot create Lua handler");
|
|
|
|
// Store it in the Registry
|
|
let handler = lua
|
|
.create_registry_value(handler_fn)
|
|
.expect("cannot store Lua handler");
|
|
|
|
task::LocalSet::new()
|
|
.run_until(run_server(lua, handler))
|
|
.await
|
|
.expect("cannot run server")
|
|
}
|
|
|
|
fn is_transient_error(e: &io::Error) -> bool {
|
|
e.kind() == io::ErrorKind::ConnectionRefused
|
|
|| e.kind() == io::ErrorKind::ConnectionAborted
|
|
|| e.kind() == io::ErrorKind::ConnectionReset
|
|
}
|