Implemented script reloading

Cargo.lock

@@ -1242,6 +1242,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"
@@ -1288,6 +1297,7 @@ dependencies = [
  "structopt",
  "tokio",
  "tokio-util",
+ "users",
 ]
 
 [[package]]
@@ -1432,6 +1442,7 @@ dependencies = [
  "mio",
  "num_cpus",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2 0.5.5",
  "tokio-macros",
  "windows-sys",
@@ -1550,6 +1561,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"

Cargo.toml

@@ -13,9 +13,10 @@ 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"

src/engine.rs

@@ -1,15 +1,9 @@
 use std::future::Future;
 use mlua::{FromLua, IntoLua, Lua, UserData, UserDataFields, Value};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
-use std::path::{Path, PathBuf};
+use std::path::PathBuf;
 use anyhow::anyhow;
 use async_tftp::packet::Error;
-use futures::TryFutureExt;
-use tokio::sync::oneshot::error::RecvError;
-
-pub struct LuaRunner {
-
-}
 
 #[derive(Clone, Debug)]
 pub struct Client {
@@ -181,7 +175,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 +194,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 {

src/handler.rs

@@ -1,15 +1,12 @@
 // 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 reqwest::{Body, StatusCode};
 use tokio_util::compat::TokioAsyncWriteCompatExt;
 use crate::engine::{Client, Engine, Resource};
 
@@ -50,9 +47,16 @@ impl async_tftp::server::Handler for Handler {
         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();
 

src/main.rs

@@ -5,6 +5,7 @@ use std::net::{Ipv4Addr, Ipv6Addr};
 use std::path::{Path, PathBuf};
 use anyhow::anyhow;
 use structopt::StructOpt;
+use tokio::signal::unix::SignalKind;
 
 mod handler;
 mod engine;
@@ -35,6 +36,26 @@ struct Options {
 async fn main() -> anyhow::Result<()> {
     let opts = Options::from_args();
 
+    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 +63,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 +136,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>(())
     };