POSIX.1-2024 timeout implemented in Rust

timeout.rs

// generated by GPT4o using the command summary in this README as the prompt:
// https://github.com/posix-utilities/timeout
use std::process::{Command, ExitStatus};
use std::time::Duration;
use std::os::unix::process::CommandExt;
use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
use std::thread;
use nix::sys::signal::{self, Signal};
use std::env;
use std::ffi::OsString;

fn parse_duration(s: &str) -> Result<Duration, &'static str> {
    if s.ends_with('s') {
        s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v)).map_err(|_| "Invalid seconds")
    } else if s.ends_with('m') {
        s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 60.0)).map_err(|_| "Invalid minutes")
    } else if s.ends_with('h') {
        s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 3600.0)).map_err(|_| "Invalid hours")
    } else if s.ends_with('d') {
        s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 86400.0)).map_err(|_| "Invalid days")
    } else {
        s.parse::<f64>().map(|v| Duration::from_secs_f64(v)).map_err(|_| "Invalid duration")
    }
}

fn main() {
    let args: Vec<OsString> = env::args_os().collect();

    if args.len() < 3 {
        eprintln!("Usage: timeout [options] <duration> <command> [arguments...]");
        std::process::exit(125);
    }

    let mut duration: Option<Duration> = None;
    let mut kill_after: Option<Duration> = None;
    let mut signal = Signal::SIGTERM;
    let mut preserve_exit_status = false;
    let mut for_process_only = false;
    let mut command_idx = None;

    let mut i = 1;
    while i < args.len() {
        let arg = &args[i];
        match arg.to_str() {
            Some("-f") => for_process_only = true,
            Some("-k") => {
                i += 1;
                if i < args.len() {
                    kill_after = Some(parse_duration(&args[i].to_string_lossy()).unwrap());
                }
            }
            Some("-p") => preserve_exit_status = true,
            Some("-s") => {
                i += 1;
                if i < args.len() {
                    signal = Signal::from_str(&args[i].to_string_lossy()).unwrap_or(Signal::SIGTERM);
                }
            }
            Some(s) if duration.is_none() => {
                duration = Some(parse_duration(s).unwrap());
            }
            _ => {
                command_idx = Some(i);
                break;
            }
        }
        i += 1;
    }

    let duration = duration.unwrap_or_else(|| {
        eprintln!("Duration must be specified");
        std::process::exit(125);
    });

    let command_idx = command_idx.unwrap_or_else(|| {
        eprintln!("Command must be specified");
        std::process::exit(125);
    });

    let command = &args[command_idx];
    let command_args = &args[command_idx + 1..];

    let child = Command::new(command)
        .args(command_args)
        .spawn()
        .expect("Failed to spawn process");

    let pid = child.id() as i32;

    let timeout_occurred = Arc::new(AtomicBool::new(false));
    let timeout_occurred_clone = Arc::clone(&timeout_occurred);

    let handle = thread::spawn(move || {
        thread::sleep(duration);
        timeout_occurred_clone.store(true, Ordering::SeqCst);
        if for_process_only {
            signal::kill(pid, signal).unwrap();
        } else {
            signal::killpg(pid, signal).unwrap();
        }
    });

    let status = child.wait().expect("Failed to wait on child");

    if timeout_occurred.load(Ordering::SeqCst) {
        if let Some(kill_duration) = kill_after {
            thread::sleep(kill_duration);
            signal::kill(pid, Signal::SIGKILL).unwrap();
        }
    }

    handle.join().unwrap();

    let exit_status_code = if preserve_exit_status {
        match status.code() {
            Some(code) => code,
            None => 124, // Default timeout status
        }
    } else if timeout_occurred.load(Ordering::SeqCst) {
        124
    } else {
        status.code().unwrap_or(125)
    };

    std::process::exit(exit_status_code);
}