//
// Syd: rock-solid application kernel
// src/spec.rs: Interface to Linux prctl(2) speculation misfeature interface
//
// Copyright (c) 2024, 2025 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

//! Set of functions to manage speculation misfeature

use std::fmt;

use nix::errno::Errno;

/// Speculative execution control constants
pub(crate) const PR_GET_SPECULATION_CTRL: libc::c_int = 52;
pub(crate) const PR_SET_SPECULATION_CTRL: libc::c_int = 53;

/// Speculation control identifiers
pub(crate) const PR_SPEC_STORE_BYPASS: u32 = 0;
pub(crate) const PR_SPEC_INDIRECT_BRANCH: u32 = 1;
pub(crate) const PR_SPEC_L1D_FLUSH: u32 = 2;

/// Status mask to extract bits 0-4
pub(crate) const SPECULATION_CTRL_MASK: u32 = 0x1F;

/// Speculative execution status flags
pub(crate) const PR_SPEC_NOT_AFFECTED: u32 = 0;
pub(crate) const PR_SPEC_PRCTL: u32 = 1 << 0;
pub(crate) const PR_SPEC_ENABLE: u32 = 1 << 1;
pub(crate) const PR_SPEC_DISABLE: u32 = 1 << 2;
pub(crate) const PR_SPEC_FORCE_DISABLE: u32 = 1 << 3;
pub(crate) const PR_SPEC_DISABLE_NOEXEC: u32 = 1 << 4;

/// Represents the different speculative execution features.
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpeculationFeature {
    /// Speculative Store Bypass
    StoreBypass = PR_SPEC_STORE_BYPASS,

    /// Indirect Branch Speculation in User Processes
    IndirectBranch = PR_SPEC_INDIRECT_BRANCH,

    /// Flush L1D Cache on context switch out of the task
    L1DFlush = PR_SPEC_L1D_FLUSH,
}

impl fmt::Display for SpeculationFeature {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let feature_str = match self {
            SpeculationFeature::StoreBypass => "Store Bypass",
            SpeculationFeature::IndirectBranch => "Indirect Branch",
            SpeculationFeature::L1DFlush => "L1D Flush",
        };
        write!(f, "{}", feature_str)
    }
}

/// Represents the status of speculative execution controls using bitflags.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SpeculationStatus(pub u32);

impl SpeculationStatus {
    /// Creates a new `SpeculationStatus` from a raw value.
    ///
    /// # Arguments
    ///
    /// * `val` - The raw status value obtained from `prctl`.
    pub fn from_raw(val: u32) -> Self {
        SpeculationStatus(val & SPECULATION_CTRL_MASK)
    }

    /// Checks if the status indicates that the CPU is not affected by the speculation misfeature.
    pub fn is_not_affected(&self) -> bool {
        self.0 == PR_SPEC_NOT_AFFECTED
    }

    /// Checks if `prctl` can set speculation mitigation.
    pub fn can_prctl_set(&self) -> bool {
        self.0 & PR_SPEC_PRCTL != 0
    }

    /// Checks if the speculation feature is enabled.
    pub fn is_enabled(&self) -> bool {
        self.0 & PR_SPEC_ENABLE != 0
    }

    /// Checks if the speculation feature is disabled.
    pub fn is_disabled(&self) -> bool {
        self.0 & PR_SPEC_DISABLE != 0
    }

    /// Checks if the speculation feature is force-disabled.
    pub fn is_force_disabled(&self) -> bool {
        self.0 & PR_SPEC_FORCE_DISABLE != 0
    }

    /// Checks if the speculation feature is exec-disabled.
    pub fn is_disable_noexec(&self) -> bool {
        self.0 & PR_SPEC_DISABLE_NOEXEC != 0
    }

    /// Retrieves the raw status value.
    pub fn raw(&self) -> u32 {
        self.0
    }
}

impl fmt::Display for SpeculationStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_not_affected() {
            return write!(f, "Not affected by speculation");
        }

        let mut statuses = Vec::new();

        if self.is_enabled() {
            statuses.push("Speculation feature is enabled, mitigation is disabled");
        } else if self.is_disabled() {
            statuses.push("Speculation feature is disabled, mitigation is enabled");
        } else if self.is_force_disabled() {
            statuses.push("Speculation feature is force-disabled, mitigation is enabled");
        } else if self.is_disable_noexec() {
            statuses.push("Speculation feature is exec-disabled, mitigation is enabled");
        }

        if self.can_prctl_set() {
            statuses.push("(prctl can set speculation mitigation)");
        }

        if statuses.is_empty() {
            write!(f, "Speculation feature status unknown: {:#X}", self.0)
        } else {
            write!(f, "{}.", statuses.join(" "))
        }
    }
}

/// Represents the status of a speculative execution control for a specific feature.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SpeculationControlStatus {
    /// The speculative execution feature.
    pub feature: SpeculationFeature,

    /// The status of the speculative execution feature.
    pub status: SpeculationStatus,
}

impl fmt::Display for SpeculationControlStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} status: {}", self.feature, self.status)
    }
}

/// Retrieves the current speculative execution status for a given feature.
pub fn speculation_get(feature: SpeculationFeature) -> Result<SpeculationControlStatus, Errno> {
    // SAFETY: nix does not support the speculation interface yet.
    let ret = Errno::result(unsafe {
        libc::prctl(PR_GET_SPECULATION_CTRL, feature as libc::c_int, 0, 0, 0)
    })?;

    #[allow(clippy::cast_sign_loss)]
    let masked = (ret as u32) & SPECULATION_CTRL_MASK;
    let status = SpeculationStatus::from_raw(masked);

    Ok(SpeculationControlStatus { feature, status })
}

/// Sets the speculative execution status for a given feature.
pub fn speculation_set(
    feature: SpeculationFeature,
    status: SpeculationStatus,
) -> Result<(), Errno> {
    #[allow(clippy::cast_lossless)]
    // SAFETY: nix does not support the speculation interface yet.
    Errno::result(unsafe {
        libc::prctl(
            PR_SET_SPECULATION_CTRL,
            feature as libc::c_int,
            status.raw() as libc::c_ulong,
            0,
            0,
        )
    })
    .map(drop)
}
