//! Adversarial memory-isolation tests. Each one is an attack a process
//! with no special capabilities could try; the kernel must refuse it.

use super::*;
use crate::{SyscallOutcome, SyscallResult};
use domain::syscall::Syscall;
use types::address::{DataRef, VirtualAddress};
use types::bounded::{CapSlot, PageCount};
use types::id::ThreadId;
use types::mapping::PagePerms;

const PAGE: usize = 4096;
// Mirrors the user window the mock platform exposes.
const USER_VA_END: usize = 0x0200_0000;

fn pages(n: u16) -> PageCount {
    PageCount::new(n).unwrap()
}

fn request(k: &mut Kernel<MockPlatform>, n: u16) -> Option<usize> {
    match k.handle_syscall(Syscall::RequestPages { count: pages(n) }) {
        SyscallOutcome::Complete(SyscallResult::Value(va)) => Some(va),
        _ => None,
    }
}

fn release(k: &mut Kernel<MockPlatform>, va: usize, n: u16) -> bool {
    matches!(
        k.handle_syscall(Syscall::ReleasePages {
            addr: VirtualAddress::from_raw(va),
            count: pages(n),
        }),
        SyscallOutcome::Complete(SyscallResult::Ok)
    )
}

fn map_page(k: &mut Kernel<MockPlatform>, process: CapSlot, src: usize, dst: usize) -> bool {
    matches!(
        k.handle_syscall(Syscall::MapPage {
            process,
            src_va: VirtualAddress::from_raw(src),
            dst_va: VirtualAddress::from_raw(dst),
            executable: false,
        }),
        SyscallOutcome::Complete(SyscallResult::Ok)
    )
}

/// Boots init (A) and creates a second process B with one runnable thread.
/// Returns A's process cap for B.
fn two_processes() -> (
    Kernel<MockPlatform>,
    ProcessId,
    ThreadId,
    ProcessId,
    ThreadId,
    CapSlot,
) {
    let (mut k, pid_a) = kernel_with_process();
    let tid_a = k.current_thread().unwrap();
    let cap_b = match k.handle_syscall(Syscall::CreateProcess {
        create_cap: CapSlot::new(0).unwrap(),
        image: DataRef::empty(),
    }) {
        SyscallOutcome::Complete(SyscallResult::Slot(s)) => s,
        _ => panic!("create process"),
    };
    let pid_b = ProcessId::from_raw(1);
    let tid_b = k.spawn_thread_for_test(pid_b, Priority::DEFAULT, TimeSlice::DEFAULT);
    (k, pid_a, tid_a, pid_b, tid_b, cap_b)
}

#[test]
fn release_with_inflated_count_cannot_free_another_processes_pages() {
    let (mut k, pid_a, tid_a, pid_b, tid_b, _) = two_processes();

    let va_a = request(&mut k, 1).unwrap();

    k.switch_to_thread_for_test(tid_b);
    let va_b = request(&mut k, 8).unwrap();
    let victim: [usize; 8] =
        core::array::from_fn(|i| k.platform_mut().mapping(pid_b, va_b + i * PAGE).unwrap());

    k.switch_to_thread_for_test(tid_a);
    assert!(
        !release(&mut k, va_a, 9),
        "released more pages than it owns"
    );

    let va_new = request(&mut k, 8).unwrap();
    for i in 0..8 {
        let pa = k.platform_mut().mapping(pid_a, va_new + i * PAGE).unwrap();
        assert!(!victim.contains(&pa), "A was handed B's live page {pa:#x}");
    }
}

#[test]
fn released_page_is_unmapped_from_caller() {
    let (mut k, pid) = kernel_with_process();
    let va = request(&mut k, 1).unwrap();
    assert!(release(&mut k, va, 1));
    assert_eq!(k.platform_mut().mapping(pid, va), None);
}

#[test]
fn page_cannot_be_released_twice() {
    let (mut k, _) = kernel_with_process();
    let va = request(&mut k, 1).unwrap();
    assert!(release(&mut k, va, 1));
    let free = k.free_page_count();
    assert!(!release(&mut k, va, 1));
    assert_eq!(k.free_page_count(), free);
}

#[test]
fn killing_a_process_after_release_does_not_double_free() {
    let (mut k, _, tid_a, pid_b, tid_b, cap_b) = two_processes();

    k.switch_to_thread_for_test(tid_b);
    let va = request(&mut k, 2).unwrap();
    assert!(release(&mut k, va, 2));

    k.switch_to_thread_for_test(tid_a);
    // Only what is still charged to B (its page tables) may come back.
    let expected = k.free_page_count() + k.page_usage(pid_b);
    k.handle_syscall(Syscall::Kill { process: cap_b });
    assert_eq!(k.free_page_count(), expected);
}

#[test]
fn killed_process_pages_return_to_pool() {
    let (mut k, _, tid_a, pid_b, tid_b, cap_b) = two_processes();
    let baseline = k.free_page_count() + k.page_usage(pid_b);

    k.switch_to_thread_for_test(tid_b);
    request(&mut k, 3).unwrap();

    k.switch_to_thread_for_test(tid_a);
    k.handle_syscall(Syscall::Kill { process: cap_b });
    assert_eq!(k.free_page_count(), baseline);
    assert_eq!(k.page_usage(pid_b), 0);
}

#[test]
fn map_page_rejects_destination_outside_user_window() {
    let (mut k, _, _, pid_b, _, cap_b) = two_processes();
    let src = request(&mut k, 1).unwrap();
    let kernel_text = 0x8020_0000;
    assert!(!map_page(&mut k, cap_b, src, kernel_text));
    assert_eq!(k.platform_mut().mapping(pid_b, kernel_text), None);
}

#[test]
fn map_page_rejects_null_page_destination() {
    let (mut k, _, _, pid_b, _, cap_b) = two_processes();
    let src = request(&mut k, 1).unwrap();
    assert!(!map_page(&mut k, cap_b, src, 0));
    assert_eq!(k.platform_mut().mapping(pid_b, 0), None);
}

#[test]
fn map_page_moves_the_page_out_of_the_caller() {
    let (mut k, pid_a, _, pid_b, _, cap_b) = two_processes();
    let src = request(&mut k, 1).unwrap();
    let pa = k.platform_mut().mapping(pid_a, src).unwrap();

    assert!(map_page(&mut k, cap_b, src, 0x10000));
    assert_eq!(k.platform_mut().mapping(pid_a, src), None);
    assert_eq!(k.platform_mut().mapping(pid_b, 0x10000), Some(pa));
}

#[test]
fn map_page_cannot_overwrite_an_existing_mapping() {
    let (mut k, _, _, pid_b, _, cap_b) = two_processes();
    let first = request(&mut k, 1).unwrap();
    let second = request(&mut k, 1).unwrap();

    assert!(map_page(&mut k, cap_b, first, 0x10000));
    let pa = k.platform_mut().mapping(pid_b, 0x10000);
    assert!(!map_page(&mut k, cap_b, second, 0x10000));
    assert_eq!(k.platform_mut().mapping(pid_b, 0x10000), pa);
}

#[test]
fn map_page_source_must_be_a_page_the_caller_owns() {
    let (mut k, _, _, _, _, cap_b) = two_processes();
    // An MMIO address the caller never received through RequestPages.
    assert!(!map_page(&mut k, cap_b, 0x1000_0000, 0x10000));
}

#[test]
fn heap_never_grows_outside_the_user_window() {
    let (mut k, _) = kernel_with_process();
    for _ in 0..10_000 {
        let Some(va) = request(&mut k, 1) else { break };
        assert!(
            va >= PAGE && va + PAGE <= USER_VA_END,
            "heap page at {va:#x}"
        );
        assert!(release(&mut k, va, 1));
    }
}

#[test]
fn requested_pages_are_zeroed_before_mapping() {
    let (mut k, pid) = kernel_with_process();
    let va = request(&mut k, 2).unwrap();
    for i in 0..2 {
        let pa = k.platform_mut().mapping(pid, va + i * PAGE).unwrap();
        assert!(
            k.platform_mut().was_zeroed(pa),
            "page {pa:#x} handed out dirty"
        );
    }
}

#[test]
fn page_tables_count_against_the_process_quota() {
    let (mut k, pid) = kernel_with_process();
    let used = k.page_usage(pid);
    // Room for exactly one data page plus the leaf table it needs.
    k.set_page_limit(pid, used + 2).unwrap();
    assert!(request(&mut k, 1).is_some());
    assert!(request(&mut k, 1).is_none(), "quota should be exhausted");
}

#[test]
fn killing_a_process_returns_its_page_tables_too() {
    let (mut k, pid_a) = kernel_with_process();
    let before = k.free_page_count();
    let cap_b = match k.handle_syscall(Syscall::CreateProcess {
        create_cap: CapSlot::new(0).unwrap(),
        image: DataRef::empty(),
    }) {
        SyscallOutcome::Complete(SyscallResult::Slot(s)) => s,
        _ => panic!("create process"),
    };
    let src = request(&mut k, 1).unwrap();
    assert!(map_page(&mut k, cap_b, src, 0x10000));
    k.handle_syscall(Syscall::Kill { process: cap_b });
    // A keeps its own heap leaf table; B's root, leaf and page are all back.
    assert_eq!(k.free_page_count(), before - (k.page_usage(pid_a) - 1));
    assert!(!k.platform_mut().has_address_space(ProcessId::from_raw(1)));
}

#[test]
fn map_page_moves_the_quota_charge_to_the_target() {
    let (mut k, pid_a, _, pid_b, _, cap_b) = two_processes();
    let src = request(&mut k, 1).unwrap();
    let a_before = k.page_usage(pid_a);
    let b_before = k.page_usage(pid_b);
    assert!(map_page(&mut k, cap_b, src, 0x10000));
    assert_eq!(k.page_usage(pid_a), a_before - 1);
    // the page plus the leaf table covering 0x10000
    assert_eq!(k.page_usage(pid_b), b_before + 2);
}

#[test]
fn map_page_fails_cleanly_when_target_quota_is_full() {
    let (mut k, pid_a, _, pid_b, _, cap_b) = two_processes();
    let src = request(&mut k, 1).unwrap();
    let used = k.page_usage(pid_b);
    k.set_page_limit(pid_b, used).unwrap();
    assert!(!map_page(&mut k, cap_b, src, 0x10000));
    // the page stays with the caller
    assert!(k.platform_mut().mapping(pid_a, src).is_some());
    assert_eq!(k.platform_mut().mapping(pid_b, 0x10000), None);
}

#[test]
fn shared_region_is_visible_in_both_address_spaces() {
    let (mut k, pid_a, _, pid_b, _, _) = two_processes();
    let id = k.create_pool_region(pid_a, pages(2)).unwrap();
    k.map_region(
        id,
        pid_a,
        VirtualAddress::from_raw(0x30000),
        PagePerms::ReadWrite,
    )
    .unwrap();
    k.map_region(
        id,
        pid_b,
        VirtualAddress::from_raw(0x50000),
        PagePerms::ReadOnly,
    )
    .unwrap();
    for i in 0..2 {
        let a = k.platform_mut().mapping(pid_a, 0x30000 + i * PAGE);
        assert!(a.is_some());
        assert_eq!(a, k.platform_mut().mapping(pid_b, 0x50000 + i * PAGE));
    }
    assert_eq!(
        k.platform_mut().perms(pid_b, 0x50000),
        Some(PagePerms::ReadOnly)
    );
}

#[test]
fn heap_requests_never_land_on_a_region_mapping() {
    let (mut k, pid) = kernel_with_process();
    let id = k.create_pool_region(pid, pages(4)).unwrap();
    // Exactly where the kernel would otherwise put the first heap page.
    let heap_start = request(&mut k, 1).unwrap();
    assert!(release(&mut k, heap_start, 1));
    k.map_region(
        id,
        pid,
        VirtualAddress::from_raw(heap_start),
        PagePerms::ReadWrite,
    )
    .unwrap();
    let va = request(&mut k, 1).unwrap();
    assert!(
        va >= heap_start + 4 * PAGE || va + PAGE <= heap_start,
        "heap page {va:#x} inside region"
    );
}

#[test]
fn map_page_cannot_target_a_region_mapping() {
    let (mut k, _, _, pid_b, _, cap_b) = two_processes();
    let id = k.create_pool_region(pid_b, pages(1)).unwrap();
    k.map_region(
        id,
        pid_b,
        VirtualAddress::from_raw(0x30000),
        PagePerms::ReadWrite,
    )
    .unwrap();
    let region_pa = k.platform_mut().mapping(pid_b, 0x30000);
    let src = request(&mut k, 1).unwrap();
    assert!(!map_page(&mut k, cap_b, src, 0x30000));
    assert_eq!(k.platform_mut().mapping(pid_b, 0x30000), region_pa);
}

#[test]
fn killing_a_region_owner_unmaps_it_everywhere() {
    let (mut k, pid_a, tid_a, pid_b, tid_b, cap_b) = two_processes();
    let id = k.create_pool_region(pid_b, pages(2)).unwrap();
    k.map_region(
        id,
        pid_b,
        VirtualAddress::from_raw(0x30000),
        PagePerms::ReadWrite,
    )
    .unwrap();
    k.map_region(
        id,
        pid_a,
        VirtualAddress::from_raw(0x30000),
        PagePerms::ReadWrite,
    )
    .unwrap();
    let _ = (tid_a, tid_b);

    let before = k.free_page_count();
    k.handle_syscall(Syscall::Kill { process: cap_b });
    assert_eq!(k.platform_mut().mapping(pid_a, 0x30000), None);
    assert_eq!(k.platform_mut().mapping(pid_a, 0x31000), None);
    // region frames, B's root and B's leaf table are all freed
    assert_eq!(k.free_page_count(), before + 4);
}

#[test]
fn region_cannot_overlap_existing_pages() {
    let (mut k, pid) = kernel_with_process();
    let va = request(&mut k, 1).unwrap();
    let id = k.create_pool_region(pid, pages(1)).unwrap();
    assert!(
        k.map_region(id, pid, VirtualAddress::from_raw(va), PagePerms::ReadWrite)
            .is_err()
    );
}

#[test]
fn region_mapping_must_fit_the_user_window() {
    let (mut k, pid) = kernel_with_process();
    let id = k.create_pool_region(pid, pages(2)).unwrap();
    let last = USER_VA_END - PAGE;
    assert!(
        k.map_region(
            id,
            pid,
            VirtualAddress::from_raw(last),
            PagePerms::ReadWrite
        )
        .is_err()
    );
    assert_eq!(k.platform_mut().mapping(pid, last), None);
}