From 016d7595121b06999908778b155759e6925c6c64 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Mon, 26 Aug 2024 12:59:57 +0100
Subject: [PATCH] [Feature] add umask option

---
 Cargo.lock               |  1 +
 crates/rqbit/Cargo.toml  |  1 +
 crates/rqbit/src/main.rs | 59 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index ed8e74f..918d812 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2522,6 +2522,7 @@ dependencies = [
  "clap_complete",
  "console-subscriber",
  "futures",
+ "libc",
  "librqbit",
  "openssl",
  "parking_lot",
diff --git a/crates/rqbit/Cargo.toml b/crates/rqbit/Cargo.toml
index d8eb602..2adb816 100644
--- a/crates/rqbit/Cargo.toml
+++ b/crates/rqbit/Cargo.toml
@@ -45,6 +45,7 @@ size_format = "1"
 bytes = "1.5.0"
 openssl = { version = "0.10", features = ["vendored"], optional = true }
 upnp-serve = { path = "../upnp-serve" }
+libc = "0.2.158"
 
 [dev-dependencies]
 futures = { version = "0.3" }
diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs
index ead5efa..f681c2e 100644
--- a/crates/rqbit/src/main.rs
+++ b/crates/rqbit/src/main.rs
@@ -28,6 +28,31 @@ enum LogLevel {
     Error,
 }
 
+fn parse_umask(value: &str) -> anyhow::Result<libc::mode_t> {
+    fn parse_oct_digit(d: u8) -> Option<libc::mode_t> {
+        Some(match d {
+            b'0' => 0,
+            b'1' => 1,
+            b'2' => 2,
+            b'3' => 3,
+            b'4' => 4,
+            b'5' => 5,
+            b'6' => 6,
+            b'7' => 7,
+            _ => return None,
+        })
+    }
+    if value.len() != 3 {
+        bail!("expected 3 digits")
+    }
+    let mut output = 0;
+    for digit in value.as_bytes() {
+        let digit = parse_oct_digit(*digit).context("expected 3 digits")?;
+        output = output * 8 + digit;
+    }
+    Ok(output)
+}
+
 #[derive(Parser)]
 #[command(version, author, about)]
 struct Opts {
@@ -162,6 +187,15 @@ struct Opts {
     /// How many torrents can be initializing (rehashing) at the same time
     #[arg(long, default_value = "5", env = "RQBIT_CONCURRENT_INIT_LIMIT")]
     concurrent_init_limit: usize,
+
+    /// Set the process umask to this value.
+    ///
+    /// Default is inherited from your environment (usually 022).
+    /// This will affect the file mode of created files.
+    ///
+    /// Read more at https://man7.org/linux/man-pages/man2/umask.2.html
+    #[arg(long, env = "RQBIT_UMASK", value_parser=parse_umask)]
+    umask: Option<libc::mode_t>,
 }
 
 #[derive(Parser)]
@@ -294,6 +328,10 @@ fn _start_deadlock_detector_thread() {
 fn main() -> anyhow::Result<()> {
     let opts = Opts::parse();
 
+    if let Some(umask) = opts.umask {
+        unsafe { libc::umask(umask) };
+    }
+
     let mut rt_builder = match opts.single_thread_runtime {
         true => tokio::runtime::Builder::new_current_thread(),
         false => {
@@ -714,3 +752,24 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::parse_umask;
+
+    #[test]
+    fn test_parse_umask() {
+        let range = b'0'..=b'7';
+        for d0 in range.clone() {
+            for d1 in range.clone() {
+                for d2 in range.clone() {
+                    let inp = [d0, d1, d2];
+                    let inp_str = std::str::from_utf8(&inp).unwrap();
+                    let parsed = parse_umask(inp_str).expect(inp_str);
+                    let expected = format!("{parsed:03o}");
+                    assert_eq!(inp_str, expected);
+                }
+            }
+        }
+    }
+}