rust-kernel

//! A tiny "libc" for user programs: syscall wrappers, print!/println!,
//! eprintln!, a small heap (so `Vec` and `String` work), 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, unused_macros)]

use core::alloc::{GlobalAlloc, Layout};
use core::arch::asm;
use core::cell::UnsafeCell;
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_IOCTL: u64 = 16;
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_UNAME: u64 = 63;
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 SYS_MOUNTS: u64 = 501;

pub const O_CREAT: u64 = 0o100;
pub const O_TRUNC: u64 = 0o1000;
pub const O_APPEND: u64 = 0o2000;
pub const TTY_SET_RAW: u64 = 1;
pub const WNOHANG: u64 = 1;
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 EAGAIN: i64 = 11;
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)
}
/// Like `wait`, but don't block: -EAGAIN while the child is still running.
pub fn try_wait(pid: i64) -> i64 {
    syscall3(SYS_WAIT, pid as u64, WNOHANG, 0)
}
/// Console raw mode on/off (keys arrive one by one, no echo). Returns the
/// old mode, or -ENOTTY if `fd` isn't the console.
pub fn set_raw(fd: i64, raw: bool) -> i64 {
    syscall3(SYS_IOCTL, fd as u64, TTY_SET_RAW, raw as u64)
}
/// Linux's `struct utsname`: sysname, nodename, release, version, machine,
/// domainname. Returns them as strings.
pub fn uname(buf: &mut [u8; 6 * 65]) -> [&str; 6] {
    syscall3(SYS_UNAME, buf.as_mut_ptr() as u64, 0, 0);
    core::array::from_fn(|i| {
        let field = &buf[i * 65..(i + 1) * 65];
        let len = field.iter().position(|&b| b == 0).unwrap_or(65);
        core::str::from_utf8(&field[..len]).unwrap_or("?")
    })
}
/// The mount table as text ("path fstype" lines); returns the length.
pub fn mounts(buf: &mut [u8]) -> i64 {
    syscall3(SYS_MOUNTS, buf.as_mut_ptr() as u64, buf.len() as u64, 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()
    }
}

/// Write all of `buf` (pipes and files may take less than all at once).
pub fn write_all(fd: i64, mut buf: &[u8]) -> i64 {
    while !buf.is_empty() {
        let n = write(fd, buf);
        if n <= 0 {
            return n;
        }
        buf = &buf[n as usize..];
    }
    0
}

/// Call `f` for each line read from `fd` (without the '\n'). Lines longer
/// than 1 KiB come in pieces. Returns 0, or -errno if a read fails.
pub fn for_each_line(fd: i64, mut f: impl FnMut(&[u8])) -> i64 {
    let mut line = [0u8; 1024];
    let mut len = 0;
    let mut buf = [0u8; 512];
    loop {
        let n = read(fd, &mut buf);
        if n < 0 {
            return n;
        }
        if n == 0 {
            if len > 0 {
                f(&line[..len]); // last line without a newline
            }
            return 0;
        }
        for &b in &buf[..n as usize] {
            if b == b'\n' || len == line.len() {
                f(&line[..len]);
                len = 0;
                if b == b'\n' {
                    continue;
                }
            }
            line[len] = b;
            len += 1;
        }
    }
}

/// `ls`, `cat` etc. print errors like "cat: x: No such file or directory".
pub fn strerror(errno: i64) -> &'static str {
    match errno.abs() {
        2 => "No such file or directory",
        5 => "I/O error",
        7 => "Argument list too long",
        8 => "Exec format error",
        9 => "Bad file descriptor",
        10 => "No child processes",
        11 => "Try again",
        12 => "Out of memory",
        14 => "Bad address",
        16 => "Device or resource busy",
        17 => "File exists",
        20 => "Not a directory",
        21 => "Is a directory",
        22 => "Invalid argument",
        24 => "Too many open files",
        25 => "Not a terminal",
        28 => "No space left on device",
        29 => "Illegal seek",
        32 => "Broken pipe",
        34 => "Result too large",
        36 => "File name too long",
        38 => "Function not implemented",
        39 => "Directory not empty",
        _ => "Unknown error",
    }
}

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)*);
    }};
}

pub struct Stderr;

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

/// Write formatted text to any fd: `fprint!(fd, "...")`.
pub struct Fd(pub i64);

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

macro_rules! eprintln {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let _ = writeln!($crate::ulib::Stderr, $($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)
}

// --- heap ----------------------------------------------------------------
//
// A tiny allocator so programs can use `Vec`, `String`, `format!` (with
// `extern crate alloc;`). Memory comes from a 64 KiB array in .bss, which
// the ELF loader maps as zeroed pages -- no `brk`/`mmap` syscall needed.
// Same idea as the kernel's fixed-size-block allocator: sizes are rounded up
// to a power of two (16 B .. 32 KiB) and freed blocks go on one free list
// per size, to be reused by the next allocation of that size. Programs are
// single-threaded, so no lock.

const HEAP_SIZE: usize = 64 * 1024;
const CLASSES: usize = 12; // 16 << 11 = 32 KiB

#[repr(C, align(4096))]
struct Heap {
    memory: UnsafeCell<[u8; HEAP_SIZE]>,
    used: UnsafeCell<usize>,               // bump pointer: bytes handed out so far
    free: UnsafeCell<[usize; CLASSES]>,    // free list heads (addresses, 0 = empty)
}

unsafe impl Sync for Heap {} // single-threaded programs

#[global_allocator]
static HEAP: Heap = Heap {
    memory: UnsafeCell::new([0; HEAP_SIZE]),
    used: UnsafeCell::new(0),
    free: UnsafeCell::new([0; CLASSES]),
};

/// Size class for a layout: 0 = 16 bytes, 1 = 32, ...
fn class_of(layout: Layout) -> usize {
    let size = layout.size().max(layout.align()).max(16).next_power_of_two();
    size.trailing_zeros() as usize - 4
}

unsafe impl GlobalAlloc for Heap {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let class = class_of(layout);
        if class >= CLASSES {
            return core::ptr::null_mut(); // too big: "out of memory"
        }
        unsafe {
            let free = &mut *self.free.get();
            if free[class] != 0 {
                // Pop a freed block; its first 8 bytes hold the next one.
                let block = free[class] as *mut u8;
                free[class] = *(block as *const usize);
                return block;
            }
            // Carve a new block. Blocks are aligned to their size, which
            // covers any alignment up to that size.
            let size = 16 << class;
            let used = &mut *self.used.get();
            let start = (*used).next_multiple_of(size);
            if start + size > HEAP_SIZE {
                return core::ptr::null_mut();
            }
            *used = start + size;
            (self.memory.get() as *mut u8).add(start)
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let class = class_of(layout);
        unsafe {
            let free = &mut *self.free.get();
            *(ptr as *mut usize) = free[class];
            free[class] = ptr as usize;
        }
    }
}