roundabout

use std::{

    env, io,

    os::fd::{AsRawFd, FromRawFd, OwnedFd},

    os::unix::process::CommandExt,

    process::{Child, Command},

    time::Duration,

};

use futures_lite::StreamExt;

use zbus::{

    blocking::{Connection, Proxy},

    proxy::SignalStream,

    zvariant::{OwnedObjectPath, Value},

};

use crate::{

    config::{Overrun, RamOverrun},

    units::Bytes,

};

const EXEC: &str = "__roundabout_exec";

const MANAGER: &str = "org.freedesktop.systemd1.Manager";

const BUS_NAME: &str = "org.freedesktop.systemd1";

const BUS_PATH: &str = "/org/freedesktop/systemd1";

pub struct Scope {

    connection: Connection,

    name: String,

}

fn properties(

    pid: u32,

    cores: u32,

    ram: Bytes,

    overrun: &Overrun,

) -> zbus::Result<Vec<(&'static str, Value<'static>)>> {

    let max = ram.0 as f64 * overrun.ram_slack.max(1.0);

    if !overrun.ram_slack.is_finite() || max >= u64::MAX as f64 {

        return Err(zbus::Error::Failure("RAM limit must be finite".into()));

    }

    let mut properties = vec![

        ("Description", Value::from("roundabout build")),

        ("PIDs", Value::from(vec![pid])),

        ("CollectMode", Value::from("inactive-or-failed")),

        ("CPUAccounting", Value::from(true)),

        ("MemoryAccounting", Value::from(true)),

        (

            "CPUQuotaPerSecUSec",

            Value::from(u64::from(cores.max(1)) * 1_000_000),

        ),

        ("MemoryMax", Value::from((max as u64).max(1))),

        ("OOMPolicy", Value::from("kill")),

        ("TimeoutStopUSec", Value::from(10_000_000_u64)),

    ];

    if overrun.ram == RamOverrun::Throttle {

        properties.push(("MemoryHigh", Value::from(ram.0.max(1))));

    }

    Ok(properties)

}

fn wait_job(mut signals: SignalStream<'_>, job: &OwnedObjectPath) -> zbus::Result<()> {

    async_io::block_on(futures_lite::future::race(

        async {

            while let Some(signal) = signals.next().await {

                let (_, path, _, result): (u32, OwnedObjectPath, String, String) =

                    signal.body().deserialize()?;

                if path == *job {

                    return if result == "done" {

                        Ok(())

                    } else {

                        Err(zbus::Error::Failure(format!("systemd job {result}")))

                    };

                }

            }

            Err(zbus::Error::Failure(

                "systemd disconnected before the job finished".into(),

            ))

        },

        async {

            async_io::Timer::after(Duration::from_secs(15)).await;

            Err(zbus::Error::Failure("systemd job timed out".into()))

        },

    ))

}

impl Scope {

    fn manager(&self) -> zbus::Result<Proxy<'_>> {

        Proxy::new(&self.connection, BUS_NAME, BUS_PATH, MANAGER)

    }

    fn start(

        connection: Connection,

        pid: u32,

        lease: u64,

        cores: u32,

        ram: Bytes,

        overrun: &Overrun,

    ) -> zbus::Result<Self> {

        let scope = Self {

            connection,

            name: format!("roundabout-{}-{lease}-{pid}.scope", std::process::id()),

        };

        let manager = scope.manager()?;

        let signals = async_io::block_on(

            manager

                .inner()

                .receive_signal_with_args("JobRemoved", &[(2, &scope.name)]),

        )?;

        let _: () = manager.call("Subscribe", &())?;

        let auxiliary: Vec<(&str, Vec<(&str, Value<'_>)>)> = Vec::new();

        let job: OwnedObjectPath = manager.call(

            "StartTransientUnit",

            &(

                &scope.name,

                "fail",

                properties(pid, cores, ram, overrun)?,

                auxiliary,

            ),

        )?;

        wait_job(signals, &job)?;

        drop(manager);

        Ok(scope)

    }

    pub fn signal(&self, signal: i32) -> zbus::Result<()> {

        self.manager()?

            .call("KillUnit", &(&self.name, "all", signal))

    }

    fn stop(&self) -> zbus::Result<()> {

        let manager = self.manager()?;

        let signals = async_io::block_on(

            manager

                .inner()

                .receive_signal_with_args("JobRemoved", &[(2, &self.name)]),

        )?;

        let job: OwnedObjectPath = manager.call("StopUnit", &(&self.name, "replace"))?;

        wait_job(signals, &job)

    }

}

impl Drop for Scope {

    fn drop(&mut self) {

        let _ = self.stop();

    }

}

fn gate() -> io::Result<(OwnedFd, OwnedFd)> {

    let mut descriptors = [-1; 2];

    if unsafe { libc::pipe2(descriptors.as_mut_ptr(), libc::O_CLOEXEC) } != 0 {

        return Err(io::Error::last_os_error());

    }

    Ok(unsafe {

        (

            OwnedFd::from_raw_fd(descriptors[0]),

            OwnedFd::from_raw_fd(descriptors[1]),

        )

    })

}

pub fn spawn(

    command: &Command,

    lease: u64,

    cores: u32,

    ram: Bytes,

    overrun: &Overrun,

) -> Result<(Child, Scope), String> {

    let connection =

        Connection::session().map_err(|error| format!("can't connect to the user bus, {error}"))?;

    let (read, write) = gate().map_err(|error| error.to_string())?;

    let descriptor = read.as_raw_fd();

    let mut runner = Command::new(env::current_exe().map_err(|error| error.to_string())?);

    runner

        .arg(EXEC)

        .arg(descriptor.to_string())

        .arg(command.get_program())

        .args(command.get_args())

        .process_group(0);

    for (name, value) in command.get_envs() {

        if let Some(value) = value {

            runner.env(name, value);

        } else {

            runner.env_remove(name);

        }

    }

    let parent = unsafe { libc::getpid() };

    unsafe {

        runner.pre_exec(move || {

            if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) != 0

                || libc::fcntl(descriptor, libc::F_SETFD, 0) != 0

            {

                return Err(io::Error::last_os_error());

            }

            if libc::getppid() != parent {

                return Err(io::Error::other("parent exited before launch"));

            }

            Ok(())

        });

    }

    let mut child = runner.spawn().map_err(|error| error.to_string())?;

    drop(read);

    let result = Scope::start(connection, child.id(), lease, cores, ram, overrun)

        .map_err(|error| format!("can't start the transient scope, {error}"))

        .and_then(|scope| {

            if unsafe { libc::write(write.as_raw_fd(), b"1".as_ptr().cast(), 1) } != 1 {

                return Err(io::Error::last_os_error().to_string());

            }

            Ok(scope)

        });

    drop(write);

    match result {

        Ok(scope) => Ok((child, scope)),

        Err(error) => {

            let _ = child.kill();

            let _ = child.wait();

            Err(error)

        }

    }

}

pub fn exec() -> Option<i32> {

    let mut arguments = env::args_os().skip(1);

    if arguments.next()?.as_encoded_bytes() != EXEC.as_bytes() {

        return None;

    }

    let Some(descriptor) = arguments

        .next()

        .and_then(|arg| arg.to_str()?.parse::<i32>().ok())

        .filter(|fd| *fd >= 3)

    else {

        return Some(1);

    };

    let descriptor = unsafe { OwnedFd::from_raw_fd(descriptor) };

    let mut byte = [0_u8];

    loop {

        let count = unsafe { libc::read(descriptor.as_raw_fd(), byte.as_mut_ptr().cast(), 1) };

        if count == 1 && byte[0] == b'1' {

            break;

        }

        if count != -1 || io::Error::last_os_error().kind() != io::ErrorKind::Interrupted {

            return Some(1);

        }

    }

    drop(descriptor);

    let Some(program) = arguments.next() else {

        return Some(1);

    };

    let error = Command::new(&program).args(arguments).exec();

    eprintln!(

        "roundabout: couldn't start {}: {error}",

        program.to_string_lossy()

    );

    Some(127)

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn limits_follow_the_grant_with_memory_headroom() {

        let overrun = Overrun::default();

        let properties = properties(123, 8, Bytes::gib(4), &overrun).unwrap();

        assert!(properties.contains(&("PIDs", Value::from(vec![123_u32]))));

        assert!(properties.contains(&("CPUQuotaPerSecUSec", Value::from(8_000_000_u64))));

        assert!(properties.contains(&("MemoryMax", Value::from(Bytes::gib(5).0))));

        assert!(!properties.iter().any(|(name, _)| *name == "MemoryHigh"));

    }

    #[test]

    fn throttle_starts_at_the_granted_memory() {

        let overrun = Overrun {

            ram: RamOverrun::Throttle,

            ..Overrun::default()

        };

        let properties = properties(123, 2, Bytes::gib(2), &overrun).unwrap();

        assert!(properties.contains(&("MemoryHigh", Value::from(Bytes::gib(2).0))));

        assert!(properties.contains(&("MemoryMax", Value::from(Bytes::mib(2560).0))));

    }

    #[test]

    fn small_grants_never_become_unlimited() {

        let overrun = Overrun {

            ram_slack: 0.5,

            ..Overrun::default()

        };

        let properties = properties(123, 0, Bytes::ZERO, &overrun).unwrap();

        assert!(properties.contains(&("CPUQuotaPerSecUSec", Value::from(1_000_000_u64))));

        assert!(properties.contains(&("MemoryMax", Value::from(1_u64))));

    }

    #[test]

    fn nonfinite_limits_are_rejected() {

        for ram_slack in [f64::INFINITY, f64::NEG_INFINITY, f64::NAN, f64::MAX] {

            let overrun = Overrun {

                ram_slack,

                ..Overrun::default()

            };

            assert!(properties(123, 1, Bytes::gib(2), &overrun).is_err());

        }

    }

}