proka_kernel/
test.rs

1//! # Proka Kernel - A kernel for ProkaOS
2//! Copyright (C) RainSTR Studio 2025, All rights reserved.
3//!
4//! This provides the test trait and runner.
5
6use crate::{serial_print, serial_println};
7use core::arch::asm;
8use spin::Mutex;
9
10/// The Jump Buffer to save the CPU state.
11#[derive(Debug, Clone, Copy, Default)]
12#[repr(C)]
13pub struct JumpBuffer {
14    rbx: u64,
15    rsp: u64,
16    rbp: u64,
17    r12: u64,
18    r13: u64,
19    r14: u64,
20    r15: u64,
21    rip: u64,
22}
23
24static TEST_JMP_BUF: Mutex<Option<JumpBuffer>> = Mutex::new(None);
25static FAIL_COUNT: Mutex<usize> = Mutex::new(0);
26
27/// Save the current context into the jump buffer.
28/// Returns 0 when saving, and 1 when returning from long_jmp.
29///
30/// # Safety
31/// This function is unsafe because it manipulates CPU registers directly
32/// and relies on the caller to manage the jump buffer correctly.
33#[unsafe(no_mangle)]
34pub unsafe extern "C" fn set_jmp() -> u64 {
35    let mut jmp_buf = JumpBuffer::default();
36    let res: u64;
37
38    asm!(
39        // 保存寄存器状态
40        "mov [rcx + 0], rbx",
41        "mov [rcx + 8], rsp",
42        "mov [rcx + 16], rbp",
43        "mov [rcx + 24], r12",
44        "mov [rcx + 32], r13",
45        "mov [rcx + 40], r14",
46        "mov [rcx + 48], r15",
47        // 保存返回地址
48        "lea rdx, [rip + 2f]",
49        "mov [rcx + 56], rdx",
50        // 首次调用返回 0
51        "mov rax, 0",
52        "jmp 3f",
53        // longjmp 返回点
54        "2:",
55        // longjmp 后返回 1
56        "mov rax, 1",
57        "3:",
58        in("rcx") &mut jmp_buf,
59        out("rax") res,
60        out("rdx") _,
61        clobber_abi("sysv64")
62    );
63
64    if res == 0 {
65        *TEST_JMP_BUF.lock() = Some(jmp_buf);
66    }
67    res
68}
69
70/// Restore the context and jump back to the set_jmp location.
71pub fn long_jmp() -> ! {
72    let jmp_buf = TEST_JMP_BUF.lock().expect("No jump buffer set!");
73
74    // 在恢复寄存器前增加错误计数
75    // 这样错误计数不会因为寄存器恢复而被覆盖
76    let mut fail_count = FAIL_COUNT.lock();
77    *fail_count += 1;
78    drop(fail_count); // 释放锁，防止死锁
79
80    unsafe {
81        asm!(
82            // 恢复寄存器状态
83            "mov rbx, [rcx + 0]",
84            "mov rsp, [rcx + 8]",
85            "mov rbp, [rcx + 16]",
86            "mov r12, [rcx + 24]",
87            "mov r13, [rcx + 32]",
88            "mov r14, [rcx + 40]",
89            "mov r15, [rcx + 48]",
90            // 跳转回保存的地址
91            "jmp [rcx + 56]",
92            in("rcx") &jmp_buf,
93            options(noreturn)
94        );
95    }
96}
97
98/// The trait that assign the function is testable.
99pub trait Testable {
100    /// The things will run
101    fn run(&self) -> ();
102}
103
104// This is the default implementation of this trait
105impl<T> Testable for T
106where
107    T: Fn(),
108{
109    fn run(&self) {
110        serial_print!("Testing {}... ", core::any::type_name::<T>());
111
112        // 使用 setjmp/longjmp 进行测试
113        unsafe {
114            if set_jmp() == 0 {
115                // 第一次执行测试
116                self();
117                serial_println!("[OK]");
118            } else {
119                // longjmp 返回，测试失败
120                // [FAILED] 已在 panic 处理程序中打印
121            }
122        }
123    }
124}
125
126/// This is the test runner, which will run if the test begins.
127pub fn test_runner(tests: &[&dyn Testable]) {
128    serial_println!("Running {} tests", tests.len());
129
130    for test in tests {
131        test.run();
132    }
133
134    let final_fail_count = *FAIL_COUNT.lock();
135    serial_println!("Total failures: {}", final_fail_count);
136
137    if final_fail_count > 0 {
138        serial_println!("\n[DONE] FAILED: {} tests failed", final_fail_count);
139        exit_qemu(QemuExitCode::Failed);
140    } else {
141        serial_println!("\n[DONE] SUCCESS: All tests passed!");
142        exit_qemu(QemuExitCode::Success);
143    }
144}
145
146/// This is the QEMU exit code
147pub enum QemuExitCode {
148    Success = 0x10,
149    Failed = 0x11,
150}
151
152/// The function to quit the QEMU
153pub fn exit_qemu(exit_code: QemuExitCode) {
154    use x86_64::instructions::port::Port;
155
156    unsafe {
157        let mut port = Port::new(0xf4);
158        port.write(exit_code as u32);
159    }
160}
161
162// The kernel entry, which will start up the test
163#[cfg(test)]
164#[unsafe(no_mangle)]
165pub extern "C" fn kernel_main() -> ! {
166    crate::memory::init(); // Initialize memory management
167    crate::drivers::init_devices(); // Initialize devices
168    crate::libs::time::init(); // Init time system
169    crate::libs::logger::init_logger(); // Init log system
170    crate::libs::initrd::load_initrd(); // Load initrd
171    crate::interrupts::gdt::init(); // Initialize GDT
172    crate::interrupts::idt::init_idt(); // Initialize IDT
173    crate::interrupts::pic::init(); // Initialize PI
174    x86_64::instructions::interrupts::enable(); // Enable interrupts
175    crate::test_main();
176    loop {
177        x86_64::instructions::hlt();
178    }
179}
proka_kernel/test.rs

proka_kernel/
test.rs
