Improve privilege dropping; run cargo fmt

aquatic_common/src/cpu_pinning.rs

@@ -194,7 +194,9 @@ pub mod glommio {
                     // 15 -> 14 and 15
                     // 14 -> 12 and 13
                     // 13 -> 10 and 11
-                    CpuPinningDirection::Descending => num_cpu_cores - 2 * (num_cpu_cores - core_index),
+                    CpuPinningDirection::Descending => {
+                        num_cpu_cores - 2 * (num_cpu_cores - core_index)
+                    }
                 };
 
                 get_cpu_set()?

aquatic_common/src/privileges.rs

@@ -1,22 +1,22 @@
 use std::{
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
-    time::Duration,
+    path::PathBuf,
+    sync::{Arc, Barrier},
 };
 
-use aquatic_toml_config::TomlConfig;
 use privdrop::PrivDrop;
 use serde::Deserialize;
 
+use aquatic_toml_config::TomlConfig;
+
 #[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
 #[serde(default)]
 pub struct PrivilegeConfig {
-    /// Chroot and switch user after binding to sockets
+    /// Chroot and switch group and user after binding to sockets
     pub drop_privileges: bool,
     /// Chroot to this path
-    pub chroot_path: String,
+    pub chroot_path: PathBuf,
+    /// Group to switch to after chrooting
+    pub group: String,
     /// User to switch to after chrooting
     pub user: String,
 }
@@ -25,41 +25,37 @@ impl Default for PrivilegeConfig {
     fn default() -> Self {
         Self {
             drop_privileges: false,
-            chroot_path: ".".to_string(),
+            chroot_path: ".".into(),
             user: "nobody".to_string(),
+            group: "nobody".to_string(),
         }
     }
 }
 
-pub fn drop_privileges_after_socket_binding(
-    config: &PrivilegeConfig,
-    num_bound_sockets: Arc<AtomicUsize>,
-    target_num: usize,
-) -> anyhow::Result<()> {
-    if config.drop_privileges {
-        let mut counter = 0usize;
+#[derive(Clone)]
+pub struct PrivilegeDropper {
+    barrier: Arc<Barrier>,
+    config: Arc<PrivilegeConfig>,
+}
 
-        loop {
-            let num_bound = num_bound_sockets.load(Ordering::SeqCst);
+impl PrivilegeDropper {
+    pub fn new(config: PrivilegeConfig, num_sockets: usize) -> Self {
+        Self {
+            barrier: Arc::new(Barrier::new(num_sockets)),
+            config: Arc::new(config),
+        }
+    }
 
-            if num_bound == target_num {
+    pub fn after_socket_creation(&self) {
+        if self.config.drop_privileges {
+            if self.barrier.wait().is_leader() {
                 PrivDrop::default()
-                    .chroot(config.chroot_path.clone())
-                    .user(config.user.clone())
-                    .apply()?;
-
-                break;
-            }
-
-            ::std::thread::sleep(Duration::from_millis(10));
-
-            counter += 1;
-
-            if counter == 500 {
-                panic!("Sockets didn't bind in time for privilege drop.");
+                    .chroot(self.config.chroot_path.clone())
+                    .user(self.config.user.clone())
+                    .user(self.config.user.clone())
+                    .apply()
+                    .expect("drop privileges");
             }
         }
     }
-
-    Ok(())
 }