From fabbf742062a2b47283ede4d360192e3cfacc83b Mon Sep 17 00:00:00 2001
From: TQ Hirsch <thequux@thequux.com>
Date: Sun, 19 Nov 2023 01:30:39 +0100
Subject: [PATCH] Implemented privilege dropping

---
 Cargo.lock  | 11 +++++++++++
 Cargo.toml  |  3 ++-
 src/main.rs | 48 ++++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 297b20a..f031a65 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1288,6 +1288,7 @@ dependencies = [
  "structopt",
  "tokio",
  "tokio-util",
+ "users",
 ]
 
 [[package]]
@@ -1550,6 +1551,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"
diff --git a/Cargo.toml b/Cargo.toml
index dacce0c..5479102 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,4 +18,5 @@ 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"
\ No newline at end of file
+libc = "0.2.150"
+users = "0.11.0"
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index fb384e0..20b03f4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -35,6 +35,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()?;
@@ -106,15 +126,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>(())
     };