foundry

//! A tiny "libc" for user programs: syscall wrappers, print!/println!,
//! the `_start` entry point and a panic handler.
//!
//! Each user program includes it with
//!     #[macro_use] #[path = "ulib.rs"] mod ulib;
//! and defines `fn main() -> i64` (the exit code).
//!
//! User programs are built for the `x86_64-unknown-none` target, which has
//! SSE/floating point disabled. That matters: the kernel doesn't save SSE
//! registers on a context switch.

#![allow(dead_code)]

use core::arch::asm;
use core::fmt;

pub const SYS_READ: u64 = 0;
pub const SYS_WRITE: u64 = 1;
pub const SYS_OPEN: u64 = 2;
pub const SYS_CLOSE: u64 = 3;
pub const SYS_STAT: u64 = 4;
pub const SYS_FSTAT: u64 = 5;
pub const SYS_LSEEK: u64 = 8;
pub const SYS_PIPE: u64 = 22;
pub const SYS_SCHED_YIELD: u64 = 24;
pub const SYS_SLEEP_MS: u64 = 35;
pub const SYS_GETPID: u64 = 39;
pub const SYS_EXIT: u64 = 60;
pub const SYS_WAIT: u64 = 61;
pub const SYS_GETCWD: u64 = 79;
pub const SYS_CHDIR: u64 = 80;
pub const SYS_MKDIR: u64 = 83;
pub const SYS_RMDIR: u64 = 84;
pub const SYS_UNLINK: u64 = 87;
pub const SYS_GETDENTS64: u64 = 217;
pub const SYS_SPAWN: u64 = 500;

pub const O_CREAT: u64 = 0o100;
pub const O_TRUNC: u64 = 0o1000;
pub const SEEK_SET: u64 = 0;
pub const SEEK_CUR: u64 = 1;
pub const SEEK_END: u64 = 2;
pub const SPAWN_TRACE: u64 = 1;

pub const ENOENT: i64 = 2;
pub const EBADF: i64 = 9;
pub const ECHILD: i64 = 10;
pub const EFAULT: i64 = 14;
pub const EBUSY: i64 = 16;
pub const EEXIST: i64 = 17;
pub const ENOTDIR: i64 = 20;
pub const EISDIR: i64 = 21;
pub const ESPIPE: i64 = 29;
pub const ENOTEMPTY: i64 = 39;

/// Same layout as the kernel's `fs::file::Stat`.
#[repr(C)]
#[derive(Default, Clone, Copy)]
pub struct Stat {
    pub kind: u64, // KIND_*
    pub size: u64,
}
pub const KIND_FILE: u64 = 1;
pub const KIND_DIR: u64 = 2;
pub const KIND_CONSOLE: u64 = 3;
pub const KIND_PIPE: u64 = 4;

/// Same layout as the kernel's `syscall::SpawnOptions`.
#[repr(C)]
pub struct SpawnOptions {
    pub stdin: i64, // my fd to give the child as its fd 0 (-1 = my fd 0)
    pub stdout: i64,
    pub stderr: i64,
    pub flags: u64, // SPAWN_TRACE
}

impl SpawnOptions {
    pub const INHERIT: SpawnOptions = SpawnOptions { stdin: -1, stdout: -1, stderr: -1, flags: 0 };
}

/// Raw system call: number in rax, args in rdi/rsi/rdx, result in rax.
/// The `syscall` instruction overwrites rcx and r11.
pub fn syscall3(nr: u64, a1: u64, a2: u64, a3: u64) -> i64 {
    let ret: i64;
    unsafe {
        asm!(
            "syscall",
            inlateout("rax") nr as i64 => ret,
            in("rdi") a1,
            in("rsi") a2,
            in("rdx") a3,
            lateout("rcx") _,
            lateout("r11") _,
            options(nostack),
        );
    }
    ret
}

/// 5-argument version (4th argument goes in r10, 5th in r8).
pub fn syscall5(nr: u64, a1: u64, a2: u64, a3: u64, a4: u64, a5: u64) -> i64 {
    let ret: i64;
    unsafe {
        asm!(
            "syscall",
            inlateout("rax") nr as i64 => ret,
            in("rdi") a1,
            in("rsi") a2,
            in("rdx") a3,
            in("r10") a4,
            in("r8") a5,
            lateout("rcx") _,
            lateout("r11") _,
            options(nostack),
        );
    }
    ret
}

pub fn read(fd: i64, buf: &mut [u8]) -> i64 {
    syscall3(SYS_READ, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64)
}
pub fn write(fd: i64, buf: &[u8]) -> i64 {
    syscall3(SYS_WRITE, fd as u64, buf.as_ptr() as u64, buf.len() as u64)
}
pub fn open(path: &str, flags: u64) -> i64 {
    syscall3(SYS_OPEN, path.as_ptr() as u64, path.len() as u64, flags)
}
pub fn close(fd: i64) -> i64 {
    syscall3(SYS_CLOSE, fd as u64, 0, 0)
}
pub fn stat(path: &str, st: &mut Stat) -> i64 {
    syscall3(SYS_STAT, path.as_ptr() as u64, path.len() as u64, st as *mut Stat as u64)
}
pub fn fstat(fd: i64, st: &mut Stat) -> i64 {
    syscall3(SYS_FSTAT, fd as u64, st as *mut Stat as u64, 0)
}
pub fn lseek(fd: i64, offset: i64, whence: u64) -> i64 {
    syscall3(SYS_LSEEK, fd as u64, offset as u64, whence)
}
/// Returns (read end, write end), or the negative errno.
pub fn pipe() -> Result<(i64, i64), i64> {
    let mut fds = [0u32; 2];
    match syscall3(SYS_PIPE, fds.as_mut_ptr() as u64, 0, 0) {
        0 => Ok((fds[0] as i64, fds[1] as i64)),
        e => Err(e),
    }
}
pub fn sched_yield() {
    syscall3(SYS_SCHED_YIELD, 0, 0, 0);
}
pub fn sleep_ms(ms: u64) {
    syscall3(SYS_SLEEP_MS, ms, 0, 0);
}
pub fn getpid() -> i64 {
    syscall3(SYS_GETPID, 0, 0, 0)
}
pub fn exit(code: i64) -> ! {
    syscall3(SYS_EXIT, code as u64, 0, 0);
    loop {} // not reached
}
/// Wait for child `pid` to end; returns its exit code (139 = killed).
pub fn wait(pid: i64) -> i64 {
    syscall3(SYS_WAIT, pid as u64, 0, 0)
}
/// Current directory into `buf`; returns its length (or -errno).
pub fn getcwd(buf: &mut [u8]) -> i64 {
    syscall3(SYS_GETCWD, buf.as_mut_ptr() as u64, buf.len() as u64, 0)
}
pub fn chdir(path: &str) -> i64 {
    syscall3(SYS_CHDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn mkdir(path: &str) -> i64 {
    syscall3(SYS_MKDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn rmdir(path: &str) -> i64 {
    syscall3(SYS_RMDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn unlink(path: &str) -> i64 {
    syscall3(SYS_UNLINK, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn getdents64(fd: i64, buf: &mut [u8]) -> i64 {
    syscall3(SYS_GETDENTS64, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64)
}
/// Start program `path` with arguments `argv` (NUL-separated, e.g.
/// "child\0hi"). Returns the child's pid (or -errno).
pub fn spawn(path: &str, argv: &str, opts: &SpawnOptions) -> i64 {
    syscall5(
        SYS_SPAWN,
        path.as_ptr() as u64,
        path.len() as u64,
        argv.as_ptr() as u64,
        argv.len() as u64,
        opts as *const SpawnOptions as u64,
    )
}

/// Walk the `linux_dirent64` records getdents64 put in `buf[..len]`,
/// calling `f(name, d_type)` for each (d_type: 4 = directory, 8 = file).
pub fn for_each_dirent(buf: &[u8], len: usize, mut f: impl FnMut(&str, u8)) {
    let mut pos = 0;
    while pos + 19 < len {
        let reclen = u16::from_le_bytes([buf[pos + 16], buf[pos + 17]]) as usize;
        let d_type = buf[pos + 18];
        let name = &buf[pos + 19..pos + reclen];
        let name = &name[..name.iter().position(|&b| b == 0).unwrap_or(name.len())];
        f(core::str::from_utf8(name).unwrap_or("?"), d_type);
        pos += reclen;
    }
}

// --- program arguments ---------------------------------------------------

/// `argv` as the kernel passes it: (pointer, length) pairs.
static mut ARGC: usize = 0;
static mut ARGV: *const [usize; 2] = core::ptr::null();

/// The program's arguments; `arg(0)` is the program name.
pub fn argc() -> usize {
    unsafe { ARGC }
}
pub fn arg(i: usize) -> Option<&'static str> {
    unsafe {
        if i >= ARGC {
            return None;
        }
        let [ptr, len] = *ARGV.add(i);
        core::str::from_utf8(core::slice::from_raw_parts(ptr as *const u8, len)).ok()
    }
}

pub struct Stdout;

impl fmt::Write for Stdout {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        write(1, s.as_bytes());
        Ok(())
    }
}

macro_rules! print {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let _ = write!($crate::ulib::Stdout, $($arg)*);
    }};
}

macro_rules! println {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let _ = writeln!($crate::ulib::Stdout, $($arg)*);
    }};
}

/// The kernel jumps here (see process.rs) with rdi = argc, rsi = argv.
/// Save the arguments, run main, exit with its result.
#[unsafe(no_mangle)]
pub extern "C" fn _start(argc: usize, argv: *const [usize; 2]) -> ! {
    unsafe {
        ARGC = argc;
        ARGV = argv;
    }
    exit(crate::main())
}

#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
    println!("[user panic] {}", info);
    exit(101)
}