//! mmaptest -- checks mmap / munmap / mprotect / brk and demand paging.
//!
//!   mmaptest         run all the checks (exit code = failed checks)
//!   mmaptest crash   write to read-only memory without a handler: the
//!                    default action of SIGSEGV ends the program (139)
//!
//! Faults are caught with a SIGSEGV handler that records what happened
//! (address, SEGV_MAPERR = nothing mapped / SEGV_ACCERR = not allowed) and
//! then *fixes* it (mprotect / mmap the page), so the faulting instruction
//! can run again when the handler returns. Garbage collectors and
//! user-space pagers use the same trick.
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

use core::sync::atomic::{AtomicU64, Ordering::SeqCst};
use ulib::*;

const PAGE: u64 = 4096;
const SEGV_MAPERR: i32 = 1;
const SEGV_ACCERR: i32 = 2;

static FAULTS: AtomicU64 = AtomicU64::new(0);
static ADDR: AtomicU64 = AtomicU64::new(0);
static CODE: AtomicU64 = AtomicU64::new(0);

extern "C" fn on_segv(sig: i32, info: *const SigInfo, _context: *mut u8) {
    let info = unsafe { &*info };
    FAULTS.fetch_add(1, SeqCst);
    ADDR.store(info.addr_or_pid, SeqCst);
    CODE.store(info.code as u64, SeqCst);
    let page = info.addr_or_pid & !(PAGE - 1);
    // Fix it: make the page readable and writable (mapping it if needed).
    if sig != SIGSEGV || mprotect(page, PAGE, PROT_READ | PROT_WRITE) < 0 {
        mmap(page, PAGE, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
    }
}

fn peek(addr: u64) -> u8 {
    unsafe { (addr as *const u8).read_volatile() }
}
fn poke(addr: u64, value: u8) {
    unsafe { (addr as *mut u8).write_volatile(value) }
}

fn main() -> i64 {
    if arg(1) == Some("crash") {
        let p = mmap(0, PAGE, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0) as u64;
        println!("[mmaptest] writing to read-only memory at {:#x}...", p);
        poke(p, 1);
        println!("[mmaptest] not reached");
        return 1;
    }
    let mut c = Checks::new("mmaptest");
    let before = sysinfo().frames_in_use;

    // 1. Anonymous memory: zeroed, readable, writable.
    let a = mmap(0, 3 * PAGE, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0) as u64;
    c.check(a as i64 > 0 && a % PAGE == 0, "mmap 3 pages of anonymous memory");
    c.check((0..3 * PAGE).all(|i| peek(a + i) == 0), "it starts zeroed");
    for i in 0..3 * PAGE {
        poke(a + i, i as u8);
    }
    c.check((0..3 * PAGE).all(|i| peek(a + i) == i as u8), "it keeps what we write");

    // 2. Demand paging: a big mapping costs (almost) nothing until touched.
    let big = mmap(0, 64 << 20, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0) as u64;
    let after_map = sysinfo().frames_in_use;
    poke(big, 1);
    poke(big + (32 << 20), 2);
    poke(big + (64 << 20) - 1, 3);
    let after_touch = sysinfo().frames_in_use;
    c.check(big as i64 > 0 && after_map < before + 8, "mmap 64 MiB: no memory used yet (pages come on first touch)");
    c.check(after_touch >= after_map + 3 && after_touch < after_map + 16, "touching 3 pages of it uses ~3 frames");
    c.check(munmap(big, 64 << 20) == 0, "munmap the 64 MiB");

    // 3. Faults become signals our handler sees.
    c.check(sigaction_info(SIGSEGV, on_segv, 0) == 0, "install a SIGSEGV handler (SA_SIGINFO)");
    c.check(mprotect(a + PAGE, PAGE, PROT_READ) == 0, "mprotect the middle page read-only");
    c.check(peek(a + PAGE + 5) == 5 && FAULTS.load(SeqCst) == 0, "reading it still works");
    poke(a + PAGE + 7, 99);
    c.check(FAULTS.load(SeqCst) == 1 && ADDR.load(SeqCst) == a + PAGE + 7 && CODE.load(SeqCst) == SEGV_ACCERR as u64, "writing it: SIGSEGV, si_addr = that byte, SEGV_ACCERR");
    c.check(peek(a + PAGE + 7) == 99, "the handler's mprotect fixed it and the write was retried");

    c.check(munmap(a + PAGE, PAGE) == 0, "munmap the middle page");
    c.check(peek(a) == 0 && peek(a + 2 * PAGE + 1) == 1, "the pages around the hole are untouched");
    let v = peek(a + PAGE + 7);
    c.check(FAULTS.load(SeqCst) == 2 && CODE.load(SeqCst) == SEGV_MAPERR as u64 && v == 0, "reading the hole: SEGV_MAPERR; the handler maps a fresh (zero) page");
    c.check(mprotect(a, PAGE, PROT_NONE) == 0 && { poke(a, 4); FAULTS.load(SeqCst) == 3 }, "PROT_NONE: even reading faults");
    c.check(munmap(a, 3 * PAGE) == 0, "munmap all of it");

    // 4. A file mapped read-only (private): the same bytes read() gives.
    let fd = open("/mnt/disk0/readme.txt", 0);
    let mut buf = [0u8; 211];
    let n = read(fd, &mut buf);
    let f = mmap(0, 2 * PAGE, PROT_READ, MAP_PRIVATE, fd, 0) as u64;
    c.check(fd >= 0 && n == 211 && f as i64 > 0, "mmap /mnt/disk0/readme.txt");
    c.check((0..211).all(|i| peek(f + i) == buf[i as usize]), "the mapping has the file's bytes");
    c.check(peek(f + 211) == 0 && peek(f + PAGE - 1) == 0, "after the end of the file: zeros");
    let faults = FAULTS.load(SeqCst);
    poke(f, b'J');
    c.check(FAULTS.load(SeqCst) == faults + 1 && peek(f) == b'J', "writing it faults; after mprotect we write our private copy");
    let mut again = [0u8; 1];
    lseek(fd, 0, SEEK_SET);
    read(fd, &mut again);
    c.check(again[0] == b'H', "the file itself is unchanged");
    c.check(mmap(0, PAGE, PROT_READ, MAP_SHARED, fd, 0) < 0, "MAP_SHARED of a file isn't supported (an error, not silent loss)");
    munmap(f, 2 * PAGE);
    close(fd);

    // 5. brk: grow the heap by 1 MiB, use the end, shrink back.
    let start = brk(0);
    let end = brk(start + (1 << 20));
    c.check(end == start + (1 << 20), "brk grows the heap by 1 MiB");
    poke(end - 1, 7);
    c.check(peek(end - 1) == 7, "the new end is usable");
    c.check(brk(start) == start, "brk shrinks it back");

    // 6. Bad requests.
    c.check(mmap(0, 0, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0) == -22, "mmap of 0 bytes: EINVAL");
    c.check(mmap(0x1000, PAGE, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0) == -22, "MAP_FIXED outside user space: EINVAL");

    sigaction_raw(SIGSEGV, SIG_DFL, 0, 0);
    let leaked = sysinfo().frames_in_use as i64 - before as i64;
    c.check(leaked < 16, "memory use is back to (about) where it started");
    c.done()
}