pg-embed-setup-unpriv 0.5.2

//! Windows Job Object failsafe for process-exit cleanup.

use std::{ffi::c_void, ptr::NonNull};

use super::{
    PostmasterPid,
    PostmasterProcess,
    ProcessHandle,
    open_assignable_descendant_processes,
    process_tree,
};

const JOB_OBJECT_EXTENDED_LIMIT_INFORMATION: u32 = 9;
const JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE: u32 = 0x0000_2000;

#[link(name = "kernel32")]
unsafe extern "system" {
    fn CreateJobObjectW(attributes: *mut c_void, name: *const u16) -> *mut c_void;
    fn SetInformationJobObject(
        job: *mut c_void,
        info_class: u32,
        info: *mut c_void,
        info_length: u32,
    ) -> i32;
    fn AssignProcessToJobObject(job: *mut c_void, process: *mut c_void) -> i32;
    fn CloseHandle(handle: *mut c_void) -> i32;
}

pub(super) struct JobHandle(NonNull<c_void>);

// SAFETY: `JobHandle` owns a kernel handle. Windows permits closing and using
// job handles from any thread, and access is synchronized by the kernel.
unsafe impl Send for JobHandle {}

impl JobHandle {
    pub(super) fn create_for_process_tree(root: PostmasterProcess) -> Option<Self> {
        let Some(root_process) = ProcessHandle::open_assign_to_job(root.pid()) else {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                pid = root.pid(),
                "skipping Windows Job Object failsafe because root process could not be opened"
            );
            return None;
        };
        if !root_process.matches_postmaster(root) {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                pid = root.pid(),
                "skipping Windows Job Object failsafe because root identity changed"
            );
            return None;
        }

        let Some(job) = Self::create_kill_on_close(root.pid()) else {
            return None;
        };
        let assigned = job.assign_process_tree(root.pid(), &root_process);
        tracing::debug!(
            target: crate::observability::LOG_TARGET,
            pid = root.pid(),
            assigned,
            "prepared Windows Job Object failsafe"
        );
        assigned.then_some(job)
    }

    fn create_kill_on_close(root: PostmasterPid) -> Option<Self> {
        // SAFETY:
        // - null security attributes request the default descriptor.
        // - null name creates an unnamed private job.
        // - a null return is handled as failure.
        let raw_handle = unsafe { CreateJobObjectW(std::ptr::null_mut(), std::ptr::null()) };
        let Some(job) = NonNull::new(raw_handle).map(Self) else {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                root_pid = root,
                "failed to create Windows Job Object for process-exit failsafe"
            );
            return None;
        };
        job.enable_kill_on_close(root).then_some(job)
    }

    fn enable_kill_on_close(&self, root: PostmasterPid) -> bool {
        let mut info = JobObjectExtendedLimitInformation::kill_on_close();
        let Ok(info_length) =
            u32::try_from(std::mem::size_of::<JobObjectExtendedLimitInformation>())
        else {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                root_pid = root,
                "failed to size Windows Job Object limit information"
            );
            return false;
        };
        let info_ptr = std::ptr::addr_of_mut!(info).cast::<c_void>();

        // SAFETY:
        // - `self.0` is a non-null job handle owned by this wrapper.
        // - `info_ptr` points to an initialized job-information value with a valid byte length for
        //   this process architecture.
        // - the callee reads the buffer only for the duration of the call.
        let enabled = unsafe {
            SetInformationJobObject(
                self.0.as_ptr(),
                JOB_OBJECT_EXTENDED_LIMIT_INFORMATION,
                info_ptr,
                info_length,
            ) != 0
        };
        if !enabled {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                root_pid = root,
                "failed to enable kill-on-close for Windows Job Object"
            );
        }
        enabled
    }

    fn assign_process_tree(&self, root: PostmasterPid, root_process: &ProcessHandle) -> bool {
        let assigned_root = self.assign_process_handle(root_process);
        if !assigned_root {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                pid = root,
                "skipping Windows Job Object descendant assignment because root assignment failed"
            );
            return false;
        }

        let tree = process_tree(root);
        for process in open_assignable_descendant_processes(root, &tree) {
            self.assign_process(&process);
        }

        true
    }

    fn assign_process(&self, process: &ProcessHandle) -> bool {
        if !process.is_active_postgres() {
            tracing::debug!(
                target: crate::observability::LOG_TARGET,
                pid = process.pid(),
                "skipping Windows Job Object assignment because process is no longer active postgres"
            );
            return false;
        }
        self.assign_process_handle(process)
    }

    fn assign_process_handle(&self, process: &ProcessHandle) -> bool {
        // SAFETY:
        // - `self.0` is a valid job handle configured before assignment.
        // - `process` owns a process handle opened with the rights required by
        //   `AssignProcessToJobObject`.
        let assigned = unsafe { AssignProcessToJobObject(self.0.as_ptr(), process.raw()) != 0 };
        tracing::debug!(
            target: crate::observability::LOG_TARGET,
            pid = process.pid(),
            assigned,
            "attempted Windows Job Object process assignment"
        );
        assigned
    }
}

impl Drop for JobHandle {
    fn drop(&mut self) {
        // SAFETY: `JobHandle` owns the non-null handle returned by
        // `CreateJobObjectW`; closing it exactly once releases the OS resource.
        unsafe {
            CloseHandle(self.0.as_ptr());
        }
    }
}

#[repr(C)]
struct JobObjectBasicLimitInformation {
    per_process_user_time_limit: i64,
    per_job_user_time_limit: i64,
    limit_flags: u32,
    minimum_working_set_size: usize,
    maximum_working_set_size: usize,
    active_process_limit: u32,
    affinity: usize,
    priority_class: u32,
    scheduling_class: u32,
}

#[repr(C)]
struct IoCounters {
    read_operation_count: u64,
    write_operation_count: u64,
    other_operation_count: u64,
    read_transfer_count: u64,
    write_transfer_count: u64,
    other_transfer_count: u64,
}

#[repr(C)]
struct JobObjectExtendedLimitInformation {
    basic_limit_information: JobObjectBasicLimitInformation,
    io_info: IoCounters,
    process_memory_limit: usize,
    job_memory_limit: usize,
    peak_process_memory_used: usize,
    peak_job_memory_used: usize,
}

impl JobObjectExtendedLimitInformation {
    fn kill_on_close() -> Self {
        Self {
            basic_limit_information: JobObjectBasicLimitInformation {
                per_process_user_time_limit: 0,
                per_job_user_time_limit: 0,
                limit_flags: JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
                minimum_working_set_size: 0,
                maximum_working_set_size: 0,
                active_process_limit: 0,
                affinity: 0,
                priority_class: 0,
                scheduling_class: 0,
            },
            io_info: IoCounters {
                read_operation_count: 0,
                write_operation_count: 0,
                other_operation_count: 0,
                read_transfer_count: 0,
                write_transfer_count: 0,
                other_transfer_count: 0,
            },
            process_memory_limit: 0,
            job_memory_limit: 0,
            peak_process_memory_used: 0,
            peak_job_memory_used: 0,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::{
        ProcessEntry,
        ProcessTreeIndex,
        collect_process_tree,
        descendant_is_still_in_root_tree,
        process_has_root_ancestor,
    };

    fn entry(process_id: u32, parent_process_id: u32) -> ProcessEntry {
        ProcessEntry {
            process_id,
            parent_process_id,
        }
    }

    fn process_ids(entries: &[ProcessEntry]) -> Vec<u32> {
        entries.iter().map(|entry| entry.process_id).collect()
    }

    #[test]
    fn collect_process_tree_returns_root_when_no_descendants() {
        let entries = [entry(20, 10), entry(30, 20)];

        assert_eq!(process_ids(&collect_process_tree(99, &entries)), vec![99]);
    }

    #[test]
    fn collect_process_tree_includes_nested_descendants() {
        let entries = [entry(40, 30), entry(20, 10), entry(30, 20), entry(50, 99)];

        assert_eq!(
            process_ids(&collect_process_tree(10, &entries)),
            vec![10, 20, 30, 40]
        );
    }

    #[test]
    fn descendant_validation_rejects_reused_pid_from_different_tree() {
        let descendant = entry(20, 10);
        let entries = [entry(20, 99), entry(99, 99)];
        let index = ProcessTreeIndex::new(&entries);

        assert!(!descendant_is_still_in_root_tree(10, descendant, &index));
    }

    #[test]
    fn descendant_validation_rejects_parent_cycles() {
        let entries = [entry(20, 30), entry(30, 20)];
        let index = ProcessTreeIndex::new(&entries);

        assert!(!process_has_root_ancestor(10, 20, &index));
    }
}