//
// Syd: rock-solid application kernel
// src/pool.rs: Self growing / shrinking `ThreadPool` implementation
//
// Copyright (c) 2024, 2025, 2026 Ali Polatel <alip@chesswob.org>
// Based in part upon rusty_pool which is:
//     Copyright (c) Robin Friedli <robinfriedli@icloud.com>
//     SPDX-License-Identifier: Apache-2.0
//
// SPDX-License-Identifier: GPL-3.0

// Last sync with rusty_pool:
// Version 0.7.0
// Commit:d56805869ba3cbe47021d5660bbaf19ac5ec4bfb

use std::{
    env,
    fs::OpenOptions,
    io::Write,
    option::Option,
    os::{
        fd::{AsRawFd, RawFd},
        unix::fs::OpenOptionsExt,
    },
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, RwLock,
    },
    thread,
};

use dur::Duration;
use libseccomp::ScmpFilterContext;
use nix::{
    errno::Errno,
    fcntl::OFlag,
    sched::{unshare, CloneFlags},
    sys::signal::{sigaction, SaFlags, SigAction, SigHandler, SigSet, Signal},
    unistd::{getpid, pipe2, read, Gid, Uid},
};

use crate::{
    alert,
    config::*,
    confine::ExportMode,
    debug,
    err::{err2no, scmp2no, SydJoinHandle, SydResult},
    error,
    fd::closeexcept,
    fs::{block_signal, seccomp_export_pfc},
    hook::HandlerMap,
    info,
    retry::{retry_on_eintr, retry_on_intr},
    rwrite, rwriteln,
    sandbox::{Capability, Flags, Sandbox},
    workers::{
        aes::{AesLock, AesWorker},
        emu::Worker,
        int::Interrupter,
        ipc::IpcWorker,
        out::Timeouter,
        WorkerCache, WorkerData,
    },
};

// Signal handler function for SIGALRM.
extern "C" fn handle_sigalrm(_: libc::c_int) {}

/// Self growing / shrinking `ThreadPool` implementation.
#[derive(Clone)]
pub(crate) struct ThreadPool {
    core_size: usize,
    keep_alive: u16,
    flags: Flags,
    seccomp_fd: RawFd,
    pub(crate) cache: Arc<WorkerCache>,
    sandbox: Arc<RwLock<Sandbox>>,
    handlers: Arc<HandlerMap>,
    should_exit: Arc<AtomicBool>,
    worker_data: Arc<WorkerData>,
}

impl ThreadPool {
    /// Construct a new `ThreadPool` with the specified core pool size,
    /// max pool size and keep_alive time for non-core threads.
    ///
    /// `core_size` specifies the amount of threads to keep alive for as
    /// long as the `ThreadPool` exists and the seccomp fd remains open.
    ///
    /// `keep_alive` specifies the duration in milliseconds for which to
    /// keep non-core pool worker threads alive while they do not
    /// receive any work.
    #[expect(clippy::too_many_arguments)]
    pub(crate) fn new(
        seccomp_fd: RawFd,
        flags: Flags,
        core_size: usize,
        keep_alive: u16,
        sandbox: Arc<RwLock<Sandbox>>,
        handlers: Arc<HandlerMap>,
        should_exit: Arc<AtomicBool>,
        crypt_map: Option<AesLock>,
    ) -> Self {
        Self {
            sandbox,
            handlers,
            core_size,
            keep_alive,
            flags,
            should_exit,
            seccomp_fd,
            cache: Arc::new(WorkerCache::new(crypt_map)),
            worker_data: Arc::new(WorkerData::default()),
        }
    }

    /// Boot the thread pool. This is the main entry point.
    pub(crate) fn boot(self) -> SydResult<SydJoinHandle<()>> {
        // Export seccomp rules if requested.
        // We have to prepare the filter twice if exporting,
        // as we cannot move it safely between threads...
        #[expect(clippy::disallowed_methods)]
        match ExportMode::from_env() {
            Some(ExportMode::BerkeleyPacketFilter) => {
                // Worker rules
                let is_crypt = self.cache.crypt_map.is_some();
                let ctx = Worker::prepare_confine(self.seccomp_fd, self.flags, is_crypt, &[], &[])?;
                let file = OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .mode(0o400)
                    .open("syd_emu.bpf")?;
                ctx.export_bpf(file)?;

                // Interrupter rules
                // We pass dry_run=true to avoid Landlock confinement.
                let ctx = Interrupter::prepare_confine(
                    self.seccomp_fd,
                    getpid(),
                    self.flags,
                    &[],
                    &[],
                    true,
                )?;
                let file = OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .mode(0o400)
                    .open("syd_int.bpf")?;
                ctx.export_bpf(file)?;

                // IPC thread rules
                // We pass dummy RawFd=2525 for epoll FD.
                // We pass dry_run=true to avoid Landlock confinement.
                let ctx = IpcWorker::prepare_confine(2525, self.flags, &[], &[], true)?;
                let file = OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .mode(0o400)
                    .open("syd_ipc.bpf")?;
                ctx.export_bpf(file)?;

                // Aes worker rules
                let ctx = AesWorker::prepare_confine(self.flags, &[], &[], true)?;
                let file = OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .mode(0o400)
                    .open("syd_aes.bpf")?;
                ctx.export_bpf(file)?;
            }
            Some(ExportMode::PseudoFiltercode) => {
                // Lock stdout to prevent concurrent access.
                let mut stdout = std::io::stdout().lock();

                rwriteln!(
                    stdout,
                    "# Syd monitor rules with seccomp fd {}",
                    self.seccomp_fd
                )?;
                let is_crypt = self.cache.crypt_map.is_some();
                let ctx = Worker::prepare_confine(self.seccomp_fd, self.flags, is_crypt, &[], &[])?;
                rwrite!(stdout, "{}", seccomp_export_pfc(&ctx)?)?;

                // We pass dry_run=true to avoid Landlock confinement.
                rwriteln!(
                    stdout,
                    "# Syd interrupter rules with seccomp fd {}",
                    self.seccomp_fd
                )?;
                let ctx = Interrupter::prepare_confine(
                    self.seccomp_fd,
                    getpid(),
                    self.flags,
                    &[],
                    &[],
                    true,
                )?;
                rwrite!(stdout, "{}", seccomp_export_pfc(&ctx)?)?;

                // We pass dummy RawFd=2525 for epoll FD.
                // We pass dry_run=true to avoid Landlock confinement.
                rwriteln!(stdout, "# Syd ipc rules")?;
                let ctx = IpcWorker::prepare_confine(2525, self.flags, &[], &[], true)?;
                rwrite!(stdout, "{}", seccomp_export_pfc(&ctx)?)?;

                rwriteln!(stdout, "# Syd encryptor rules")?;
                let ctx = AesWorker::prepare_confine(self.flags, &[], &[], true)?;
                rwrite!(stdout, "{}", seccomp_export_pfc(&ctx)?)?;
            }
            _ => {}
        }

        // Ensure the lazy num_cpus::get is called before
        // the CPU pinning otherwise it may report incorrect
        // value.
        let nproc = *NPROC;
        info!("ctx": "boot", "op": "check_num_cpus",
            "msg": format!("detected {nproc} CPUs on the system"),
            "num_cpus": nproc);

        // Spawn the monitor thread which may confine itself, and spawn
        // emulator threads. Note, this will panic if it cannot spawn
        // the initial emulator thread which is going to tear everything
        // down. Return a join handle to the main thread so it can wait
        // for the monitor thread to gracefully exit which in turn is
        // going to wait for the AES threads to gracefully exit.
        self.monitor()
    }

    /// Spawn a monitor thread that watches the worker pool busy count,
    /// and spawns new helper threads as necessary. This is done to
    /// ensure a sandbox process cannot DOS Syd by merely exhausting
    /// workers by e.g. opening the read end of a FIFO over and over
    /// again.
    #[expect(clippy::cognitive_complexity)]
    pub(crate) fn monitor(self) -> SydResult<SydJoinHandle<()>> {
        thread::Builder::new()
            .name("syd_mon".to_string())
            .stack_size(MON_STACK_SIZE)
            .spawn(move || {
                let sandbox = self.sandbox.read().unwrap_or_else(|err| err.into_inner());

                // SAFETY: We use exit_group(2) here to bail,
                // because this unsharing is a critical safety feature.
                if let Err(errno) = unshare(CloneFlags::CLONE_FS | CloneFlags::CLONE_FILES) {
                    alert!("ctx": "boot", "op": "unshare_monitor_thread",
                        "msg": format!("failed to unshare(CLONE_FS|CLONE_FILES): {errno}"),
                        "err": errno as i32);
                    std::process::exit(101);
                }

                // SAFETY: The monitor thread needs to inherit FDs.
                // We have to sort the set as the FDs are randomized.
                #[expect(clippy::cast_sign_loss)]
                let mut set = vec![
                    ROOT_FD() as libc::c_uint,
                    PROC_FD() as libc::c_uint,
                    NULL_FD() as libc::c_uint,
                    sandbox.fpid as libc::c_uint,
                    self.seccomp_fd as libc::c_uint,
                    crate::log::LOG_FD.load(Ordering::Relaxed) as libc::c_uint,
                ];

                let crypt = if sandbox.enabled(Capability::CAP_CRYPT) {
                    Some((sandbox.crypt_setup()?, sandbox.crypt_tmp))
                } else {
                    None
                };
                #[expect(clippy::cast_sign_loss)]
                if let Some((crypt_fds, crypt_tmp)) = crypt {
                    set.push(crypt_fds.0 as libc::c_uint);
                    set.push(crypt_fds.1 as libc::c_uint);
                    if let Some(crypt_tmp) = crypt_tmp {
                        set.push(crypt_tmp as libc::c_uint);
                    }
                }
                set.sort_unstable();
                closeexcept(&set)?;
                drop(set);

                // Spawn the interrupt thread which will confine itself.
                self.try_spawn_interrupt(&sandbox.transit_uids, &sandbox.transit_gids)?;

                // Spawn the AES thread if encryption is on.
                let crypt_handle = if let Some((fds, tmp)) = crypt {
                    let map = self.cache.crypt_map.as_ref().map(Arc::clone).ok_or(Errno::ENOKEY)?;
                    let should_exit = Arc::clone(&self.should_exit);
                    Some(self.try_spawn_aes(
                            fds,
                            map,
                            tmp.is_none(),
                            should_exit,
                            &sandbox.transit_uids,
                            &sandbox.transit_gids)?)
                } else {
                    None
                };

                if let Some(tmout) = sandbox.tmout {
                    // Spawn the timeouter thread which will confine itself.
                    self.try_spawn_timeout(tmout, &sandbox.transit_uids, &sandbox.transit_gids)?;
                }

                info!("ctx": "boot", "op": "start_monitor_thread",
                    "msg": format!("started monitor thread with pool size set to {} threads and keep alive set to {} seconds",
                        self.core_size,
                        self.keep_alive.saturating_div(1000)),
                    "core_size": self.core_size,
                    "keep_alive": self.keep_alive);

                // SAFETY:
                // 1. If sandbox is locked, confine right away.
                //    Pass confined parameter to try_spawn so subsequent
                //    spawned threads don't need to reapply the same filter
                //    as it is inherited.
                // 2. If sandbox is not locked yet, build the seccomp context anyway,
                //    precompute it and pass it to emulator threads for fast confinement.
                let dry_run = env::var_os(ENV_SKIP_SCMP).is_some() || ExportMode::from_env().is_some();
                let safe_setid = self.flags.intersects(Flags::FL_ALLOW_SAFE_SETUID | Flags::FL_ALLOW_SAFE_SETGID);
                let is_crypt = self.cache.crypt_map.is_some();

                let mut ctx = if !dry_run {
                    let ctx = Worker::prepare_confine(
                        self.seccomp_fd,
                        self.flags,
                        is_crypt,
                        &sandbox.transit_uids,
                        &sandbox.transit_gids)?;

                    if Sandbox::locked_once() {
                        // Sandbox locked, confine right away.
                        //
                        // SAFETY: We use exit_group(2) here to bail,
                        // because this confinement is a critical safety feature.
                        if let Err(error) = ctx.load() {
                            let errno = scmp2no(&error).unwrap_or(Errno::ENOSYS);
                            alert!("ctx": "boot", "op": "confine_monitor_thread",
                                "msg": format!("failed to confine: {error}"),
                                "err": errno as i32);
                            std::process::exit(101);
                        }

                        info!("ctx": "confine", "op": "confine_monitor_thread",
                            "msg": format!("monitor thread confined with{} SROP mitigation",
                                if safe_setid { "out" } else { "" }));

                        None
                    } else {
                        // Sandbox not locked yet, precompute and save filter.
                        //
                        // SAFETY: We use exit_group(2) here to bail,
                        // because this confinement is a critical safety feature.
                        #[cfg(libseccomp_v2_6)]
                        if let Err(error) = ctx.precompute() {
                            let errno = scmp2no(&error).unwrap_or(Errno::ENOSYS);
                            alert!("ctx": "boot", "op": "confine_monitor_thread",
                                "msg": format!("failed to precompute: {error}"),
                                "err": errno as i32);
                            std::process::exit(101);
                        }

                        info!("ctx": "confine", "op": "confine_monitor_thread",
                            "msg": "monitor thread is running unconfined because sandbox isn't locked yet");

                        Some(ctx)
                    }
                } else {
                    error!("ctx": "confine", "op": "confine_monitor_thread",
                        "msg": "monitor thread is running unconfined in debug mode");
                    None
                };
                drop(sandbox); // release the read-lock.

                info!("ctx": "boot", "op": "start_core_emulator_threads",
                    "msg": format!("starting {} core emulator thread{}, sandboxing started!",
                        self.core_size,
                        if self.core_size > 1 { "s" } else { "" }),
                    "core_size": self.core_size,
                    "keep_alive": self.keep_alive);

                // Spawn the initial core emulator thread.
                self.try_spawn(ctx.as_ref()).map(drop)?;

                // Wait for grace period to give the initial
                // core emulator thread a chance to spawn itself.
                std::thread::sleep(MON_GRACE_TIME.into());

                loop {
                    // Confine and drop filter if sandbox is locked.
                    if let Some(ref filter) = ctx {
                        if Sandbox::locked_once() {
                            // SAFETY: We use exit_group(2) here to bail,
                            // because this confinement is a critical safety feature.
                            if let Err(error) = filter.load() {
                                let errno = scmp2no(&error).unwrap_or(Errno::ENOSYS);
                                alert!("ctx": "boot", "op": "confine_monitor_thread",
                                    "msg": format!("failed to confine: {error}"),
                                    "err": errno as i32);
                                std::process::exit(101);
                            }

                            info!("ctx": "confine", "op": "confine_monitor_thread",
                                "msg": format!("monitor thread confined with{} SROP mitigation",
                                    if safe_setid { "out" } else { "" }));

                            // SAFETY: We cannot free the seccomp context here,
                            // because it may have references in emulator
                            // threads.
                            std::mem::forget(ctx);

                            ctx = None;
                        }
                    }

                    if self.should_exit.load(Ordering::Relaxed) {
                        // Time to exit.
                        break;
                    }

                    // Spawn a new thread if all others are busy.
                    match self.try_spawn(ctx.as_ref()) {
                        Ok(Some(_)) => {
                            // We have spawned a new emulator thread,
                            // wait for one cycle before reattempting.
                            std::thread::sleep(MON_CYCLE_TIME.into());
                        }
                        Ok(None) => {
                            // We have idle threads, no need to spawn a new worker.
                            // Wait for grace period before reattempting.
                            std::thread::sleep(MON_GRACE_TIME.into());
                        }
                        Err(errno) => {
                            alert!("ctx": "spawn",
                                "msg": format!("spawn emulator failed: {errno}!"),
                                "err": errno as i32,
                                "core_size": self.core_size,
                                "keep_alive": self.keep_alive);

                            // Be defensive and signal stuck emulator threads to make
                            // better use of available resources.
                            self.signal_int();

                            // Wait for grace period before reattempting.
                            std::thread::sleep(MON_GRACE_TIME.into());
                        }
                    }
                }

                // Wait for AES threads to gracefully exit.
                if let Some(crypt_handle) = crypt_handle {
                    crypt_handle.join().or(Err(Errno::EAGAIN))??;
                }

                Ok(())
            })
            .map_err(|err| err2no(&err).into())
    }

    /// Spawn an interrupt handler thread to unblock Syd syscall
    /// handler threads when the respective sandbox process
    /// receives a non-restarting signal.
    pub(crate) fn try_spawn_interrupt(
        &self,
        transit_uids: &[(Uid, Uid)],
        transit_gids: &[(Gid, Gid)],
    ) -> SydResult<SydJoinHandle<()>> {
        // Block SIGALRM, this mask will be inherited by emulators.
        block_signal(Signal::SIGALRM)?;

        // Set up the signal handler for SIGALRM.
        let sig_action = SigAction::new(
            SigHandler::Handler(handle_sigalrm),
            SaFlags::empty(),
            SigSet::empty(),
        );

        // SAFETY: Register the handler for SIGALRM.
        // This handler is per-process.
        unsafe { sigaction(Signal::SIGALRM, &sig_action) }?;

        // Set up a notification pipe and wait for
        // the interrupt worker to start and unshare CLONE_F{ILE,}S.
        let (pipe_rd, pipe_wr) = pipe2(OFlag::O_CLOEXEC)?;

        let handle = retry_on_intr(|| {
            Interrupter::new(
                self.seccomp_fd,
                self.flags,
                transit_uids,
                transit_gids,
                Arc::clone(&self.should_exit),
                Arc::clone(&self.cache),
            )
            .try_spawn((pipe_rd.as_raw_fd(), pipe_wr.as_raw_fd()))
        })?;

        // Wait for startup notification.
        let mut buf = [0u8; 1];
        match retry_on_eintr(|| read(&pipe_rd, &mut buf[..]))? {
            0 => {
                // Interrupt thread died before unshare.
                // This should ideally never happen.
                return Err(Errno::EIO.into());
            }
            1 if buf[0] == 42 => {
                // Interrupt thread unshared successfully.
                // We can go ahead and close our copies now.
            }
            _ => unreachable!("BUG: The meaning of life is not {:#x}!", buf[0]),
        }

        Ok(handle)
    }

    /// Spawn an timeout handler thread to unblock Syd syscall
    /// handler threads when the respective sandbox process
    /// receives a non-restarting signal.
    pub(crate) fn try_spawn_timeout(
        &self,
        tmout: Duration,
        transit_uids: &[(Uid, Uid)],
        transit_gids: &[(Gid, Gid)],
    ) -> SydResult<SydJoinHandle<()>> {
        // Set up a notification pipe and wait for
        // the timeout worker to start and unshare CLONE_F{ILE,}S.
        let (pipe_rd, pipe_wr) = pipe2(OFlag::O_CLOEXEC)?;

        let handle = retry_on_intr(|| {
            Timeouter::new(
                tmout,
                self.flags,
                transit_uids,
                transit_gids,
                Arc::clone(&self.should_exit),
            )
            .try_spawn((pipe_rd.as_raw_fd(), pipe_wr.as_raw_fd()))
        })?;

        // Wait for startup notification.
        let mut buf = [0u8; 1];
        match retry_on_eintr(|| read(&pipe_rd, &mut buf[..]))? {
            0 => {
                // timeout thread died before unshare.
                // This should ideally never happen.
                return Err(Errno::EIO.into());
            }
            1 if buf[0] == 42 => {
                // timeout thread unshared successfully.
                // We can go ahead and close our copies now.
            }
            _ => unreachable!("BUG: The meaning of life is not {:#x}!", buf[0]),
        }

        Ok(handle)
    }

    /// Try to create a new encryption thread.
    pub(crate) fn try_spawn_aes(
        &self,
        fdalg: (RawFd, RawFd),
        files: AesLock,
        memfd: bool,
        should_exit: Arc<AtomicBool>,
        transit_uids: &[(Uid, Uid)],
        transit_gids: &[(Gid, Gid)],
    ) -> Result<SydJoinHandle<()>, Errno> {
        let worker = AesWorker::new(
            fdalg,
            files,
            self.flags,
            memfd,
            should_exit,
            transit_uids,
            transit_gids,
        );

        // AesWorker has only RawFds as Fds which
        // we do _not_ want to duplicate on clone,
        // so we can get away with a clone here...
        retry_on_intr(|| worker.clone().try_spawn())
    }

    /// Try to create a new worker thread as needed.
    /// Returns Ok(Some((SydJoinHandle, bool))) if spawn succeeded, Ok(None) if no spawn was needed.
    /// The boolean in the success case is true if the thread we spawned was a core thread.
    #[expect(clippy::cognitive_complexity)]
    #[expect(clippy::type_complexity)]
    pub(crate) fn try_spawn(
        &self,
        ctx: Option<&ScmpFilterContext>,
    ) -> Result<Option<(SydJoinHandle<()>, bool)>, Errno> {
        // Create a new worker if there are no idle threads and the
        // current worker count is lower than the max pool size.
        let worker_count_val = self.worker_data.0.load(Ordering::Relaxed);
        let (curr_worker_count, busy_worker_count) = WorkerData::split(worker_count_val);

        let keep_alive = if curr_worker_count < self.core_size {
            // Create a new core worker if current pool size is below
            // core size during the invocation of this function.
            debug!("ctx": "spawn",
                "msg": "creating new core emulator",
                "busy_worker_count": busy_worker_count,
                "curr_worker_count": curr_worker_count,
                "core_size": self.core_size);
            None
        } else if busy_worker_count < curr_worker_count {
            // We have idle threads, no need to spawn a new worker.
            debug!("ctx": "spawn",
                "msg": "idle emulator exists, no need to spawn",
                "busy_worker_count": busy_worker_count,
                "curr_worker_count": curr_worker_count,
                "core_size": self.core_size);
            return Ok(None);
        } else if curr_worker_count < *EMU_MAX_SIZE {
            // Create a new helper worker if the current worker count is
            // below the EMU_MAX_SIZE and the pool has been observed to
            // be busy (no idle workers) during the invocation of this
            // function.
            debug!("ctx": "spawn",
                "msg": "creating new idle emulator",
                "busy_worker_count": busy_worker_count,
                "curr_worker_count": curr_worker_count,
                "core_size": self.core_size,
                "keep_alive": self.keep_alive);
            Some(self.keep_alive)
        } else {
            // We cannot spawn anymore workers!
            // Ideally, this should never happen.
            alert!("ctx": "spawn",
                "msg": "emulator capacity exceeded, can not spawn new",
                "busy_worker_count": busy_worker_count,
                "curr_worker_count": curr_worker_count,
                "core_size": self.core_size,
                "keep_alive": self.keep_alive);
            return Err(Errno::ERANGE);
        };

        // Try to spawn a new worker.
        Ok(Some((
            retry_on_intr(|| {
                Worker::new(
                    self.seccomp_fd,
                    Arc::clone(&self.cache),
                    Arc::clone(&self.sandbox),
                    Arc::clone(&self.handlers),
                    keep_alive,
                    Arc::clone(&self.should_exit),
                    Arc::clone(&self.worker_data),
                )
                .try_spawn(ctx)
            })?,
            keep_alive.is_none(),
        )))
    }

    fn signal_int(&self) {
        // Unblock stuck emulator threads with manual signal.
        let mut nsig = 0usize;

        {
            let (ref lock, ref cvar) = *self.cache.sysint_map.sys_block;
            let mut map = lock.lock().unwrap_or_else(|err| err.into_inner());

            for interrupt in map.iter_mut() {
                if !interrupt.signal {
                    interrupt.signal = true;
                    nsig = nsig.saturating_add(1);
                }
            }

            // Notify interrupt thread.
            cvar.notify_one();
        }

        if nsig > 0 {
            alert!("ctx": "spawn",
                "msg": format!("signaled interrupt to unstuck {nsig} emulators"),
                "nsig": nsig,
                "core_size": self.core_size,
                "keep_alive": self.keep_alive);
        }
    }
}
