ax-cpu 0.10.0

//! Lazy x86 LinuxCurrent userspace register ownership.

use core::mem::offset_of;

use super::entry_state::CpuEntryState;

const IA32_FS_BASE: u32 = 0xc000_0100;
const IA32_KERNEL_GS_BASE: u32 = 0xc000_0102;

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
struct UserTlsValues {
    fs_base: usize,
    gs_base: usize,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct UserTlsWrites {
    fs_base: bool,
    gs_base: bool,
}

const CPU_USER_FS_BASE_OFFSET: usize = offset_of!(CpuEntryState, fs_base);

const CPU_USER_GS_BASE_OFFSET: usize = offset_of!(CpuEntryState, gs_base);

const CPU_USER_TLS_GENERATION_OFFSET: usize = offset_of!(CpuEntryState, tls_generation);

const CPU_USER_FP_OWNER_OFFSET: usize = offset_of!(CpuEntryState, user_fp_owner);
#[cfg(feature = "fp-simd")]
const CPU_USER_XSAVE_CONFIG_OFFSET: usize = offset_of!(CpuEntryState, xsave_config);

#[cfg(feature = "fp-simd")]
const USER_XSAVEOPT_ENABLED: usize = 1 << (usize::BITS - 1);

#[cfg(feature = "fp-simd")]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum UserFpOwnerMatch {
    Unowned,
    Current,
    Foreign,
}

#[cfg(feature = "fp-simd")]
fn classify_user_fp_owner(owner: usize, current: usize) -> UserFpOwnerMatch {
    if owner == 0 {
        UserFpOwnerMatch::Unowned
    } else if owner == current {
        UserFpOwnerMatch::Current
    } else {
        UserFpOwnerMatch::Foreign
    }
}

fn changed_user_tls(
    previous: UserTlsValues,
    next: UserTlsValues,
    initialized: bool,
) -> UserTlsWrites {
    UserTlsWrites {
        fs_base: !initialized || previous.fs_base != next.fs_base,
        gs_base: !initialized || previous.gs_base != next.gs_base,
    }
}

fn next_generation(previous: usize) -> usize {
    match previous.wrapping_add(1) {
        0 => 1,
        generation => generation,
    }
}

fn current_cpu_user_tls() -> (UserTlsValues, usize) {
    let fs_base: usize;
    let gs_base: usize;
    let generation: usize;
    // SAFETY: this runs with local IRQs disabled after the CPU area has been
    // installed. The three fields are owned by this CPU and no remote path
    // reads or writes the architecture reserve.
    unsafe {
        core::arch::asm!(
            "mov {fs_base}, gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {fs_offset}]",
            "mov {gs_base}, gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {gs_offset}]",
            "mov {generation}, gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {generation_offset}]",
            fs_base = out(reg) fs_base,
            gs_base = out(reg) gs_base,
            generation = out(reg) generation,
            fs_offset = const CPU_USER_FS_BASE_OFFSET,
            gs_offset = const CPU_USER_GS_BASE_OFFSET,
            generation_offset = const CPU_USER_TLS_GENERATION_OFFSET,
            options(nostack, preserves_flags, readonly),
        );
    }
    (UserTlsValues { fs_base, gs_base }, generation)
}

fn publish_current_cpu_user_tls(values: UserTlsValues, generation: usize) {
    // SAFETY: the caller retains the same IRQ-disabled CPU ownership used by
    // current_cpu_user_tls. Publishing the generation last makes a future
    // diagnostic reader reject a partially updated cache image.
    unsafe {
        core::arch::asm!(
            "mov gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {fs_offset}], {fs_base}",
            "mov gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {gs_offset}], {gs_base}",
            "mov gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {generation_offset}], {generation}",
            fs_offset = const CPU_USER_FS_BASE_OFFSET,
            gs_offset = const CPU_USER_GS_BASE_OFFSET,
            generation_offset = const CPU_USER_TLS_GENERATION_OFFSET,
            fs_base = in(reg) values.fs_base,
            gs_base = in(reg) values.gs_base,
            generation = in(reg) generation,
            options(nostack, preserves_flags),
        );
    }
}

#[cfg(feature = "fp-simd")]
fn current_cpu_user_fp_owner() -> usize {
    let owner: usize;
    // SAFETY: local IRQs are disabled after CPU-area installation. The owner
    // word is private to this physical CPU and is never accessed remotely.
    unsafe {
        core::arch::asm!(
            "mov {owner}, gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {owner_offset}]",
            owner = out(reg) owner,
            owner_offset = const CPU_USER_FP_OWNER_OFFSET,
            options(nostack, preserves_flags, readonly),
        );
    }
    owner
}

fn publish_current_cpu_user_fp_owner(owner: usize) {
    // SAFETY: the caller retains the same IRQ-disabled CPU ownership used by
    // `current_cpu_user_fp_owner` for the complete hardware-state transition.
    unsafe {
        core::arch::asm!(
            "mov gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {owner_offset}], {owner}",
            owner_offset = const CPU_USER_FP_OWNER_OFFSET,
            owner = in(reg) owner,
            options(nostack, preserves_flags),
        );
    }
}

#[cfg(feature = "fp-simd")]
pub(super) fn current_cpu_user_xsave_config() -> Option<(u64, bool)> {
    let config: usize;
    // SAFETY: userspace CPU initialization publishes this immutable word
    // before this CPU can schedule a userspace context. It is thereafter read
    // only by the same IRQ-disabled CPU during FP save and restore.
    unsafe {
        core::arch::asm!(
            "mov {config}, gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {config_offset}]",
            config = out(reg) config,
            config_offset = const CPU_USER_XSAVE_CONFIG_OFFSET,
            options(nostack, preserves_flags, readonly),
        );
    }
    let mask = config & !USER_XSAVEOPT_ENABLED;
    (mask != 0).then_some((mask as u64, config & USER_XSAVEOPT_ENABLED != 0))
}

#[cfg(feature = "fp-simd")]
fn publish_current_cpu_user_xsave_config(mask: u64, xsaveopt_enabled: bool) {
    let config = mask as usize | usize::from(xsaveopt_enabled) * USER_XSAVEOPT_ENABLED;
    // SAFETY: this runs once for the current CPU after its GS CPU-area base and
    // XCR0 policy are installed, before that CPU can enter userspace.
    unsafe {
        core::arch::asm!(
            "mov gs:[offset __AX_CPU_AREA_ARCH_STATE_OFFSET + {config_offset}], {config}",
            config_offset = const CPU_USER_XSAVE_CONFIG_OFFSET,
            config = in(reg) config,
            options(nostack, preserves_flags),
        );
    }
}

fn write_changed_user_tls(previous: UserTlsValues, next: UserTlsValues, initialized: bool) {
    let writes = changed_user_tls(previous, next, initialized);
    if writes.fs_base {
        write_user_tls_msr(IA32_FS_BASE, next.fs_base);
    }
    if writes.gs_base {
        write_user_tls_msr(IA32_KERNEL_GS_BASE, next.gs_base);
    }
}

fn write_user_tls_msr(msr: u32, value: usize) {
    let value = value as u64;
    // SAFETY: LinuxCurrent does not use FS for kernel TLS, and SWAPGS keeps
    // IA32_KERNEL_GS_BASE inactive while Rust executes. The caller holds this
    // CPU with IRQs disabled, so both writes update only its user return image.
    unsafe {
        core::arch::asm!(
            "wrmsr",
            in("ecx") msr,
            in("eax") value as u32,
            in("edx") (value >> 32) as u32,
            options(nostack, preserves_flags),
        );
    }
}

/// Establishes the physical reset image before this CPU can enter userspace.
pub(super) fn initialize_cpu_user_tls() {
    debug_assert!(!super::asm::irqs_enabled());
    let values = UserTlsValues::default();
    write_changed_user_tls(UserTlsValues::default(), values, false);
    publish_current_cpu_user_tls(values, 1);
    publish_current_cpu_user_fp_owner(0);
    #[cfg(feature = "fp-simd")]
    {
        let config = {
            let cr4 = unsafe { x86::controlregs::cr4() };
            if cr4.contains(x86::controlregs::Cr4::CR4_ENABLE_OS_XSAVE) {
                let mask = unsafe { x86::controlregs::xcr0().bits() };
                let xsaveopt_enabled = core::arch::x86_64::__cpuid_count(0x0d, 1).eax & 1 != 0;
                (mask, xsaveopt_enabled)
            } else {
                (0, false)
            }
        };

        publish_current_cpu_user_xsave_config(config.0, config.1);
    }
}

/// Lazily installs a user context without disturbing it in kernel-only tasks.
pub(super) fn install_current_user_tls(fs_base: usize, gs_base: usize) {
    debug_assert!(!super::asm::irqs_enabled());
    let (previous, generation) = current_cpu_user_tls();
    let next = UserTlsValues { fs_base, gs_base };
    let initialized = generation != 0;
    let writes = changed_user_tls(previous, next, initialized);
    if !writes.fs_base && !writes.gs_base {
        return;
    }
    write_changed_user_tls(previous, next, initialized);
    publish_current_cpu_user_tls(next, next_generation(generation));
}

/// Reports whether `current` owns the physical FPU image that must be saved.
#[cfg(feature = "fp-simd")]
pub(super) fn current_user_fp_is_owner(current: usize) -> bool {
    debug_assert!(!super::asm::irqs_enabled());
    match classify_user_fp_owner(current_cpu_user_fp_owner(), current) {
        UserFpOwnerMatch::Unowned => false,
        UserFpOwnerMatch::Current => true,
        UserFpOwnerMatch::Foreign => {
            panic!("x86 user FPU owner does not match the outgoing current context")
        }
    }
}

/// Clears `current` only after its physical FPU image has reached task memory.
#[cfg(feature = "fp-simd")]
pub(super) fn clear_current_user_fp_owner_after_save(_current: usize) {
    {
        debug_assert!(!super::asm::irqs_enabled());
        debug_assert_eq!(current_cpu_user_fp_owner(), _current);
    }
    publish_current_cpu_user_fp_owner(0);
}

/// Verifies that a context without a scheduler identity owns no user FPU image.
#[cfg(feature = "fp-simd")]
pub(super) fn assert_current_user_fp_unowned() {
    debug_assert!(!super::asm::irqs_enabled());
    assert_eq!(
        current_cpu_user_fp_owner(),
        0,
        "an unbound context cannot own the physical user FPU image",
    );
}

/// Reports whether `current` must restore its user FPU image before user mode.
#[cfg(feature = "fp-simd")]
pub(super) fn current_user_fp_needs_restore(current: usize) -> bool {
    debug_assert!(!super::asm::irqs_enabled());
    match classify_user_fp_owner(current_cpu_user_fp_owner(), current) {
        UserFpOwnerMatch::Unowned => true,
        UserFpOwnerMatch::Current => false,
        UserFpOwnerMatch::Foreign => {
            panic!("x86 user FPU owner does not match the return-to-user context")
        }
    }
}

/// Validates that `current` may replace the physical user FPU image.
#[cfg(feature = "fp-simd")]
pub(super) fn assert_current_user_fp_resettable(current: usize) {
    debug_assert!(!super::asm::irqs_enabled());
    match classify_user_fp_owner(current_cpu_user_fp_owner(), current) {
        UserFpOwnerMatch::Unowned | UserFpOwnerMatch::Current => {}
        UserFpOwnerMatch::Foreign => {
            panic!("x86 user FPU owner does not match the resetting current context")
        }
    }
}

/// Publishes `current` after its user FPU image has reached hardware.
#[cfg(feature = "fp-simd")]
pub(super) fn publish_current_user_fp_owner(current: usize) {
    debug_assert!(!super::asm::irqs_enabled());
    assert_ne!(current, 0, "a user FPU owner requires a context identity");
    publish_current_cpu_user_fp_owner(current);
}