ax-cpu 0.10.0

use x86::{controlregs::cr2, irq::*};
use x86_64::{registers::rflags::RFlags, structures::idt::PageFaultErrorCode};

use super::context::TrapFrame;
use crate::{TrapOrigin, trap::PageFaultFlags};

// In 64-bit mode the CPU saves SS:RSP even for same-privilege traps. Use
// the complete hardware return frame, including when IST changes the stack.
// Linux arch/x86/entry/entry_64.S, restore_regs_and_return_to_kernel, relies
// on the same SS:RSP contract. A frame-address calculation loses alignment
// and IST information and is not the interrupted stack pointer.
type RawTrapFrame = TrapFrame;

/// Lifetime-bound view of a kernel-origin x86 trap frame.
///
/// The view deliberately provides no mutable dereference to the assembly
/// image. Probe/debug integrations can edit a copy and apply it through
/// [`Self::apply_registers`], which preserves the trap origin and vector
/// metadata.
pub struct KernelTrapFrame<'a> {
    raw: &'a mut RawTrapFrame,
    _not_send: core::marker::PhantomData<*mut ()>,
}

impl<'a> KernelTrapFrame<'a> {
    /// Returns the privilege domain represented by this view.
    pub const fn origin(&self) -> TrapOrigin {
        TrapOrigin::Kernel
    }

    /// Copies the saved register image for inspection or probe emulation.
    pub fn snapshot(&self) -> TrapFrame {
        *self.raw
    }

    /// Applies task-register changes while preserving trap-origin metadata.
    pub fn apply_registers(&mut self, updated: &TrapFrame) {
        self.raw.regs = updated.regs;
        self.raw.rip = updated.rip;
        self.raw.rflags = updated.rflags;
    }

    /// Returns the saved instruction pointer.
    pub const fn ip(&self) -> usize {
        self.raw.rip as usize
    }

    /// Sets the saved instruction pointer.
    pub const fn set_ip(&mut self, ip: usize) {
        self.raw.rip = ip as u64;
    }

    /// Creates the typed view at the assembly boundary.
    ///
    /// # Safety
    ///
    /// `raw` must be the uniquely borrowed, live kernel-origin frame built by
    /// the x86 trap entry and must remain valid for `'a`.
    unsafe fn from_raw(raw: &'a mut RawTrapFrame) -> Self {
        debug_assert_eq!(raw.cs & 0b11, 0);
        Self {
            raw,
            _not_send: core::marker::PhantomData,
        }
    }
}

impl core::fmt::Debug for KernelTrapFrame<'_> {
    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        self.snapshot().fmt(formatter)
    }
}

pub(super) const LEGACY_SYSCALL_VECTOR: u8 = 0x80;
pub(super) const IRQ_VECTOR_START: u8 = 0x20;
pub(super) const IRQ_VECTOR_END: u8 = 0xff;

fn handle_page_fault(tf: &mut KernelTrapFrame<'_>) {
    let access_flags = err_code_to_flags(tf.raw.error_code)
        .unwrap_or_else(|e| panic!("Invalid #PF error code: {:#x}", e));
    let vaddr = va!(unsafe { cr2() });
    #[cfg(feature = "exception-table")]
    {
        let mut updated = tf.snapshot();
        if updated.fixup_nofault_exception() {
            tf.apply_registers(&updated);
            return;
        }
    }
    if crate::trap::call_page_fault_handler_with_parent_irqs(
        vaddr,
        access_flags,
        RFlags::from_bits_truncate(tf.raw.rflags).contains(RFlags::INTERRUPT_FLAG),
    ) {
        return;
    }
    #[cfg(feature = "exception-table")]
    {
        let mut updated = tf.snapshot();
        if updated.fixup_exception() {
            tf.apply_registers(&updated);
            return;
        }
    }
    let snapshot = tf.snapshot();
    let bt = crate::trap::diagnostics::BacktraceDisplay(snapshot.backtrace_registers());
    panic!(
        "Unhandled #PF @ {:#x}, fault_vaddr={:#x}, error_code={:#x} ({:?}):\n{:#x?}\n{}",
        tf.raw.rip, vaddr, tf.raw.error_code, access_flags, snapshot, bt
    );
}

fn handle_breakpoint(tf: &mut KernelTrapFrame<'_>) {
    debug!("#BP @ {:#x} ", tf.raw.rip);
    let _ = crate::trap::breakpoint_handler(tf);
}

fn handle_debug(tf: &mut KernelTrapFrame<'_>) {
    debug!("#DB @ {:#x} ", tf.raw.rip);
    if crate::trap::debug_handler(tf) {
        return;
    }
    // Kernel-mode #DB was not claimed by any handler.
    // Unclaimed user-mode #DB is routed through the user-space exception loop
    // (.Ltrap_user → .Lexit_user in trap.S), so `x86_trap_handler` is only
    // reached for kernel-mode traps. An unhandled kernel #DB is a fatal
    // condition: if resumed the CPU re-executes the faulting instruction,
    // likely looping into a triple fault.
    warn!("Unhandled kernel #DB @ {:#x}", tf.raw.rip);
    let snapshot = tf.snapshot();
    let bt = crate::trap::diagnostics::BacktraceDisplay(snapshot.backtrace_registers());
    panic!(
        "Unhandled #DB @ {:#x}, error_code={:#x}:\n{:#x?}\n{}",
        tf.raw.rip, tf.raw.error_code, snapshot, bt
    );
}

#[unsafe(no_mangle)]
unsafe extern "C" fn x86_trap_handler(raw: *mut RawTrapFrame) {
    // SAFETY: every x86 trap vector allocates one complete, exclusively owned
    // frame and passes its aligned stack address in the C argument register.
    let raw = unsafe { &mut *raw };
    let mut tf = unsafe { KernelTrapFrame::from_raw(raw) };
    match tf.raw.vector as u8 {
        PAGE_FAULT_VECTOR => handle_page_fault(&mut tf),
        BREAKPOINT_VECTOR => handle_breakpoint(&mut tf),
        DEBUG_VECTOR => handle_debug(&mut tf),
        GENERAL_PROTECTION_FAULT_VECTOR => {
            let snapshot = tf.snapshot();
            let bt = crate::trap::diagnostics::BacktraceDisplay(snapshot.backtrace_registers());
            panic!(
                "#GP @ {:#x}, error_code={:#x}:\n{:#x?}\n{}",
                tf.raw.rip, tf.raw.error_code, snapshot, bt
            );
        }
        IRQ_VECTOR_START..=IRQ_VECTOR_END => {
            crate::trap::dispatch_irq(
                tf.raw.vector as _,
                crate::trap::TrapOrigin::Kernel,
                Some(tf.snapshot().interrupted_context()),
            );
        }
        _ => {
            let snapshot = tf.snapshot();
            let bt = crate::trap::diagnostics::BacktraceDisplay(snapshot.backtrace_registers());
            panic!(
                "Unhandled exception {} ({}, error_code={:#x}) @ {:#x}:\n{:#x?}\n{}",
                tf.raw.vector,
                vec_to_str(tf.raw.vector),
                tf.raw.error_code,
                tf.raw.rip,
                snapshot,
                bt
            );
        }
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn x86_double_fault_handler(raw: *mut RawTrapFrame) -> ! {
    let fault_vaddr = unsafe { cr2() };
    // SAFETY: vector 8 enters through its dedicated IST stack, builds the same
    // register prefix as every other x86 trap, and never returns or aliases it.
    let raw = unsafe { &*raw };
    panic!(
        "Fatal #DF @ {:#x}, fault_vaddr={:#x}, error_code={:#x}, cs={:#x}, rflags={:#x}",
        raw.rip, fault_vaddr, raw.error_code, raw.cs, raw.rflags,
    );
}

fn vec_to_str(vec: u64) -> &'static str {
    if vec < 32 {
        EXCEPTIONS[vec as usize].mnemonic
    } else {
        "Unknown"
    }
}

pub(super) fn err_code_to_flags(err_code: u64) -> Result<PageFaultFlags, u64> {
    let code = PageFaultErrorCode::from_bits_truncate(err_code);
    let reserved_bits = (PageFaultErrorCode::CAUSED_BY_WRITE
        | PageFaultErrorCode::USER_MODE
        | PageFaultErrorCode::INSTRUCTION_FETCH
        | PageFaultErrorCode::PROTECTION_VIOLATION)
        .complement();
    if code.intersects(reserved_bits) {
        Err(err_code)
    } else {
        let mut flags = PageFaultFlags::empty();
        if code.contains(PageFaultErrorCode::CAUSED_BY_WRITE) {
            flags |= PageFaultFlags::WRITE;
        } else {
            flags |= PageFaultFlags::READ;
        }
        if code.contains(PageFaultErrorCode::USER_MODE) {
            flags |= PageFaultFlags::USER;
        }
        if code.contains(PageFaultErrorCode::INSTRUCTION_FETCH) {
            flags |= PageFaultFlags::EXECUTE;
        }
        Ok(flags)
    }
}