diff --git a/Cargo.lock b/Cargo.lock index 6f641d4..1d865db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -208,6 +208,16 @@ dependencies = [ "log", ] +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -386,6 +396,27 @@ dependencies = [ "common", "env_logger", "log", + "signal-hook", +] + +[[package]] +name = "signal-hook" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a0c28ca5908dbdbcd52e6fdaa00358ab88637f8ab33e1f188dd510eb44b53d" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", ] [[package]] diff --git a/server/Cargo.toml b/server/Cargo.toml index e85aa48..b346b76 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -9,7 +9,8 @@ common = { workspace = true } argon2 = "0.5.3" env_logger = "0.11.8" log = "0.4.29" +signal-hook = "0.4" [[bin]] name = "saslauthd" -path = "src/main.rs" \ No newline at end of file +path = "src/main.rs" diff --git a/server/src/server.rs b/server/src/server.rs index cc7a2e0..3ad1c00 100644 --- a/server/src/server.rs +++ b/server/src/server.rs @@ -1,17 +1,30 @@ use crate::handler::Handler; use crate::repository::PasswordDirectory; +use log::{info, warn}; +use signal_hook::consts::{SIGINT, SIGTERM}; +use signal_hook::flag; use std::fs::{metadata, set_permissions}; -use std::io::Error; +use std::io::{Error, ErrorKind}; use std::os::unix::fs::PermissionsExt; use std::os::unix::net::UnixListener; use std::path::PathBuf; -use std::thread; -use log::info; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +/// How long to sleep between polls of the non-blocking listener while idle. +const ACCEPT_POLL_INTERVAL: Duration = Duration::from_millis(50); + +/// How long to wait for in-flight client handlers to finish on shutdown. +const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(3); pub struct Server { listener: UnixListener, repository: PasswordDirectory, socket_name: PathBuf, + /// Set to the signal number (SIGTERM / SIGINT) once one has been received; 0 otherwise. + shutdown_signal: Arc, } impl Server { @@ -28,35 +41,102 @@ impl Server { perms.set_mode(0o0777); set_permissions(&socket_name, perms)?; + // Non-blocking so the accept loop can notice a shutdown request. + listener.set_nonblocking(true)?; + + let shutdown_signal = Arc::new(AtomicUsize::new(0)); + flag::register_usize(SIGTERM, Arc::clone(&shutdown_signal), SIGTERM as usize)?; + flag::register_usize(SIGINT, Arc::clone(&shutdown_signal), SIGINT as usize)?; + Ok(Server { repository, listener, socket_name, + shutdown_signal, }) } pub fn run(&mut self) -> Result<(), Error> { - // accept connections and process them, spawning a new thread for each one - for stream in self.listener.incoming() { - match stream { - Ok(stream) => { + let mut workers: Vec> = Vec::new(); + + // accept connections and process them, spawning a new thread for each one, + // until a shutdown signal arrives + loop { + let signal = self.shutdown_signal.load(Ordering::Relaxed); + if signal != 0 { + info!("received {}, shutting down", signal_name(signal)); + break; + } + + match self.listener.accept() { + Ok((stream, _addr)) => { /* connection succeeded */ + // On some platforms (e.g. macOS) the accepted stream inherits the + // listener's non-blocking mode; the handler expects blocking I/O. + stream.set_nonblocking(false)?; let repository = self.repository.clone(); - thread::spawn(|| Handler::new(stream, repository).handle_client()); + workers.push(thread::spawn(|| { + Handler::new(stream, repository).handle_client() + })); + workers.retain(|worker| !worker.is_finished()); } - Err(_err) => { + Err(err) if err.kind() == ErrorKind::WouldBlock => { + /* nothing pending */ + thread::sleep(ACCEPT_POLL_INTERVAL); + } + Err(err) => { /* connection failed */ - break; + return Err(err); } } } + + // Stop accepting new connections; the listener is dropped with the Server, + // but from here on we simply never call accept() again. + self.wait_for_workers(workers); + Ok(()) } + + /// Give in-flight client handlers a bounded amount of time to finish. + fn wait_for_workers(&self, mut workers: Vec>) { + workers.retain(|worker| !worker.is_finished()); + if workers.is_empty() { + return; + } + + info!("waiting for {} in-flight request(s) to finish", workers.len()); + let deadline = Instant::now() + SHUTDOWN_GRACE_PERIOD; + + while Instant::now() < deadline { + workers.retain(|worker| !worker.is_finished()); + if workers.is_empty() { + break; + } + thread::sleep(Duration::from_millis(10)); + } + + for worker in workers { + if worker.is_finished() { + let _ = worker.join(); + } else { + warn!("abandoning a request still in flight after {SHUTDOWN_GRACE_PERIOD:?}"); + } + } + } } impl Drop for Server { fn drop(&mut self) { let _result = std::fs::remove_file(&self.socket_name); - info!("Shutting down"); + info!("shutdown complete"); + } +} + +fn signal_name(signal: usize) -> &'static str { + match signal as i32 { + SIGTERM => "SIGTERM", + SIGINT => "SIGINT", + _ => "signal", } } diff --git a/server/tests/shutdown_test.rs b/server/tests/shutdown_test.rs new file mode 100644 index 0000000..c85a873 --- /dev/null +++ b/server/tests/shutdown_test.rs @@ -0,0 +1,130 @@ +use common::io::read_string; +use common::request::Request; +use std::os::unix::fs::FileTypeExt; +use std::os::unix::net::UnixStream; +use std::path::PathBuf; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +const PASSWORD_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../test_passwords"); + +fn temp_socket_path(tag: &str) -> PathBuf { + let mut path = std::env::temp_dir(); + path.push(format!("saslauthd-{}-{}.sock", tag, std::process::id())); + let _ = std::fs::remove_file(&path); + path +} + +/// A spawned daemon that is killed if the test panics before stopping it. +struct Daemon(Child); + +impl Drop for Daemon { + fn drop(&mut self) { + if self.0.try_wait().ok().flatten().is_none() { + let _ = self.0.kill(); + let _ = self.0.wait(); + } + } +} + +fn is_socket(path: &PathBuf) -> bool { + std::fs::metadata(path) + .map(|m| m.file_type().is_socket()) + .unwrap_or(false) +} + +fn spawn_daemon(socket: &PathBuf) -> Daemon { + let child = Command::new(env!("CARGO_BIN_EXE_saslauthd")) + .arg("--socket-name") + .arg(socket) + .arg("--password-dir") + .arg(PASSWORD_DIR) + .spawn() + .expect("failed to spawn saslauthd"); + + let deadline = Instant::now() + Duration::from_secs(5); + while !is_socket(socket) { + assert!(Instant::now() < deadline, "daemon never created {socket:?}"); + std::thread::sleep(Duration::from_millis(10)); + } + Daemon(child) +} + +fn send_signal(daemon: &Daemon, signal: &str) { + let child = &daemon.0; + let status = Command::new("kill") + .arg(format!("-{signal}")) + .arg(child.id().to_string()) + .status() + .expect("failed to run kill"); + assert!(status.success(), "kill -{signal} failed"); +} + +fn wait_for_exit(daemon: &mut Daemon, timeout: Duration) -> std::process::ExitStatus { + let child = &mut daemon.0; + let deadline = Instant::now() + timeout; + loop { + if let Some(status) = child.try_wait().expect("try_wait failed") { + return status; + } + if Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + panic!("daemon did not exit within {timeout:?}"); + } + std::thread::sleep(Duration::from_millis(10)); + } +} + +fn authenticate(socket: &PathBuf, password: &str) -> String { + let mut stream = UnixStream::connect(socket).expect("connect to daemon"); + Request::new("bp", password, "imap", "") + .write_to_stream(&mut stream) + .expect("write request"); + read_string(&mut stream).expect("read response") +} + +fn assert_clean_shutdown(signal: &str) { + let socket = temp_socket_path(signal); + let mut daemon = spawn_daemon(&socket); + + // Make sure the daemon actually works before we shut it down. + assert_eq!(authenticate(&socket, "some-rad-password"), "OK"); + assert_eq!(authenticate(&socket, "wrong-password"), "NO"); + + let started = Instant::now(); + send_signal(&daemon, signal); + let status = wait_for_exit(&mut daemon, Duration::from_secs(1)); + + assert!(status.success(), "expected exit status 0 on {signal}, got {status}"); + assert!( + started.elapsed() < Duration::from_secs(1), + "shutdown took {:?}", + started.elapsed() + ); + assert!(!socket.exists(), "socket file {socket:?} was left behind"); +} + +#[test] +fn it_shuts_down_cleanly_on_sigterm() { + assert_clean_shutdown("TERM"); +} + +#[test] +fn it_shuts_down_cleanly_on_sigint() { + assert_clean_shutdown("INT"); +} + +#[test] +fn it_recreates_a_stale_socket_on_startup() { + let socket = temp_socket_path("stale"); + std::fs::write(&socket, b"stale").expect("create stale socket file"); + + let mut daemon = spawn_daemon(&socket); + assert_eq!(authenticate(&socket, "some-rad-password"), "OK"); + + send_signal(&daemon, "TERM"); + let status = wait_for_exit(&mut daemon, Duration::from_secs(1)); + assert!(status.success()); + assert!(!socket.exists()); +}