Compare commits
8 Commits
68b656acfb
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| d014f6e2d2 | |||
| 30b479534c | |||
| 35ad2063c6 | |||
| 253e9d90ce | |||
| 559bd03de3 | |||
| fabbf74206 | |||
| c638acecc2 | |||
| e693909953 |
2
.gitmodules
vendored
2
.gitmodules
vendored
@@ -1,3 +1,3 @@
|
||||
[submodule "vendor/async-tftp-rs"]
|
||||
path = vendor/async-tftp-rs
|
||||
url = git@github.com:thequux/async-tftp-rs.git
|
||||
url = https://github.com/thequux/async-tftp-rs.git
|
||||
|
||||
28
Cargo.lock
generated
28
Cargo.lock
generated
@@ -635,6 +635,12 @@ version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
|
||||
|
||||
[[package]]
|
||||
name = "humantime"
|
||||
version = "2.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
|
||||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "0.14.27"
|
||||
@@ -1242,6 +1248,15 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "signal-hook-registry"
|
||||
version = "1.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "slab"
|
||||
version = "0.4.9"
|
||||
@@ -1280,6 +1295,7 @@ dependencies = [
|
||||
"async-tftp",
|
||||
"fern",
|
||||
"futures",
|
||||
"humantime",
|
||||
"libc",
|
||||
"listenfd",
|
||||
"log",
|
||||
@@ -1288,6 +1304,7 @@ dependencies = [
|
||||
"structopt",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"users",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1432,6 +1449,7 @@ dependencies = [
|
||||
"mio",
|
||||
"num_cpus",
|
||||
"pin-project-lite",
|
||||
"signal-hook-registry",
|
||||
"socket2 0.5.5",
|
||||
"tokio-macros",
|
||||
"windows-sys",
|
||||
@@ -1550,6 +1568,16 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "users"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24cc0f6d6f267b73e5a2cadf007ba8f9bc39c6a6f9666f8cf25ea809a153b032"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.5.0"
|
||||
|
||||
@@ -13,9 +13,11 @@ anyhow = "1"
|
||||
fern = "0.6.2"
|
||||
log = "0.4.20"
|
||||
structopt = "0.3"
|
||||
tokio = { version = "1.32.0", features = ["rt-multi-thread", "macros", "rt"] }
|
||||
tokio = { version = "1.32.0", features = ["rt-multi-thread", "macros", "rt", "signal"] }
|
||||
tokio-util = { version = "0.7.10", features = ["io", "io-util", "rt", "compat"] }
|
||||
mlua = { version = "0.9.1", features = ["luau-jit", "vendored", "async", "send"] }
|
||||
reqwest = { version = "0.11.22", features = ["stream"] }
|
||||
listenfd = "1.0.1"
|
||||
libc = "0.2.150"
|
||||
libc = "0.2.150"
|
||||
users = "0.11.0"
|
||||
humantime = "2.1.0"
|
||||
@@ -5,7 +5,7 @@
|
||||
#env.GREET = "devenv";
|
||||
|
||||
# https://devenv.sh/packages/
|
||||
packages = [ pkgs.git pkgs.openssl ];
|
||||
packages = [ pkgs.git pkgs.openssl pkgs.lua ];
|
||||
|
||||
# https://devenv.sh/scripts/
|
||||
#scripts.hello.exec = "echo hello from $GREET";
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::future::Future;
|
||||
use mlua::{FromLua, IntoLua, Lua, UserData, UserDataFields, Value};
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::path::{Path, PathBuf};
|
||||
use mlua::{FromLua, FromLuaMulti, IntoLua, Lua, UserData, UserDataFields, UserDataMethods, Value};
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
use std::path::PathBuf;
|
||||
use std::str::FromStr;
|
||||
use anyhow::anyhow;
|
||||
use async_tftp::packet::Error;
|
||||
use futures::TryFutureExt;
|
||||
use tokio::sync::oneshot::error::RecvError;
|
||||
|
||||
pub struct LuaRunner {
|
||||
|
||||
}
|
||||
mod ipaddr;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Client {
|
||||
@@ -32,6 +30,17 @@ pub enum Resource {
|
||||
Error(Error)
|
||||
}
|
||||
|
||||
impl Display for Resource {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Resource::Http(url) => write!(f, "HTTP {url}"),
|
||||
Resource::File(path) => write!(f, "FILE {path}"),
|
||||
Resource::Data(data) => write!(f, "DATA ({} bytes)", data.len()),
|
||||
Resource::Error(err) => write!(f, "ERR {err:?}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Resource {
|
||||
}
|
||||
|
||||
@@ -40,7 +49,7 @@ impl UserData for Resource {
|
||||
}
|
||||
|
||||
impl<'lua> FromLua<'lua> for Resource {
|
||||
fn from_lua(value: Value<'lua>, _lua: &'lua Lua) -> mlua::Result<Self> {
|
||||
fn from_lua(value: Value<'lua>, _: &'lua Lua) -> mlua::Result<Self> {
|
||||
value.as_userdata().ok_or(mlua::Error::UserDataTypeMismatch).and_then(|value| value.take())
|
||||
}
|
||||
}
|
||||
@@ -73,16 +82,12 @@ pub struct EngineImpl {
|
||||
chan: Option<tokio::sync::mpsc::Receiver<EngineReq>>,
|
||||
}
|
||||
|
||||
|
||||
impl EngineImpl {
|
||||
pub(crate) fn init(&mut self) -> anyhow::Result<()> {
|
||||
let lua = &* self.lua;
|
||||
let lua = self.lua;
|
||||
lua.load_from_std_lib(mlua::StdLib::ALL)?;
|
||||
|
||||
|
||||
lua.register_userdata_type::<IpAddr>(|registry| {
|
||||
registry.add_field_method_get("version", |_, ip| Ok(if ip.is_ipv4() { 4} else {6}));
|
||||
})?;
|
||||
|
||||
{
|
||||
// prepare resource types...
|
||||
let resources = lua.create_table()?;
|
||||
@@ -104,11 +109,16 @@ impl EngineImpl {
|
||||
err_tbl.set("FileAlreadyExists", Resource::Error(Error::FileAlreadyExists))?;
|
||||
err_tbl.set("NoSuchUser", Resource::Error(Error::NoSuchUser))?;
|
||||
err_tbl.set("Message", err_fn)?;
|
||||
err_tbl.set_readonly(true);
|
||||
|
||||
resources.set("ERROR", err_tbl)?;
|
||||
resources.set_readonly(true);
|
||||
|
||||
lua.globals().set("resource", resources)?;
|
||||
lua.globals().set("state", lua.create_table()?)?;
|
||||
|
||||
// Construct data types
|
||||
ipaddr::register(lua)?;
|
||||
//
|
||||
}
|
||||
|
||||
@@ -181,7 +191,7 @@ impl EngineImpl {
|
||||
.into_lua(lua)
|
||||
.map_err(|err| Error::Msg(err.to_string()))?;
|
||||
let path = path.to_str().ok_or(Error::FileNotFound)?.to_owned();
|
||||
let (resource, size): (Resource, Option<u64>) = resolver
|
||||
let resource: Resource = resolver
|
||||
.call_async((path, lua_client.clone(), size))
|
||||
.await
|
||||
.map_err(|err| Error::Msg(err.to_string()))?;
|
||||
@@ -200,8 +210,6 @@ impl Engine {
|
||||
pub fn new() -> anyhow::Result<(Self, EngineImpl)> {
|
||||
let lua = Box::leak(Box::new(mlua::Lua::new()));
|
||||
// Add stdlib
|
||||
let handler_fn = lua.create_registry_value(0)?;
|
||||
|
||||
let (req_snd, req_rcv) = tokio::sync::mpsc::channel(1);
|
||||
|
||||
let engine = Self {
|
||||
|
||||
238
src/engine/ipaddr.rs
Normal file
238
src/engine/ipaddr.rs
Normal file
@@ -0,0 +1,238 @@
|
||||
use std::fmt::Display;
|
||||
use std::net::{IpAddr, Ipv6Addr};
|
||||
use std::str::FromStr;
|
||||
use std::fmt::Formatter;
|
||||
use mlua::{FromLua, FromLuaMulti, Lua, UserData, UserDataFields, UserDataMethods, Value};
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
struct Cidr {
|
||||
addr: IpAddr,
|
||||
prefix: u8,
|
||||
}
|
||||
|
||||
impl Display for Cidr {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}/{}", self.addr, self.prefix)
|
||||
}
|
||||
}
|
||||
|
||||
impl Cidr {
|
||||
fn to_v6(&self) -> Self {
|
||||
match self.addr {
|
||||
IpAddr::V4(ip) => Cidr {
|
||||
addr: IpAddr::V6(ip.to_ipv6_mapped()),
|
||||
prefix: self.prefix + 96,
|
||||
},
|
||||
IpAddr::V6(_) => *self,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[repr(transparent)]
|
||||
#[derive(Copy, Clone)]
|
||||
struct IpAddrWrapper(IpAddr);
|
||||
|
||||
impl IpAddrWrapper {
|
||||
fn octet_len(&self) -> usize {
|
||||
match &self.0 {
|
||||
IpAddr::V4(_) => 4,
|
||||
IpAddr::V6(_) => 16,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_v6(&self) -> Ipv6Addr {
|
||||
match self.0 {
|
||||
IpAddr::V4(v4) => v4.to_ipv6_mapped(),
|
||||
IpAddr::V6(v6) => v6
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UserData for IpAddrWrapper {
|
||||
fn add_fields<'lua, F: UserDataFields<'lua, Self>>(fields: &mut F) {
|
||||
fields.add_field_method_get("version", |_, ip| Ok(if ip.0.is_ipv4() { 4} else {6}));
|
||||
fields.add_field_method_get("is_ipv4", |_, ip| Ok(ip.0.is_ipv4()));
|
||||
fields.add_field_method_get("is_ipv6", |_, ip| Ok(ip.0.is_ipv6()));
|
||||
fields.add_field_method_get("bytes", |lua, ip| match &ip.0 {
|
||||
IpAddr::V4(v4) => lua.create_string(v4.octets()),
|
||||
IpAddr::V6(v6) => lua.create_string(v6.octets()),
|
||||
});
|
||||
fields.add_meta_field("__name", "IpAddr");
|
||||
}
|
||||
|
||||
fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
|
||||
methods.add_method("to_v6", |_, ip, ()| match &ip.0 {
|
||||
IpAddr::V4(v4) => Ok(IpAddrWrapper(IpAddr::V6(v4.to_ipv6_mapped()))),
|
||||
IpAddr::V6(v6) => Ok(IpAddrWrapper(IpAddr::V6(*v6))),
|
||||
});
|
||||
methods.add_meta_method("__index", |_, ip, index: usize| match &ip.0 {
|
||||
IpAddr::V4(ip) if 1 <= index && index <= 4 => Ok(ip.octets()[index-1]),
|
||||
IpAddr::V6(ip) if 1 <= index && index <= 16 => Ok(ip.octets()[index-1]),
|
||||
_ => Err(mlua::Error::runtime("Index out of range")),
|
||||
});
|
||||
methods.add_meta_method_mut("__newindex", |_, ip, (index, value): (usize, u8)| {
|
||||
match &mut ip.0 {
|
||||
IpAddr::V4(ip) if 1 <= index && index <= 4 => {
|
||||
let mut octets = ip.octets();
|
||||
octets[index-1] = value;
|
||||
IpAddr::V4(octets.into());
|
||||
},
|
||||
IpAddr::V6(ip) if 1 <= index && index <= 16 => {
|
||||
let mut octets = ip.octets();
|
||||
octets[index-1] = value;
|
||||
IpAddr::V6(octets.into());
|
||||
},
|
||||
_ => return Err(mlua::Error::runtime("Index out of range")),
|
||||
};
|
||||
Ok(())
|
||||
});
|
||||
methods.add_meta_method("__tostring", |_, ip, ()| Ok(ip.0.to_string()));
|
||||
methods.add_meta_function("__call", |lua, args: mlua::MultiValue| Ok(IpAddrWrapper({
|
||||
eprintln!("IpAddr __call Received {} args", args.len());
|
||||
if args.len() == 4 {
|
||||
// IPV4 direct bytes
|
||||
let (a,b,c,d) = <(u8,u8,u8,u8)>::from_lua_args(args, 0, Some("sftpd.IpAddr"), lua)?;
|
||||
IpAddr::V4([a,b,c,d].into())
|
||||
} else if args.len() == 16 {
|
||||
let mut octets = [0u8;16];
|
||||
for i in 0..16 {
|
||||
octets[i] = u8::from_lua(args.get(i).unwrap().clone(), lua)?;
|
||||
}
|
||||
IpAddr::V6(octets.into())
|
||||
} else if args.len() == 1 {
|
||||
return IpAddrWrapper::from_lua(args[0].clone(), lua);
|
||||
} else {
|
||||
return Err(mlua::Error::runtime("Invalid arguments to stfptd.IpAddr"))
|
||||
}
|
||||
})));
|
||||
methods.add_meta_method("__eq", |_, me, other: IpAddrWrapper| Ok(me.to_v6() == other.to_v6()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'lua> FromLua<'lua> for IpAddrWrapper {
|
||||
fn from_lua(value: Value<'lua>, _: &'lua Lua) -> mlua::Result<Self> {
|
||||
return if let Some(ud) = value.as_userdata() {
|
||||
Ok(ud.borrow::<IpAddrWrapper>()?.clone())
|
||||
} else if let Some(s) = value.as_str() {
|
||||
IpAddr::from_str(s)
|
||||
.map(IpAddrWrapper)
|
||||
.map_err(mlua::Error::external)
|
||||
} else {
|
||||
Err(mlua::Error::FromLuaConversionError {
|
||||
from: value.type_name(),
|
||||
to: "IpAddr",
|
||||
message: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UserData for Cidr {
|
||||
fn add_fields<'lua, F: UserDataFields<'lua, Self>>(fields: &mut F) {
|
||||
fields.add_field_method_get("addr", |_lua, cidr| Ok(IpAddrWrapper(cidr.addr)));
|
||||
fields.add_field_method_set("addr", |_lua, cidr, value: IpAddrWrapper| {
|
||||
if cidr.addr.is_ipv4() ^ value.0.is_ipv4() {
|
||||
return Err(mlua::Error::runtime(format!(
|
||||
"Cannot assign v{} addr to v{} CIDR",
|
||||
if value.0.is_ipv4() { 4 } else { 6 },
|
||||
if cidr.addr.is_ipv4() { 4 } else { 6 },
|
||||
)));
|
||||
}
|
||||
cidr.addr = value.0;
|
||||
Ok(())
|
||||
});
|
||||
fields.add_field_method_get("prefix", |_lua, cidr| Ok(cidr.prefix));
|
||||
fields.add_field_method_set("prefix", |_lua, cidr, value: u8| {
|
||||
if value >= if cidr.addr.is_ipv4() { 32 } else { 128 } {
|
||||
return Err(mlua::Error::runtime("Invalid prefix length"));
|
||||
}
|
||||
cidr.prefix = value;
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
|
||||
fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
|
||||
methods.add_method("contains", |_, cidr, addr: IpAddrWrapper| {
|
||||
let addr = u128::from_be_bytes(addr.to_v6().octets());
|
||||
let cidr = cidr.to_v6();
|
||||
if cidr.prefix == 0 {
|
||||
// Global network
|
||||
return Ok(true)
|
||||
}
|
||||
let mask = (!0u128) << (128 - cidr.prefix);
|
||||
let network = u128::from_be_bytes(IpAddrWrapper(cidr.addr).to_v6().octets());
|
||||
|
||||
Ok(addr & mask == network & mask)
|
||||
});
|
||||
|
||||
methods.add_meta_method("__tostring", |_, cidr, ()| Ok(cidr.to_string()));
|
||||
methods.add_meta_method("__eq", |_, cidr, other: Cidr| {
|
||||
let cidr = cidr.to_v6();
|
||||
let other = other.to_v6();
|
||||
Ok(cidr.prefix == other.prefix && cidr.addr == other.addr)
|
||||
});
|
||||
|
||||
methods.add_meta_function("__call", |lua, args: mlua::MultiValue| {
|
||||
if args.len() == 0 {
|
||||
Ok(Cidr{addr: IpAddr::V6(Ipv6Addr::UNSPECIFIED), prefix: 0})
|
||||
} else if args.len() == 1 {
|
||||
Cidr::from_lua_args(args, 0, Some("stftpd.Cidr"), lua)
|
||||
} else {
|
||||
let (IpAddrWrapper(addr), prefix) = <(IpAddrWrapper, u8)>::from_lua_args(args, 0, Some("stftpd.Cidr"), lua)?;
|
||||
let max_prefix = if addr.is_ipv4() { 32 } else { 64 };
|
||||
if prefix > max_prefix {
|
||||
return Err(mlua::Error::runtime("Invalid prefix"));
|
||||
}
|
||||
Ok(Cidr{addr, prefix})
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl<'lua> FromLua<'lua> for Cidr {
|
||||
fn from_lua(value: Value<'lua>, _: &'lua Lua) -> mlua::Result<Self> {
|
||||
return if let Some(ud) = value.as_userdata() {
|
||||
Ok(ud.borrow::<Cidr>()?.clone())
|
||||
} else if let Some(s) = value.as_str() {
|
||||
let (addr, prefix) = s.split_once("/")
|
||||
.map(|(a,p)| (a,Some(p)))
|
||||
.unwrap_or((s, None));
|
||||
let addr = IpAddr::from_str(addr)?;
|
||||
let prefix = prefix.map(u8::from_str)
|
||||
.transpose()
|
||||
.map_err(mlua::Error::external)?;
|
||||
|
||||
let max_prefix = if addr.is_ipv4() { 32 } else { 128 };
|
||||
let prefix = prefix.unwrap_or(max_prefix);
|
||||
if prefix > max_prefix {
|
||||
return Err(mlua::Error::runtime("CIDR prefix to long"))
|
||||
}
|
||||
|
||||
Ok(Cidr {
|
||||
addr, prefix
|
||||
})
|
||||
} else {
|
||||
Err(mlua::Error::FromLuaConversionError {
|
||||
from: value.type_name(),
|
||||
to: "Cidr",
|
||||
message: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn register(lua: &'static Lua) -> anyhow::Result<()> {
|
||||
let globals = lua.globals();
|
||||
let stftpd: Value = globals.get("stftpd")?;
|
||||
let stftpd = if stftpd.is_nil() {
|
||||
let newtab = lua.create_table()?;
|
||||
globals.set("stftpd", newtab.clone())?;
|
||||
newtab
|
||||
} else {
|
||||
stftpd.as_table().unwrap().clone()
|
||||
};
|
||||
stftpd.set("Cidr", lua.create_proxy::<Cidr>()?)?;
|
||||
stftpd.set("IpAddr", lua.create_proxy::<IpAddrWrapper>()?)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,15 +1,13 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::future::Future;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::net::SocketAddr;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use async_lock::Mutex;
|
||||
use async_tftp::async_trait;
|
||||
use async_tftp::packet::Error;
|
||||
use async_tftp::server::handlers::{DirHandler, DirHandlerMode};
|
||||
use futures::{AsyncRead, AsyncWrite, TryStreamExt};
|
||||
use mlua::{FromLua, UserDataFields};
|
||||
use reqwest::Body;
|
||||
use log::info;
|
||||
use reqwest::{Body, StatusCode};
|
||||
use tokio_util::compat::TokioAsyncWriteCompatExt;
|
||||
use crate::engine::{Client, Engine, Resource};
|
||||
|
||||
@@ -47,12 +45,21 @@ impl async_tftp::server::Handler for Handler {
|
||||
// .to_str().ok_or(Error::FileNotFound)?.to_owned()
|
||||
let resource: Resource = self.engine.resolve(path.to_owned(), &mut lua_client, None).await?;
|
||||
|
||||
info!("GET {path:?} from {client:?} -> {resource}");
|
||||
|
||||
match resource {
|
||||
Resource::Http(url) => {
|
||||
// TODO: Add headers describing client
|
||||
let req = self.http.get(url).send().await.map_err(|err| Error::Msg(err.to_string()))?;
|
||||
let size = req.content_length();
|
||||
let stream = req.bytes_stream()
|
||||
let resp = self.http.get(url).send().await.map_err(|err| Error::Msg(err.to_string()))?;
|
||||
match resp.status() {
|
||||
StatusCode::OK => {},
|
||||
StatusCode::NOT_FOUND => return Err(Error::FileNotFound),
|
||||
StatusCode::UNAUTHORIZED => return Err(Error::PermissionDenied),
|
||||
x => return Err(Error::Msg(x.to_string()))
|
||||
}
|
||||
|
||||
let size = resp.content_length();
|
||||
let stream = resp.bytes_stream()
|
||||
.map_err(|e| futures::io::Error::new(futures::io::ErrorKind::Other, e))
|
||||
.into_async_read();
|
||||
|
||||
@@ -79,6 +86,7 @@ impl async_tftp::server::Handler for Handler {
|
||||
for_write: false,
|
||||
};
|
||||
let resource: Resource = self.engine.resolve(path.to_owned(), &mut lua_client, size).await?;
|
||||
info!("PUT {path:?} from {client:?} -> {resource}");
|
||||
|
||||
match resource {
|
||||
Resource::Http(url) => {
|
||||
|
||||
79
src/main.rs
79
src/main.rs
@@ -3,8 +3,10 @@ use std::ffi::c_void;
|
||||
use std::net;
|
||||
use std::net::{Ipv4Addr, Ipv6Addr};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
use anyhow::anyhow;
|
||||
use structopt::StructOpt;
|
||||
use tokio::signal::unix::SignalKind;
|
||||
|
||||
mod handler;
|
||||
mod engine;
|
||||
@@ -35,6 +37,42 @@ struct Options {
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let opts = Options::from_args();
|
||||
|
||||
// Configure logging
|
||||
fern::Dispatch::new()
|
||||
.format(|out, message, record| {
|
||||
out.finish(format_args!(
|
||||
"[{} {} {}] {}",
|
||||
humantime::format_rfc3339_seconds(SystemTime::now()),
|
||||
record.level(),
|
||||
record.target(),
|
||||
message
|
||||
))
|
||||
})
|
||||
.level(log::LevelFilter::Debug)
|
||||
.chain(std::io::stdout())
|
||||
.chain(fern::log_file("output.log")?)
|
||||
.apply()?;
|
||||
|
||||
let group = opts.group.map(|name| {
|
||||
if let Ok(gid) = libc::gid_t::from_str_radix(name.as_str(), 10) {
|
||||
Ok(gid)
|
||||
} else {
|
||||
users::get_group_by_name(name.as_str())
|
||||
.ok_or(anyhow!("Error dropping privileges: No such group"))
|
||||
.map(|group| group.gid())
|
||||
}
|
||||
}).transpose()?;
|
||||
let user = opts.user.map(|name| {
|
||||
if let Ok(uid) = libc::gid_t::from_str_radix(name.as_str(), 10) {
|
||||
Ok(uid)
|
||||
} else {
|
||||
users::get_user_by_name(name.as_str())
|
||||
.ok_or(anyhow!("Error dropping privileges: No such user"))
|
||||
.map(|user| user.uid())
|
||||
}
|
||||
}).transpose()?;
|
||||
|
||||
|
||||
let local_set = tokio::task::LocalSet::new();
|
||||
|
||||
let (engine, mut engine_impl) = engine::Engine::new()?;
|
||||
@@ -42,8 +80,17 @@ async fn main() -> anyhow::Result<()> {
|
||||
engine_impl.load_script(opts.script.clone()).await?;
|
||||
|
||||
local_set.spawn_local(engine_impl.run());
|
||||
|
||||
let mut handler = handler::Handler::new(
|
||||
{
|
||||
let engine = engine.clone();
|
||||
let script = opts.script.clone();
|
||||
let mut hup_stream = tokio::signal::unix::signal(SignalKind::hangup())?;
|
||||
local_set.spawn_local(async move {
|
||||
while let Some(_) = hup_stream.recv().await {
|
||||
engine.load_script(script.clone()).await.ok();
|
||||
}
|
||||
});
|
||||
}
|
||||
let handler = handler::Handler::new(
|
||||
engine.clone(),
|
||||
opts.serve.as_ref()
|
||||
.map(PathBuf::as_path)
|
||||
@@ -106,15 +153,31 @@ async fn main() -> anyhow::Result<()> {
|
||||
} else {
|
||||
net::UdpSocket::bind((host, port))?
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
|
||||
let server = async_tftp::server::TftpServerBuilder::with_handler(handler.clone())
|
||||
.std_socket(sock)?
|
||||
.build().await?;
|
||||
|
||||
// Drop privileges
|
||||
if let Some(group) = group {
|
||||
unsafe {
|
||||
if libc::setgid(group) < 0 {
|
||||
return Err(std::io::Error::last_os_error().into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(uid) = user {
|
||||
unsafe {
|
||||
if libc::setuid(uid) < 0 {
|
||||
return Err(std::io::Error::last_os_error().into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let main_task = async move {
|
||||
|
||||
let server = async_tftp::server::TftpServerBuilder::with_handler(handler.clone())
|
||||
.std_socket(sock)?
|
||||
.build().await?;
|
||||
|
||||
server.serve().await?;
|
||||
Ok::<(), anyhow::Error>(())
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user