ax-cpu 0.9.0

//! Structures and functions for user space.

use core::{
    mem::{align_of, offset_of, size_of},
    ops::{Deref, DerefMut},
};

use ax_memory_addr::VirtAddr;
use x86_64::{
    registers::{
        control::{Cr2, Cr4, Cr4Flags},
        model_specific::{Efer, EferFlags, LStar, SFMask, Star},
        rflags::RFlags,
    },
    structures::idt::ExceptionVector,
};

use super::{
    TrapFrame, gdt,
    trap::{IRQ_VECTOR_END, IRQ_VECTOR_START, LEGACY_SYSCALL_VECTOR, err_code_to_flags},
};
pub use crate::uspace_common::{ExceptionKind, ExceptionSyndrome, ReturnReason};

/// Context to enter user space.
#[derive(Debug, Clone, Copy)]
#[repr(C, align(16))]
pub struct UserContext {
    tf: TrapFrame,
    /// User-owned FS segment base.
    ///
    /// `CR4.FSGSBASE` stays disabled, so userspace cannot modify this value
    /// without an `arch_prctl`-style kernel operation updating this image.
    pub fs_base: u64,
    /// User-owned GS segment base.
    pub gs_base: u64,
    /// Kernel continuation stack saved while this context executes in ring 3.
    kernel_stack_pointer: u64,
    /// Explicitly initializes the tail bytes required by the 16-byte ABI alignment.
    _reserved: u64,
}

// SAFETY: `TrapFrame` and every following field are integer-only, the explicit
// tail word consumes the alignment padding, and the offset assertions below
// pin that layout.
unsafe impl bytemuck::NoUninit for UserContext {}

const _: () = {
    // A privilege transition may align TSS.RSP0 down to 16 bytes before
    // constructing the hardware frame. `enter_user` uses the end of `tf` as
    // both RSP0 and the boundary above which it saves the kernel continuation,
    // so both the object and that boundary must already be aligned.
    assert!(align_of::<UserContext>() >= 16);
    assert!(size_of::<TrapFrame>().is_multiple_of(16));
    assert!(offset_of!(UserContext, tf) == 0);
    assert!(offset_of!(UserContext, fs_base) == size_of::<TrapFrame>());
    assert!(offset_of!(UserContext, gs_base) == size_of::<TrapFrame>() + size_of::<u64>());
    assert!(
        offset_of!(UserContext, kernel_stack_pointer)
            == size_of::<TrapFrame>() + 2 * size_of::<u64>()
    );
    assert!(offset_of!(UserContext, _reserved) == size_of::<TrapFrame>() + 3 * size_of::<u64>());
    assert!(size_of::<UserContext>() == size_of::<TrapFrame>() + 4 * size_of::<u64>());
};

impl UserContext {
    /// Creates a new context with the given entry point, user stack pointer,
    /// and the argument.
    pub fn new(entry: usize, ustack_top: VirtAddr, arg0: usize) -> Self {
        use x86_64::registers::rflags::RFlags;
        Self {
            tf: TrapFrame {
                rdi: arg0 as _,
                rip: entry as _,
                cs: gdt::UCODE64.0 as _,
                rflags: RFlags::INTERRUPT_FLAG.bits(), // IOPL = 0, IF = 1
                rsp: ustack_top.as_usize() as _,
                ss: gdt::UDATA.0 as _,
                ..Default::default()
            },
            fs_base: 0,
            gs_base: 0,
            kernel_stack_pointer: 0,
            _reserved: 0,
        }
    }

    /// Normalizes a cloned user context so it can safely return to ring 3.
    pub fn prepare_clone_child_return_state(&mut self) {
        let mut flags = RFlags::from_bits_truncate(self.tf.rflags);
        flags.insert(RFlags::INTERRUPT_FLAG);
        flags.remove(RFlags::TRAP_FLAG | RFlags::NESTED_TASK | RFlags::RESUME_FLAG);
        self.tf.rflags = flags.bits();
    }

    /// Clears the single-step trap flag after a debug exception.
    ///
    /// Returns whether the flag had been set in the saved user context.
    pub fn clear_single_step_after_debug(&mut self) -> bool {
        let mut flags = RFlags::from_bits_truncate(self.tf.rflags);
        let was_set = flags.contains(RFlags::TRAP_FLAG);
        flags.remove(RFlags::TRAP_FLAG);
        self.tf.rflags = flags.bits();
        was_set
    }

    /// Returns the syscall instruction length in bytes.
    pub const fn syscall_insn_len(&self) -> usize {
        2
    }

    /// Gets the TLS area.
    pub const fn tls(&self) -> usize {
        self.fs_base as _
    }

    /// Sets the TLS area.
    pub const fn set_tls(&mut self, tls_area: usize) {
        self.fs_base = tls_area as _;
    }

    /// Returns whether this register image can be restored as an interruptible
    /// ring-3 context.
    pub fn has_interruptible_user_return_mode(&self) -> bool {
        let forbidden =
            RFlags::IOPL_LOW | RFlags::IOPL_HIGH | RFlags::NESTED_TASK | RFlags::VIRTUAL_8086_MODE;
        let flags = RFlags::from_bits_retain(self.tf.rflags);
        self.tf.cs == gdt::UCODE64.0 as u64
            && self.tf.ss == gdt::UDATA.0 as u64
            && flags.contains(RFlags::INTERRUPT_FLAG)
            && !flags.intersects(forbidden)
    }

    /// Enters user space without validating the runtime transition.
    ///
    /// It restores the user registers and jumps to the user entry point
    /// (saved in `rip`).
    ///
    /// This function returns when an exception or syscall occurs.
    ///
    /// # Safety
    ///
    /// The caller must be the runtime's prepared user-entry boundary for the
    /// current scheduler task. Its context-switch tail must be complete, no
    /// IRQ/preemption guard or hard interrupt may be active, and local IRQs
    /// must remain disabled after the final scheduler-work check. The active
    /// logical address space, hardware root and CPU footprint must match this
    /// task and keep every user address referenced by `self` valid. The saved
    /// selectors and RFLAGS must describe an interruptible ring-3 return. No
    /// code may run between those validations and this call.
    ///
    /// Safe code cannot invoke this raw boundary:
    ///
    /// ```compile_fail
    /// fn bypass_runtime(context: &mut ax_cpu::uspace::UserContext) {
    ///     context.run_unchecked();
    /// }
    /// ```
    pub unsafe fn run_unchecked(&mut self) -> ReturnReason {
        unsafe extern "C" {
            fn enter_user(uctx: &mut UserContext);
        }

        assert!(
            self.has_interruptible_user_return_mode(),
            "raw user entry requires an interruptible ring-3 register image"
        );

        assert!(
            !crate::asm::irqs_enabled(),
            "raw user entry requires the prepared IRQ-off boundary"
        );
        super::local_state::install_current_user_tls(self.fs_base as _, self.gs_base as _);

        unsafe { enter_user(self) };

        let vector = self.vector as u8;

        const PAGE_FAULT_VECTOR: u8 = ExceptionVector::Page as u8;

        let ret = match (vector, err_code_to_flags(self.error_code)) {
            (PAGE_FAULT_VECTOR, Ok(flags)) => {
                ReturnReason::PageFault(va!(Cr2::read_raw() as usize), flags)
            }
            (LEGACY_SYSCALL_VECTOR, _) => ReturnReason::Syscall,
            (IRQ_VECTOR_START..=IRQ_VECTOR_END, _) => {
                crate::trap::dispatch_irq(vector as _, crate::trap::TrapOrigin::User);
                ReturnReason::Interrupt
            }
            _ => ReturnReason::Exception(ExceptionInfo {
                vector,
                error_code: self.error_code,
                cr2: Cr2::read_raw() as usize,
            }),
        };

        crate::asm::enable_irqs();
        ret
    }
}

const _: unsafe fn(&mut UserContext) -> ReturnReason = UserContext::run_unchecked;

impl Deref for UserContext {
    type Target = TrapFrame;

    fn deref(&self) -> &Self::Target {
        &self.tf
    }
}

impl DerefMut for UserContext {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.tf
    }
}

/// Information about an exception that occurred in user space.
#[derive(Debug, Clone, Copy)]
pub struct ExceptionInfo {
    /// The exception vector.
    pub vector: u8,
    /// The error code.
    pub error_code: u64,
    /// The faulting virtual address (if applicable).
    pub cr2: usize,
}

impl ExceptionInfo {
    /// Returns the faulting virtual address when the CPU records one.
    pub const fn fault_addr(&self) -> Option<usize> {
        Some(self.cr2)
    }

    /// Returns architecture-neutral syndrome information for this exception.
    pub const fn syndrome(&self) -> ExceptionSyndrome {
        ExceptionSyndrome {
            raw: self.error_code,
            class: self.vector as u64,
            iss: 0,
        }
    }

    /// Returns a generalized kind of this exception.
    pub fn kind(&self) -> ExceptionKind {
        match ExceptionVector::try_from(self.vector) {
            Ok(ExceptionVector::Debug) => ExceptionKind::Debug,
            Ok(ExceptionVector::Breakpoint) => ExceptionKind::Breakpoint,
            Ok(ExceptionVector::InvalidOpcode) => ExceptionKind::IllegalInstruction,
            // `#DE`: integer divide-by-zero / `INT_MIN / -1`. Linux delivers this
            // as SIGFPE/FPE_INTDIV; the HotSpot JVM's x86 interpreter and JIT
            // rely on the trap to raise Java `ArithmeticException`.
            Ok(ExceptionVector::Division) => ExceptionKind::ArithmeticError,
            _ => ExceptionKind::Other,
        }
    }
}

/// Initializes syscall support and setups the syscall handler.
pub(super) fn init_syscall() {
    unsafe extern "C" {
        fn syscall_entry();
    }

    assert!(
        !Cr4::read().contains(Cr4Flags::FSGSBASE),
        "LinuxCurrent user TLS requires trapping all FS/GS base changes"
    );
    super::local_state::initialize_cpu_user_tls();
    LStar::write(x86_64::VirtAddr::new_truncate(
        syscall_entry as *const () as usize as _,
    ));
    Star::write(gdt::UCODE64, gdt::UDATA, gdt::KCODE64, gdt::KDATA).unwrap();
    SFMask::write(
        RFlags::TRAP_FLAG
            | RFlags::INTERRUPT_FLAG
            | RFlags::DIRECTION_FLAG
            | RFlags::IOPL_LOW
            | RFlags::IOPL_HIGH
            | RFlags::NESTED_TASK
            | RFlags::ALIGNMENT_CHECK,
    ); // TF | IF | DF | IOPL | AC | NT (0x47700)
    unsafe {
        Efer::update(|efer| *efer |= EferFlags::SYSTEM_CALL_EXTENSIONS);
    }
}