os/test/
guard.rs

1use core::sync::atomic::{AtomicPtr, Ordering};
2
3use crate::println;
4static MOCK_HANDLER: AtomicPtr<()> = AtomicPtr::new(core::ptr::null_mut());
5
6pub enum TestEnvironment {
7    None,
8    Interrupt(fn()),
9}
10pub struct TestEnvGuard {
11    prev_flags: usize,
12    prev_handler: Option<fn()>,
13}
14
15impl TestEnvGuard {
16    pub fn enter(env: TestEnvironment) -> Self {
17        unsafe fn read_flags() -> usize {
18            crate::arch::intr::read_and_disable_interrupts()
19        }
20
21        // 在设置新环境前,先保存当前的状态
22        let current_handler_ptr = MOCK_HANDLER.load(Ordering::Relaxed);
23        let prev_handler = if current_handler_ptr.is_null() {
24            None
25        } else {
26            // 将裸指针转换回函数指针 fn() 以便保存
27            Some(unsafe { core::mem::transmute(current_handler_ptr) })
28        };
29
30        let mut guard = TestEnvGuard {
31            prev_flags: unsafe { read_flags() },
32            prev_handler, // 保存旧的 handler
33        };
34
35        match env {
36            TestEnvironment::None => {}
37            TestEnvironment::Interrupt(handler) => unsafe {
38                crate::arch::intr::enable_interrupts();
39            },
40        }
41
42        guard
43    }
44}
45
46impl Drop for TestEnvGuard {
47    fn drop(&mut self) {
48        if let Some(old_handler) = self.prev_handler {
49            let old_ptr = old_handler as *mut ();
50            println!("[mock] restoring previous interrupt handler {:p}", old_ptr);
51            MOCK_HANDLER.store(old_ptr, core::sync::atomic::Ordering::Relaxed);
52        } else {
53            MOCK_HANDLER.store(core::ptr::null_mut(), core::sync::atomic::Ordering::Relaxed);
54        }
55    }
56}
57
58pub fn trigger_mock_interrupt() {
59    let ptr = MOCK_HANDLER.load(Ordering::Relaxed);
60    if !ptr.is_null() {
61        let f: fn() = unsafe { core::mem::transmute(ptr) };
62        println!("[mock] triggering fake interrupt...");
63        f();
64    } else {
65        println!("[mock] no handler registered");
66    }
67}