rust-kernel

//! Processes = a thread + its own address space + a file descriptor table
//! (+ arguments, a current directory and a parent).
//!
//! Starting one (`spawn_image`):
//!   1. new level-4 page table sharing the kernel mappings (memory.rs)
//!   2. load the ELF program's segments into its user part (elf.rs)
//!   3. map a user stack just below `USER_END`
//!   4. fd table: 0, 1, 2 = whatever the parent chose (console by default,
//!      or e.g. a pipe -- that's how the shell will do `a | b` and `> file`)
//!   5. put it in the process table and spawn a thread with that CR3. The
//!      thread starts in the kernel (`user_thread_entry`), copies the
//!      arguments onto the user stack and "returns" into ring 3 with `iretq`.
//!
//! Why `spawn(path, argv)` instead of Unix `fork()` + `exec()`? fork copies
//! the whole address space (or needs copy-on-write page tables to be fast),
//! and then exec throws the copy away again. spawn ("create a new process
//! running this program") is one simple step, and it's what Windows
//! (CreateProcess) and POSIX's own `posix_spawn` offer. Redirection, which
//! Unix shells do between fork and exec, is done with spawn's stdio options.
//!
//! Exit and wait: a process's exit status is stored in it and its `exited`
//! wait queue is woken; the parent's `wait(pid)` blocks on that queue, waits
//! until the thread is completely gone, then removes the child from the
//! process table ("reaps" it -- until then it's what Unix calls a zombie).
//!
//! Memory: a process's frames and page tables are freed when the last
//! `Arc<Process>` goes away (`impl Drop`). That's always after its thread is
//! gone: the thread holds an Arc until the scheduler frees it, which happens
//! only once the CPU has switched to another thread and address space. So
//! no CPU can still be using the page tables we free.
//!
//! Orphans: when a parent exits first, its running children become orphans
//! (parent = ORPHAN); nobody will wait for them, so they remove themselves
//! from the table when they exit. Unix hands them to `init` instead.
//!
//! Not done yet: signals/kill, more than one thread per process.

use crate::elf::{self, ElfError};
use crate::errno::Errno;
use crate::fs::{self, file::File};
use crate::memory::{self, USER_END};
use crate::thread::{self, WaitQueue};
use crate::{serial_println, tty::Console};
use alloc::{collections::BTreeMap, string::String, sync::Arc, vec, vec::Vec};
use core::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
use spin::Mutex;
use x86_64::VirtAddr;
use x86_64::instructions::interrupts::without_interrupts;
use x86_64::structures::paging::{Page, PageTableFlags, PhysFrame, Size4KiB};

/// The user stack: 16 KiB, ending one (unmapped) page below USER_END.
pub const USER_STACK_TOP: u64 = USER_END - 0x1000;
pub const USER_STACK_SIZE: u64 = 16 * 1024;
/// Below the stack: a 64 KiB guard zone that is never mapped, so a runaway
/// recursion hits it and gets killed with "stack overflow" (instead of
/// scribbling over the program's data). Big, so that one large stack frame
/// can't jump over it.
pub const USER_STACK_GUARD: u64 = 64 * 1024;
/// Lowest address of stack + guard: program segments must end below this.
pub const USER_STACK_ZONE_START: u64 = USER_STACK_TOP - USER_STACK_SIZE - USER_STACK_GUARD;
/// `parent` of a process whose parent exited first.
pub const ORPHAN: u64 = u64::MAX;

/// Is user address `addr` in the guard zone under the user stack?
pub fn in_user_stack_guard(addr: u64) -> bool {
    (USER_STACK_ZONE_START..USER_STACK_TOP - USER_STACK_SIZE).contains(&addr)
}
/// Max open files per process.
const MAX_FDS: usize = 16;
/// Max total size of a program's arguments (they go on its 16 KiB stack).
pub const MAX_ARGS_SIZE: usize = 2048;
/// Exit code `wait` reports for a process the kernel killed (like a shell
/// shows 128 + SIGSEGV(11) for a segfault).
pub const KILLED_EXIT_CODE: i64 = 139;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExitStatus {
    /// Called exit(code).
    Exited(i64),
    /// Killed by the kernel (e.g. "page fault").
    Killed(&'static str),
}

impl ExitStatus {
    /// The number `wait` returns.
    pub fn code(self) -> i64 {
        match self {
            ExitStatus::Exited(code) => code,
            ExitStatus::Killed(_) => KILLED_EXIT_CODE,
        }
    }
}

pub struct Process {
    pub pid: u64,
    /// Parent pid; 0 = started by the kernel itself, ORPHAN = parent exited.
    pub parent: AtomicU64,
    pub name: String,
    /// Print every syscall this process makes (`spawn` flag; `run --trace`).
    pub trace: bool,
    /// Current directory (always a normalized absolute path).
    pub cwd: Mutex<String>,
    entry: u64,
    /// Level-4 page table of its address space (freed in `drop`).
    l4: PhysFrame,
    /// Its thread's id (set right after the thread is created).
    tid: AtomicUsize,
    /// CPUs its thread has run on (bit per CPU; kept up by the scheduler).
    pub ran_on: AtomicU32,
    args: Vec<String>,
    fds: Mutex<Vec<Option<Arc<dyn File>>>>,
    /// Locked with interrupts off only (it's read in a wait_until condition).
    status: Mutex<Option<ExitStatus>>,
    exited: WaitQueue,
}

static NEXT_PID: AtomicU64 = AtomicU64::new(1);
/// All processes that haven't been reaped yet, by pid. Locked with
/// interrupts off only (exit_current takes it with interrupts off).
static PROCESSES: Mutex<BTreeMap<u64, Arc<Process>>> = Mutex::new(BTreeMap::new());

/// Everything needed to start a process besides the program itself.
pub struct SpawnArgs {
    pub args: Vec<String>, // argv; args[0] is the program name by convention
    pub stdio: [Arc<dyn File>; 3],
    pub cwd: String,
    pub parent: u64,
    pub trace: bool,
}

impl SpawnArgs {
    /// Defaults for a process started by the kernel: console stdio, cwd "/".
    pub fn kernel(name: &str) -> SpawnArgs {
        let console: Arc<dyn File> = Arc::new(Console);
        SpawnArgs {
            args: vec![String::from(name)],
            stdio: [console.clone(), console.clone(), console],
            cwd: String::from("/"),
            parent: 0,
            trace: false,
        }
    }
}

/// Load the program at `path` (absolute) from the VFS and start it.
pub fn spawn_path(path: &str, spawn: SpawnArgs) -> Result<Arc<Process>, Errno> {
    let image = fs::read_all(path)?; // ENOENT, EISDIR, EIO, ...
    let name = path.rsplit('/').next().unwrap_or(path);
    spawn_image(name, &image, spawn).map_err(|e| match e {
        ElfError::OutOfMemory => Errno::ENOMEM,
        _ => Errno::ENOEXEC, // "Exec format error"
    })
}

/// Start ELF program `elf_image` from the kernel (console stdio, no args).
pub fn spawn(name: &str, elf_image: &[u8]) -> Result<Arc<Process>, ElfError> {
    spawn_image(name, elf_image, SpawnArgs::kernel(name))
}

/// Load `elf_image` into a new address space and start it.
pub fn spawn_image(name: &str, elf_image: &[u8], spawn: SpawnArgs) -> Result<Arc<Process>, ElfError> {
    let l4 = memory::new_address_space().ok_or(ElfError::OutOfMemory)?;
    match load_into(elf_image, l4) {
        Ok(entry) => Ok(start_process(name, l4, entry, spawn)),
        Err(e) => {
            // Nobody runs in this address space yet: give its memory back.
            unsafe { memory::free_address_space(l4) };
            Err(e)
        }
    }
}

/// Load the program and map its stack; returns the entry point.
fn load_into(elf_image: &[u8], l4: PhysFrame) -> Result<u64, ElfError> {
    let entry = elf::load(elf_image, l4)?;
    let mut stack_flags = PageTableFlags::WRITABLE;
    if memory::nx_enabled() {
        stack_flags |= PageTableFlags::NO_EXECUTE;
    }
    let first = Page::<Size4KiB>::containing_address(VirtAddr::new(USER_STACK_TOP - USER_STACK_SIZE));
    let last = Page::<Size4KiB>::containing_address(VirtAddr::new(USER_STACK_TOP - 1));
    for page in Page::range_inclusive(first, last) {
        memory::map_user_page(l4, page, stack_flags).ok_or(ElfError::OutOfMemory)?;
    }
    Ok(entry)
}

/// Create the `Process` and its thread.
fn start_process(name: &str, l4: PhysFrame, entry: u64, spawn: SpawnArgs) -> Arc<Process> {
    let [stdin, stdout, stderr] = spawn.stdio;
    let process = Arc::new(Process {
        pid: NEXT_PID.fetch_add(1, Ordering::Relaxed),
        parent: AtomicU64::new(spawn.parent),
        name: String::from(name),
        trace: spawn.trace,
        cwd: Mutex::new(spawn.cwd),
        entry,
        l4,
        tid: AtomicUsize::new(usize::MAX),
        ran_on: AtomicU32::new(0),
        args: spawn.args,
        fds: Mutex::new(vec![Some(stdin), Some(stdout), Some(stderr)]),
        status: Mutex::new(None),
        exited: WaitQueue::new(),
    });
    without_interrupts(|| PROCESSES.lock().insert(process.pid, process.clone()));
    serial_println!("[process] started pid {} ({}), entry {:#x}", process.pid, name, entry);
    let tid = thread::spawn_with(name, user_thread_entry, l4, Some(process.clone()));
    process.tid.store(tid, Ordering::Release);
    process
}

/// First code of every process thread (still ring 0, but already running
/// with the process's CR3 and kernel stack).
fn user_thread_entry() {
    let process = thread::current_process().expect("process thread without process");
    // Copy the arguments onto the top of the user stack (we can write there
    // directly: this process's address space is active). Layout, top down:
    //   the argument strings' bytes (not NUL-terminated)
    //   argv: [ptr, len] pairs, one per argument   <- rsi
    //   (stack continues below)
    // and rdi = argc. Rust's `&str` is ptr+len, so pairs are handier than
    // C's NUL-terminated `char **argv`.
    let mut sp = USER_STACK_TOP;
    let mut pairs: Vec<[u64; 2]> = Vec::with_capacity(process.args.len());
    for arg in &process.args {
        sp -= arg.len() as u64;
        unsafe { core::ptr::copy_nonoverlapping(arg.as_ptr(), sp as *mut u8, arg.len()) };
        pairs.push([sp, arg.len() as u64]);
    }
    sp &= !0xf; // 16-byte align
    sp -= 16 * pairs.len() as u64;
    let argv = sp;
    unsafe { core::ptr::copy_nonoverlapping(pairs.as_ptr(), argv as *mut [u64; 2], pairs.len()) };
    // Functions expect rsp % 16 == 8 on entry (as if a return address had
    // just been pushed), so start 8 bytes below the 16-aligned spot.
    let (entry, argc) = (process.entry, pairs.len() as u64);
    drop(pairs);
    drop(process); // enter_user_mode never returns: don't leak the Arc
    unsafe { enter_user_mode(entry, sp - 8, argc, argv) }
}

/// Jump to ring 3. `iretq` pops RIP, CS, RFLAGS, RSP, SS -- we push a fake
/// interrupt frame with user selectors (RPL 3), so "returning" from it drops
/// the privilege level. RFLAGS = 0x202: interrupts enabled (+ the always-1 bit).
/// rdi/rsi carry `_start(argc, argv)`'s arguments.
unsafe fn enter_user_mode(rip: u64, rsp: u64, argc: u64, argv: u64) -> ! {
    let s = crate::gdt::selectors();
    unsafe {
        core::arch::asm!(
            "push {ss}",
            "push {rsp}",
            "push {rflags}",
            "push {cs}",
            "push {rip}",
            // Don't leak kernel values to user space in registers.
            "xor eax, eax", "xor ebx, ebx", "xor ecx, ecx", "xor edx, edx",
            "xor ebp, ebp",
            "xor r8d, r8d", "xor r9d, r9d", "xor r10d, r10d", "xor r11d, r11d",
            "xor r12d, r12d", "xor r13d, r13d", "xor r14d, r14d", "xor r15d, r15d",
            "iretq",
            ss = in(reg) s.user_data.0 as u64,
            rsp = in(reg) rsp,
            rflags = in(reg) 0x202u64,
            cs = in(reg) s.user_code.0 as u64,
            rip = in(reg) rip,
            in("rdi") argc,
            in("rsi") argv,
            options(noreturn)
        );
    }
}

impl Process {
    /// The open file behind `fd`.
    pub fn file(&self, fd: u64) -> Result<Arc<dyn File>, Errno> {
        let fds = self.fds.lock();
        fds.get(fd as usize).and_then(|f| f.clone()).ok_or(Errno::EBADF)
    }

    /// Put `file` into the lowest free fd slot (like Unix).
    pub fn add_file(&self, file: Arc<dyn File>) -> Result<u64, Errno> {
        let mut fds = self.fds.lock();
        if let Some(i) = fds.iter().position(|f| f.is_none()) {
            fds[i] = Some(file);
            return Ok(i as u64);
        }
        if fds.len() >= MAX_FDS {
            return Err(Errno::EMFILE);
        }
        fds.push(Some(file));
        Ok(fds.len() as u64 - 1)
    }

    pub fn close(&self, fd: u64) -> Result<(), Errno> {
        let file = {
            let mut fds = self.fds.lock();
            match fds.get_mut(fd as usize) {
                Some(slot @ Some(_)) => slot.take(),
                _ => return Err(Errno::EBADF),
            }
        };
        drop(file); // dropping the last Arc closes the file (after unlocking)
        Ok(())
    }

    pub fn status(&self) -> Option<ExitStatus> {
        without_interrupts(|| *self.status.lock())
    }

    /// Make `path` absolute, relative to this process's current directory.
    pub fn resolve(&self, path: &str) -> Result<String, Errno> {
        let cwd = self.cwd.lock().clone();
        fs::normalize(&cwd, path)
    }
}

impl Drop for Process {
    /// The last reference is gone, so its thread is gone too: give the
    /// address space's frames and page tables back.
    fn drop(&mut self) {
        unsafe { memory::free_address_space(self.l4) };
    }
}

/// Look up a live (or not yet reaped) process.
pub fn get(pid: u64) -> Option<Arc<Process>> {
    without_interrupts(|| PROCESSES.lock().get(&pid).cloned())
}

/// How many processes are in the table (running or not yet reaped).
pub fn count() -> usize {
    without_interrupts(|| PROCESSES.lock().len())
}

/// End the current process (exit syscall, or killed by an exception).
pub fn exit_current(status: ExitStatus) -> ! {
    if let Some(process) = thread::current_process() {
        // Close all files (a pipe's reader sees EOF once its writers are gone).
        let files = core::mem::take(&mut *process.fds.lock());
        drop(files);
        crate::tty::process_exited(process.pid);
        // Serial only: on screen the shell reports exit codes itself.
        match status {
            ExitStatus::Exited(code) => {
                serial_println!("[process] pid {} ({}) exited with code {}", process.pid, process.name, code)
            }
            ExitStatus::Killed(why) => {
                crate::kprintln!("[process] pid {} ({}) killed: {}", process.pid, process.name, why)
            }
        }
        // No preemption from here on: we're about to disappear anyway.
        x86_64::instructions::interrupts::disable();
        let removed = {
            // The table lock makes "set my status" and "look at my parent"
            // one step, so a parent exiting at the same time on another CPU
            // either sees our status (and reaps us) or orphans us first (and
            // we reap ourselves). Never both, never neither.
            let mut table = PROCESSES.lock();
            *process.status.lock() = Some(status);
            let mut removed = Vec::new();
            let children: Vec<Arc<Process>> =
                table.values().filter(|c| c.parent.load(Ordering::Relaxed) == process.pid).cloned().collect();
            for child in children {
                if child.status().is_some() {
                    removed.extend(table.remove(&child.pid)); // exited: reap it
                } else {
                    child.parent.store(ORPHAN, Ordering::Relaxed); // still running
                }
            }
            if process.parent.load(Ordering::Relaxed) == ORPHAN {
                removed.extend(table.remove(&process.pid)); // nobody will wait for us
            }
            removed
        };
        drop(removed); // (frees their memory if they're completely gone)
        process.exited.wake_all(); // the parent may be blocked in wait()
        drop(process);
    }
    thread::exit_current()
}

/// Block (the calling thread) until `process` has ended and its thread is
/// gone, then remove it from the process table. Used by the kernel and by
/// the `wait` syscall. Once the caller drops its own `Arc`, all of the
/// process's memory is free again.
pub fn wait(process: &Process) -> ExitStatus {
    let status = process.exited.wait_until(|| *process.status.lock());
    reap(process);
    status
}

/// Wait for the (exited) process's thread to be freed, then remove the
/// process from the table.
fn reap(process: &Process) {
    // (The tid is stored right after the thread is created; on another CPU
    // the thread may have run and exited before that.)
    let tid = loop {
        match process.tid.load(Ordering::Acquire) {
            usize::MAX => thread::yield_now(),
            tid => break tid,
        }
    };
    thread::join(tid);
    let removed = without_interrupts(|| PROCESSES.lock().remove(&process.pid));
    drop(removed);
}

/// `wait` flag: don't block; `EAGAIN` if the child is still running.
pub const WNOHANG: u64 = 1;

/// The `wait(pid, flags)` syscall: only a process's parent may wait for it.
pub fn wait_child(parent: u64, pid: u64, flags: u64) -> Result<ExitStatus, Errno> {
    let child = get(pid).filter(|c| c.parent.load(Ordering::Relaxed) == parent).ok_or(Errno::ECHILD)?;
    if flags & WNOHANG != 0 {
        let status = child.status().ok_or(Errno::EAGAIN)?;
        reap(&child);
        return Ok(status);
    }
    Ok(wait(&child))
}