//
// Syd: rock-solid unikernel
// src/pool.rs: Self growing / shrinking `ThreadPool` implementation
//
// Copyright (c) 2024 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-or-later

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

use std::{
    mem::MaybeUninit,
    option::Option,
    os::fd::{BorrowedFd, RawFd},
    sync::{
        atomic::{AtomicBool, AtomicU32, Ordering},
        Arc,
    },
    thread,
};

use libseccomp::{error::SeccompError, ScmpAction, ScmpFilterContext, ScmpSyscall};
use libseccomp_sys::{
    seccomp_notif, seccomp_notif_resp, seccomp_notify_receive, seccomp_notify_respond,
};
use nix::{
    errno::Errno,
    poll::PollTimeout,
    sys::epoll::{Epoll, EpollCreateFlags, EpollEvent, EpollFlags},
    unistd::close,
};
use parking_lot::RwLock;

use crate::{
    cache::{AddrCache, PathCache},
    compat::epoll_ctl_safe,
    config::*,
    hook::{HandlerMap, UNotifyEventRequest},
    proc::proc_mmap,
    sandbox::Sandbox,
    ScmpNotifReq, SydArch, Sydcall,
};

const EOWNERDEAD: i32 = -nix::libc::EOWNERDEAD;

/// The absolute maximum number of workers. This corresponds to the
/// maximum value that can be stored within half the bits of u32, as two
/// counters (total workers and busy workers) are stored in one
/// AtomicU32.
pub(crate) const MAX_SIZE: u16 = u16::MAX;

/// Self growing / shrinking `ThreadPool` implementation.
#[derive(Clone)]
pub(crate) struct ThreadPool {
    core_size: u16,
    keep_alive: u16,
    safe_setid: bool,
    fd: RawFd,
    epoll: Arc<Epoll>,
    cache: Arc<(PathCache, AddrCache)>,
    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. This
    /// function creates an epoll instance and adds the seccomp fd to it
    /// but it does not spawn any 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.
    pub(crate) fn new(
        safe_setid: bool,
        core_size: u16,
        keep_alive: u16,
        fd: RawFd,
        cache: Arc<(PathCache, AddrCache)>,
        sandbox: Arc<RwLock<Sandbox>>,
        handlers: Arc<HandlerMap>,
    ) -> Result<Self, Errno> {
        // Create epoll instance.
        let epoll = Epoll::new(EpollCreateFlags::EPOLL_CLOEXEC)?;

        // SAFETY: Borrow FD to make I/O safe API hippie.
        let seccomp_fd = unsafe { BorrowedFd::borrow_raw(fd) };

        // Add seccomp fd to epoll (becomes readable when system call is interrupted).
        // We do not need the data field as we operate on a single fd.
        // Quoting: https://idea.popcount.org/2017-02-20-epoll-is-fundamentally-broken-12/
        // """
        // The best and the only scalable approach is to use recent
        // Kernel 4.5+ and use level-triggered events with
        // EPOLLEXCLUSIVE flag. This will ensure only one thread is
        // woken for an event, avoid "thundering herd" issue and scale
        // properly across multiple CPU's.
        // """
        let event = EpollEvent::new(EpollFlags::EPOLLIN | EpollFlags::EPOLLEXCLUSIVE, 0);
        epoll_ctl_safe(&epoll, &seccomp_fd, Some(event))?;

        Ok(Self {
            fd,
            cache,
            sandbox,
            handlers,
            core_size,
            keep_alive,
            safe_setid,
            epoll: Arc::new(epoll),
            should_exit: Arc::new(AtomicBool::new(false)),
            worker_data: Arc::new(WorkerData::default()),
        })
    }

    /// Boot the thread pool. This is the main entry point.
    pub(crate) fn boot(self) -> Result<(), std::io::Error> {
        // Export seccomp rules if requested.
        // We have to prepare the filter twice if exporting,
        // as we cannot move it safely between threads...
        if std::env::var("SYD_SECX").is_ok() {
            println!("# Syd monitor rules");
            if let Ok(ctx) = Worker::prepare_confine(self.safe_setid) {
                let _ = ctx.export_pfc(&mut std::io::stdout());
            }
        }

        // 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.
        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.
    #[allow(clippy::cognitive_complexity)]
    pub(crate) fn monitor(self) -> Result<(), std::io::Error> {
        thread::Builder::new()
            .name("syd_mon".to_string())
            .stack_size(MON_STACK_SIZE)
            .spawn(move || {
                // SAFETY: 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.
                let mut confined = if Sandbox::locked_once() {
                    Worker::prepare_confine(self.safe_setid)
                        .expect("prepare monitor thread confinement")
                        .load()
                        .expect("confine monitor thread");
                    crate::info!("ctx": "mon", "op": "confine",
                        "nnp": true, "tsync": false,
                        "sigreturn": self.safe_setid);
                    true
                } else {
                    false
                };
                // SAFETY: Panic if we cannot spawn the initial core thread.
                // There's little sense in continuing in this case.
                self.try_spawn(confined)
                    .expect("spawn initial emulator thread");
                std::thread::sleep(MON_GRACE_TIME);
                loop {
                    if !confined && Sandbox::locked_once() {
                        Worker::prepare_confine(self.safe_setid)
                            .expect("prepare monitor thread confinement")
                            .load()
                            .expect("confine monitor thread");
                        crate::info!("ctx": "mon", "op": "confine",
                            "nnp": true, "tsync": false,
                            "sigreturn": self.safe_setid);
                        confined = true;
                    }
                    if self.should_exit.load(Ordering::Relaxed) {
                        break; // Time to exit.
                    } else if let Ok(false) = self.try_spawn(confined) {
                        // When no idle threads are present we wait for cycle period.
                        std::thread::sleep(MON_CYCLE_TIME);
                    } else {
                        // When we spawned a new thread or there's an error spawning
                        // a thread we wait for the grace period.
                        std::thread::sleep(MON_GRACE_TIME);
                    }
                }
            })
            .map(drop)
    }

    /// Try to create a new worker thread as needed.
    /// Returns Ok(true) if spawn succeeded, Ok(false) if no spawn was needed.
    pub(crate) fn try_spawn(&self, confined: bool) -> Result<bool, std::io::Error> {
        // 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);

        if busy_worker_count < curr_worker_count {
            // We have idle threads, no need to spawn a new worker.
            return Ok(false);
        }

        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.
            None
        } else if curr_worker_count < MAX_SIZE {
            // Create a new helper worker if the current worker count is
            // below the MAX_SIZE and the pool has been observed to be
            // busy (no idle workers) during the invocation of this
            // function.
            Some(self.keep_alive)
        } else {
            // We cannot spawn anymore workers!
            // Ideally, this should never happen.
            return Ok(false);
        };

        // Try to spawn a new worker.
        Worker::new(
            self.fd,
            self.safe_setid,
            Arc::clone(&self.epoll),
            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(confined)?;

        Ok(true)
    }
}

#[derive(Clone)]
struct Worker {
    fd: RawFd,
    safe_setid: bool,
    epoll: Arc<Epoll>,
    cache: Arc<(PathCache, AddrCache)>,
    sandbox: Arc<RwLock<Sandbox>>,
    handlers: Arc<HandlerMap>,
    keep_alive: Option<u16>,
    should_exit: Arc<AtomicBool>,
    worker_data: Arc<WorkerData>,
}

impl Worker {
    fn new(
        fd: RawFd,
        safe_setid: bool,
        epoll: Arc<Epoll>,
        cache: Arc<(PathCache, AddrCache)>,
        sandbox: Arc<RwLock<Sandbox>>,
        handlers: Arc<HandlerMap>,
        keep_alive: Option<u16>,
        should_exit: Arc<AtomicBool>,
        worker_data: Arc<WorkerData>,
    ) -> Self {
        Worker {
            fd,
            safe_setid,
            epoll,
            cache,
            sandbox,
            handlers,
            keep_alive,
            should_exit,
            worker_data,
        }
    }

    fn try_spawn(self, mut confined: bool) -> Result<(), std::io::Error> {
        thread::Builder::new()
            .name("syd_emu".to_string())
            .stack_size(EMU_STACK_SIZE)
            .spawn(move || {
                // Create sentinel.
                let mut sentinel = Sentinel::new(&self);

                // Thread successfully started, increment total worker count.
                self.worker_data.increment_worker_total();

                loop {
                    // SAFETY: Confine if/once locked.
                    if !confined && Sandbox::locked_once() {
                        // SAFETY: Panic if we cannot confine the thread.
                        Self::prepare_confine(self.safe_setid)
                            .expect("prepare emulator thread confinement")
                            .load()
                            .expect("confine emulator thread");
                        crate::info!("ctx": "emu", "op": "confine",
                            "nnp": true, "tsync": false,
                            "sigreturn": self.safe_setid);
                        confined = true;
                    }

                    // Wait for the request to become ready as necessary.
                    // epoll_wait(2) will timeout and exit for non-core threads.
                    if self.poll().is_err() {
                        // Timeout or critical error.
                        // Decrement worker total and exit.
                        self.worker_data.decrement_worker_total();
                        break;
                    }

                    // Receive seccomp notification.
                    let request = if let Ok(request) = self.receive() {
                        request
                    } else {
                        // Critical error, decrement worker total and exit.
                        self.worker_data.decrement_worker_total();
                        break;
                    };

                    if let Some(request) = request {
                        // Mark thread busy and handle request.
                        sentinel.is_working = true;
                        self.worker_data.increment_worker_busy();
                        self.handle(request);
                        sentinel.is_working = false;
                        self.worker_data.decrement_worker_busy();
                    } // else process died-midway, continue.
                }
            })
            .map(drop)
    }

    fn receive(&self) -> Result<Option<ScmpNotifReq>, Errno> {
        // Receive and return request.
        // Break if file descriptor was closed.
        // Ignore rest of the errors as we cannot handle them,
        // e.g: EINTR|ENOENT: task is killed mid-way.
        match self.read() {
            Ok(request) => Ok(Some(request)),
            Err(Errno::EBADF) => Err(Errno::EBADF),
            Err(_) => Ok(None),
        }
    }

    fn read(&self) -> Result<ScmpNotifReq, Errno> {
        // Use libc::seccomp_notif rather than libseccomp_sys's.
        // The latter is opaque and required us to do a heap
        // allocation which we don't always want.
        let mut req: MaybeUninit<nix::libc::seccomp_notif> = MaybeUninit::zeroed();
        let ptr = req.as_mut_ptr() as *mut seccomp_notif;

        // SAFETY: libseccomp's wrapper allocates each call.
        // Note: EINTR means child killed by signal!
        match unsafe { seccomp_notify_receive(self.fd, ptr) } {
            0 => {
                // SAFETY: seccomp_notify_receive returned success.
                // Request is populated and accessing it is safe.
                ScmpNotifReq::from_sys(unsafe { req.assume_init() })
            }
            _ => Err(Errno::last()),
        }
    }

    #[allow(clippy::cognitive_complexity)]
    fn handle(&self, req: ScmpNotifReq) {
        let request = UNotifyEventRequest::new(
            req,
            self.fd,
            Arc::clone(&self.cache),
            Arc::clone(&self.sandbox),
        );
        let sydcall = Sydcall(req.data.syscall, req.data.arch);
        let handler = if let Some(handler) = self.handlers.get(&sydcall) {
            handler
        } else {
            unreachable!(
                "BUG: Missing hook for syscall {}!",
                i32::from(req.data.syscall)
            )
        };
        let resp = handler(&request);
        if resp.id == 0 && resp.val == 0 && resp.error == 0 && resp.flags == 0 {
            // Dummy seccomp response.
            // A previous addfd request has
            // already replied to the request.
            // Nothing left to do here.
            return;
        } else if resp.error == EOWNERDEAD {
            // EOWNERDEAD is a pseudo errno used by
            // the stat handler thread to close the
            // seccomp notify fd upon receiving the
            // "ghost" command.
            let req = request.get_request();
            crate::warn!("ctx": "main", "op": "ghost",
                "pid": req.pid, "sys": sydcall,
                "arch": SydArch(req.data.arch),
                "args": req.data.args,
                "src": proc_mmap(req.pid(), req.data.arch, req.data.instr_pointer));
        }
        let mut resp = nix::libc::seccomp_notif_resp {
            id: resp.id,
            val: resp.val,
            error: resp.error,
            flags: resp.flags,
        };
        // libc's data type is not opaque unlike libseccomp's.
        let ptr: *mut seccomp_notif_resp = std::ptr::addr_of_mut!(resp) as *mut seccomp_notif_resp;

        // SAFETY:
        // 1. libseccomp's version allocates needlessly
        // 2. Nothing we can do on errors, EINTR means child dead.
        unsafe { seccomp_notify_respond(self.fd, ptr) };

        // See above.
        if resp.error == EOWNERDEAD {
            // Note, threads blocked on epoll_wait will not
            // wake up even if we close the epoll fd or
            // delete the seccomp fd from epoll wait-list here.
            // That said, they'll never ever wake up again,
            // and therefore will not consume system resources.
            let _ = close(self.fd);

            // Inform the monitor thread to exit.
            self.should_exit.store(true, Ordering::Relaxed);
        }
    }

    fn poll(&self) -> Result<(), Errno> {
        let timeout = if let Some(keep_alive) = self.keep_alive {
            PollTimeout::from(keep_alive)
        } else {
            PollTimeout::NONE
        };

        // Wait for an event and handle EINTR.
        let mut events = [EpollEvent::empty(); 1];
        loop {
            return match self.epoll.wait(&mut events, timeout) {
                Ok(0) if self.keep_alive.is_some() => Err(Errno::ETIMEDOUT),
                Ok(0) | Err(Errno::EINTR) => continue,
                Ok(_) if events[0].events().contains(EpollFlags::EPOLLIN) => Ok(()),
                Ok(_) => Err(Errno::EBADF), // POLLHUP!
                Err(errno) => Err(errno),
            };
        }
    }

    /// Confine Worker thread.
    #[allow(clippy::cognitive_complexity)]
    fn prepare_confine(safe_setid: bool) -> Result<ScmpFilterContext, SeccompError> {
        // Create seccomp filter with default action.
        let mut ctx = ScmpFilterContext::new_filter(ScmpAction::KillProcess)?;

        // Enforce the NO_NEW_PRIVS functionality before
        // loading the seccomp filter into the kernel.
        ctx.set_ctl_nnp(true)?;

        // DO NOT synchronize filter to all threads.
        // Main thread will confine itself.
        ctx.set_ctl_tsync(false)?;

        // We kill for bad system call and bad arch.
        ctx.set_act_badarch(ScmpAction::KillProcess)?;

        // SAFETY: Do NOT add supported architectures to the filter.
        // This ensures Syd can never run a non-native system call,
        // which we do not need at all.
        // seccomp_add_architectures(&mut ctx)?;

        // Allow safe system calls.
        for sysname in EMU_SYSCALLS {
            match ScmpSyscall::from_name(sysname) {
                Ok(syscall) => ctx.add_rule(ScmpAction::Allow, syscall)?,
                Err(_) => {
                    crate::info!("ctx": "mon", "op": "allow_syscall",
                        "error": "scmp_error", "sys": sysname);
                }
            }
        }

        // Allow get id system calls.
        for sysname in GET_ID_SYSCALLS {
            match ScmpSyscall::from_name(sysname) {
                Ok(syscall) => ctx.add_rule(ScmpAction::Allow, syscall)?,
                Err(_) => {
                    crate::info!("ctx": "mon", "op": "allow_syscall",
                        "error": "scmp_error", "sys": sysname);
                }
            }
        }

        if safe_setid {
            // SAFETY: Main thread confines these further.
            // As these system calls as per-process,
            // the main thread's seccomp rules will apply
            // to us even without TSYNC.
            for sysname in SET_ID_SYSCALLS {
                match ScmpSyscall::from_name(sysname) {
                    Ok(syscall) => ctx.add_rule(ScmpAction::Allow, syscall)?,
                    Err(_) => {
                        crate::info!("ctx": "mon", "op": "allow_syscall",
                            "error": "scmp_error", "sys": sysname);
                    }
                }
            }

            // SAFETY: cap{g,s}et is necessary to drop CAP_SET{U,G}ID after
            // changing {U,G}ID.
            // SAFETY: We cannot deny _sigreturn_(2) with SafeSetID!
            // NPTL uses reserved-signals during UID/GID change!
            for sysname in [
                "capget",
                "capset",
                "sigaction",
                "rt_sigaction",
                "sigreturn",
                "rt_sigreturn",
            ] {
                match ScmpSyscall::from_name(sysname) {
                    Ok(syscall) => ctx.add_rule(ScmpAction::Allow, syscall)?,
                    Err(_) => {
                        crate::info!("ctx": "mon", "op": "allow_syscall",
                            "error": "scmp_error", "sys": sysname);
                    }
                }
            }
        }

        Ok(ctx)
    }
}

/// Type that exists to manage worker exit on panic.
///
/// This type is constructed once per `Worker` and implements `Drop` to
/// handle proper worker exit in case the worker panics when executing
/// the current task or anywhere else in its work loop. If the
/// `Sentinel` is dropped at the end of the worker's work loop and the
/// current thread is panicking, handle worker exit the same way as if
/// the task completed normally (if the worker panicked while executing
/// a submitted task) then clone the worker and start it with an initial
/// task of `None`.
struct Sentinel<'s> {
    is_working: bool,
    worker_ref: &'s Worker,
}

impl Sentinel<'_> {
    fn new(worker_ref: &Worker) -> Sentinel<'_> {
        Sentinel {
            is_working: false,
            worker_ref,
        }
    }
}

impl Drop for Sentinel<'_> {
    fn drop(&mut self) {
        if thread::panicking() {
            if self.is_working {
                // Busy thread panicked.
                self.worker_ref.worker_data.decrement_both();
            } else {
                // Idle thread panicked.
                self.worker_ref.worker_data.decrement_worker_total();
            }
        }
    }
}

const WORKER_BUSY_MASK: u32 = 0x0000FFFF;
const INCREMENT_TOTAL: u32 = 0x00010000;
const INCREMENT_BUSY: u32 = 0x00000001;

/// 1. Struct containing data shared between workers.
/// 2. Struct that stores and handles an `AtomicU32` that stores the
///    total worker count in the higher half of bits and the busy worker
///    count in the lower half of bits. This allows to to increment /
///    decrement both counters in a single atomic operation.
#[derive(Default)]
struct WorkerData(AtomicU32);

impl WorkerData {
    fn increment_both(&self) -> (u16, u16) {
        let old_val = self
            .0
            .fetch_add(INCREMENT_TOTAL | INCREMENT_BUSY, Ordering::Relaxed);
        Self::split(old_val)
    }

    fn decrement_both(&self) -> (u16, u16) {
        let old_val = self
            .0
            .fetch_sub(INCREMENT_TOTAL | INCREMENT_BUSY, Ordering::Relaxed);
        Self::split(old_val)
    }

    fn increment_worker_total(&self) -> u16 {
        let old_val = self.0.fetch_add(INCREMENT_TOTAL, Ordering::Relaxed);
        Self::total(old_val)
    }

    #[allow(dead_code)]
    fn decrement_worker_total(&self) -> u16 {
        let old_val = self.0.fetch_sub(INCREMENT_TOTAL, Ordering::Relaxed);
        Self::total(old_val)
    }

    fn increment_worker_busy(&self) -> u16 {
        let old_val = self.0.fetch_add(INCREMENT_BUSY, Ordering::Relaxed);
        Self::busy(old_val)
    }

    fn decrement_worker_busy(&self) -> u16 {
        let old_val = self.0.fetch_sub(INCREMENT_BUSY, Ordering::Relaxed);
        Self::busy(old_val)
    }

    fn get_total_count(&self) -> u16 {
        Self::total(self.0.load(Ordering::Relaxed))
    }

    fn get_busy_count(&self) -> u16 {
        Self::busy(self.0.load(Ordering::Relaxed))
    }

    #[inline]
    fn split(val: u32) -> (u16, u16) {
        let total_count = (val >> 16) as u16;
        let busy_count = (val & WORKER_BUSY_MASK) as u16;
        (total_count, busy_count)
    }

    #[inline]
    fn total(val: u32) -> u16 {
        (val >> 16) as u16
    }

    #[inline]
    fn busy(val: u32) -> u16 {
        (val & WORKER_BUSY_MASK) as u16
    }
}
