Improve privilege dropping; run cargo fmt

aquatic_udp/src/lib.rs

@@ -5,13 +5,12 @@ pub mod workers;
 use config::Config;
 
 use std::collections::BTreeMap;
-use std::sync::{atomic::AtomicUsize, Arc};
 use std::thread::Builder;
 
 use anyhow::Context;
 #[cfg(feature = "cpu-pinning")]
 use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
-use aquatic_common::privileges::drop_privileges_after_socket_binding;
+use aquatic_common::privileges::PrivilegeDropper;
 use crossbeam_channel::{bounded, unbounded};
 
 use aquatic_common::access_list::update_access_list;
@@ -32,7 +31,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
 
     let mut signals = Signals::new(::std::iter::once(SIGUSR1))?;
 
-    let num_bound_sockets = Arc::new(AtomicUsize::new(0));
+    let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers);
 
     let mut request_senders = Vec::new();
     let mut request_receivers = BTreeMap::new();
@@ -96,7 +95,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
         let request_sender =
             ConnectedRequestSender::new(SocketWorkerIndex(i), request_senders.clone());
         let response_receiver = response_receivers.remove(&i).unwrap();
-        let num_bound_sockets = num_bound_sockets.clone();
+        let priv_dropper = priv_dropper.clone();
 
         Builder::new()
             .name(format!("socket-{:02}", i + 1))
@@ -115,7 +114,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
                     i,
                     request_sender,
                     response_receiver,
-                    num_bound_sockets,
+                    priv_dropper,
                 );
             })
             .with_context(|| "spawn socket worker")?;
@@ -141,13 +140,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
             .with_context(|| "spawn statistics worker")?;
     }
 
-    drop_privileges_after_socket_binding(
-        &config.privileges,
-        num_bound_sockets,
-        config.socket_workers,
-    )
-    .unwrap();
-
     #[cfg(feature = "cpu-pinning")]
     pin_current_if_configured_to(
         &config.cpu_pinning,

aquatic_udp/src/workers/socket.rs

@@ -1,12 +1,10 @@
 use std::collections::BTreeMap;
 use std::io::{Cursor, ErrorKind};
-use std::sync::{
-    atomic::{AtomicUsize, Ordering},
-    Arc,
-};
+use std::sync::atomic::Ordering;
 use std::time::{Duration, Instant};
 use std::vec::Drain;
 
+use aquatic_common::privileges::PrivilegeDropper;
 use crossbeam_channel::Receiver;
 use mio::net::UdpSocket;
 use mio::{Events, Interest, Poll, Token};
@@ -157,12 +155,12 @@ pub fn run_socket_worker(
     token_num: usize,
     request_sender: ConnectedRequestSender,
     response_receiver: Receiver<(ConnectedResponse, CanonicalSocketAddr)>,
-    num_bound_sockets: Arc<AtomicUsize>,
+    priv_dropper: PrivilegeDropper,
 ) {
     let mut rng = StdRng::from_entropy();
     let mut buffer = [0u8; MAX_PACKET_SIZE];
 
-    let mut socket = UdpSocket::from_std(create_socket(&config));
+    let mut socket = UdpSocket::from_std(create_socket(&config, priv_dropper));
     let mut poll = Poll::new().expect("create poll");
 
     let interests = Interest::READABLE;
@@ -171,8 +169,6 @@ pub fn run_socket_worker(
         .register(&mut socket, Token(token_num), interests)
         .unwrap();
 
-    num_bound_sockets.fetch_add(1, Ordering::SeqCst);
-
     let mut events = Events::with_capacity(config.network.poll_event_capacity);
     let mut connections = ConnectionMap::default();
     let mut pending_scrape_responses = PendingScrapeResponseSlab::default();
@@ -520,7 +516,7 @@ fn send_response(
     }
 }
 
-pub fn create_socket(config: &Config) -> ::std::net::UdpSocket {
+pub fn create_socket(config: &Config, priv_dropper: PrivilegeDropper) -> ::std::net::UdpSocket {
     let socket = if config.network.address.is_ipv4() {
         Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))
     } else {
@@ -542,6 +538,8 @@ pub fn create_socket(config: &Config) -> ::std::net::UdpSocket {
         .bind(&config.network.address.into())
         .unwrap_or_else(|err| panic!("socket: bind to {}: {:?}", config.network.address, err));
 
+    priv_dropper.after_socket_creation();
+
     let recv_buffer_size = config.network.socket_recv_buffer_size;
 
     if recv_buffer_size != 0 {