//! Stage 19: static PIE. The ELF loader now also takes ET_DYN files with
//! no PT_INTERP (what stock `musl-gcc -static-pie` makes), loads them at a
//! random base and reports it in the auxiliary vector; dynamically linked
//! programs (PT_INTERP) get ENOEXEC; classic ET_EXEC still works.
#![no_std]
#![no_main]
#![feature(custom_test_frameworks)]
#![test_runner(rustbelt::test_runner)]
#![reexport_test_harness_main = "test_main"]

extern crate alloc;

use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec::Vec;
use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use rustbelt::errno::Errno;
use rustbelt::fs::file::File;
use rustbelt::process::{self, ExitStatus, SpawnArgs};
use rustbelt::{elf, fs, memory, thread};

entry_point!(main);

fn main(boot_info: &'static BootInfo) -> ! {
    rustbelt::init();
    rustbelt::init_memory(boot_info);
    rustbelt::init_smp(); // piethreads' threads spread over 4 CPUs
    fs::init();
    rustbelt::init_storage();
    test_main();
    rustbelt::hlt_loop();
}

/// Run /mnt/disk0/bin/`name` with `args`; its exit status and stdout.
fn run(name: &str, args: &[&str]) -> (ExitStatus, String) {
    let (reader, writer) = fs::pipe::new();
    let mut spawn = SpawnArgs::kernel(name);
    spawn.args.extend(args.iter().map(|a| String::from(*a)));
    spawn.stdio[1] = Arc::new(writer);
    let p = process::spawn_path(&alloc::format!("/mnt/disk0/bin/{}", name), spawn).expect("spawn failed");
    let mut out = Vec::new();
    let mut buf = [0u8; 256];
    loop {
        match reader.read(&mut buf).unwrap() {
            0 => break,
            n => out.extend_from_slice(&buf[..n]),
        }
    }
    (process::wait(&p), String::from_utf8_lossy(&out).into_owned())
}

fn counts() -> (usize, usize, usize) {
    thread::sleep_ticks(3);
    (thread::thread_count(), process::count(), memory::frames_in_use())
}

/// "hello" from a static PIE: its relocations, AT_PHDR/AT_ENTRY/AT_BASE.
#[test_case]
fn static_pie_hello() {
    let (status, out) = run("piehello", &[]);
    assert!(out.starts_with("hello from a static PIE\n"), "output: {}", out);
    assert_eq!(status, ExitStatus::Exited(0), "output: {}", out);
}

/// Every run lands at a different, page-aligned address in user space.
#[test_case]
fn static_pie_base_is_random() {
    let parse = |s: &str| u64::from_str_radix(s.trim().trim_start_matches("0x"), 16).unwrap();
    let (s1, a) = run("piehello", &["where"]);
    let (s2, b) = run("piehello", &["where"]);
    assert_eq!((s1, s2), (ExitStatus::Exited(0), ExitStatus::Exited(0)));
    let (a, b) = (parse(&a), parse(&b));
    assert_ne!(a, b, "same base twice");
    assert_eq!(a % 4096, b % 4096, "main moved within its page");
    assert!(a > memory::USER_START && a < rustbelt::mm::MMAP_TOP);
}

/// pthreads + thread-local storage (PT_TLS) + stdio + musl's ttyname on a
/// pty, as a static PIE. Nothing leaks.
#[test_case]
fn static_pie_threads_and_stdio() {
    let (status, out) = run("piethreads", &[]);
    assert_eq!(status, ExitStatus::Exited(0), "output: {}", out);
    let before = counts();
    let (status, out) = run("piethreads", &[]);
    assert_eq!(status, ExitStatus::Exited(0), "output: {}", out);
    assert!(out.contains("piethreads: 0 failed"));
    assert_eq!(counts(), before, "piethreads leaked something");
}

/// A dynamically linked program (PT_INTERP) is refused with ENOEXEC.
#[test_case]
fn dynamic_program_is_enoexec() {
    let image = fs::read_all("/mnt/disk0/bin/dynhello").unwrap();
    assert_eq!(elf::check(&image), Err(elf::ElfError::NeedsInterpreter));
    let err = process::spawn_path("/mnt/disk0/bin/dynhello", SpawnArgs::kernel("dynhello")).err();
    assert_eq!(err, Some(Errno::ENOEXEC));
}

/// The loader still takes the classic ET_EXEC files (Rust programs, the
/// no-libc C program, a musl ET_EXEC), and checks PIEs at their base.
#[test_case]
fn exec_files_still_load() {
    for name in ["hello", "rawabi", "vdsotest"] {
        let image = fs::read_all(&alloc::format!("/mnt/disk0/bin/{}", name)).unwrap();
        assert_eq!(u16::from_le_bytes([image[16], image[17]]), 2, "{} is not ET_EXEC", name);
        assert!(elf::check(&image).is_ok(), "{} doesn't load", name);
    }
    let pie = fs::read_all("/mnt/disk0/bin/piehello").unwrap();
    assert_eq!(u16::from_le_bytes([pie[16], pie[17]]), 3);
    let entry = elf::check(&pie).unwrap();
    assert!(entry > memory::USER_START + 0x40_0000, "PIE entry {:#x} not biased", entry);
    assert_eq!(run("rawabi", &["/tmp"]).0, ExitStatus::Exited(0));
}

#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    rustbelt::test_panic_handler(info)
}