sec-mem 0.1.0

use core::{any, fmt::{self, Debug}};
use zeroize::{Zeroize, ZeroizeOnDrop};

use std::sync::OnceLock;

#[inline(always)]
fn global_canary() -> usize {
    static CANARY: OnceLock<usize> = OnceLock::new();
    *CANARY.get_or_init(|| {
        #[cfg(all(feature = "sec_mem", target_os = "linux"))]
        {
            let mut canary = 0usize;
            unsafe {
                if libc::getrandom(
                    &mut canary as *mut _ as *mut libc::c_void,
                    core::mem::size_of::<usize>(),
                    0,
                ) == core::mem::size_of::<usize>() as isize {
                    return canary;
                }
            }
        }

        #[cfg(target_arch = "x86_64")]
        {
            let mut canary = 0u64;
            unsafe {
                if core::arch::x86_64::_rdrand64_step(&mut canary) == 1 {
                    return canary as usize;
                }
            }
        }

        #[cfg(target_arch = "x86")]
        {
            let mut canary = 0u32;
            unsafe {
                if core::arch::x86::_rdrand32_step(&mut canary) == 1 {
                    return canary as usize;
                }
            }
        }

        // Fallback: ASLR pointer mixed with timestamp
        use std::time::SystemTime;
        let time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_nanos() as usize;
        let aslr = global_canary as *const () as usize;
        time.wrapping_mul(0x9E37_79B9_7F4A_7C15).wrapping_add(aslr)
    })
}

/// A pure-software, stack-native hardened wrapper for sensitive values (`no_std` compatible).
///
/// `SecretBox` protects inline data by surrounding it with dynamically generated, randomized canaries
/// (verified via compiler-optimization-proof `read_volatile` checks) to immediately panic during stack
/// smashing buffer overflows. 
///
/// It implements mathematically strict `&mut self` concurrency locks, forces zeroization on drop, and 
/// mandates that data can only be touched inside strictly scoped `.with_secret()` closures to 
/// eliminate accidental memory leaks. It strictly forbids `Clone` to prevent key scattering.
///
/// # Examples
/// ```
/// use sec_mem::SecretBox;
///
/// // Wrap a basic 32-bit integer or a 256-bit array entirely on the stack
/// let mut secret = SecretBox::new([0x42u8; 32]);
///
/// // The data can only be accessed exclusively inside a closure
/// secret.with_secret_mut(|s| {
///     s[0] = 0x99;
/// }); 
/// // Any stack corruption that occurred is detected here, panicking instantly.
/// ```
#[repr(C)]
pub struct SecretBox<S: Zeroize> {
    canary_front: usize,
    inner_secret: S,
    canary_back: usize,
}

impl<S: Zeroize> Zeroize for SecretBox<S> {
    fn zeroize(&mut self) {
        self.inner_secret.zeroize()
    }
}

impl<S: Zeroize> Drop for SecretBox<S> {
    fn drop(&mut self) {
        self.zeroize()
    }
}

impl<S: Zeroize> ZeroizeOnDrop for SecretBox<S> {}

impl<S: Zeroize> From<S> for SecretBox<S> {
    fn from(source: S) -> Self {
        Self::new(source)
    }
}

impl<S: Zeroize> SecretBox<S> {
    #[inline(always)]
    fn front_canary_expected() -> usize {
        global_canary()
    }

    #[inline(always)]
    fn back_canary_expected() -> usize {
        !global_canary()
    }

    /// Create a secret value using an inline value.
    pub fn new(secret: S) -> Self {
        Self {
            canary_front: Self::front_canary_expected(),
            inner_secret: secret,
            canary_back: Self::back_canary_expected(),
        }
    }

    #[inline(always)]
    fn check_canaries(&self) {
        unsafe {
            // Read volatile to absolutely prevent LLVM from dead-code-eliminating the check
            // The compiler cannot mathematically optimize away volatile reads!
            let front = core::ptr::read_volatile(&self.canary_front);
            let back = core::ptr::read_volatile(&self.canary_back);

            if front != Self::front_canary_expected() || back != Self::back_canary_expected() {
                panic!("SECURITY ALERT: SecretBox stack memory corruption detected!");
            }
        }
    }

    /// Safer alternative: Provide an exclusive closure-based interface.
    /// Requires `&mut self` to mathematically prevent concurrent multi-threaded 
    /// aliasing or reference sharing.
    pub fn with_secret<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&S) -> R,
    {
        self.check_canaries();
        f(&self.inner_secret)
    }

    /// Safer mutable alternative: Provide a closure-based interface.
    pub fn with_secret_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut S) -> R,
    {
        self.check_canaries();
        f(&mut self.inner_secret)
    }
}

impl<S: Zeroize + Default> SecretBox<S> {
    /// Create a secret value using a function that can initialize the value in-place.
    pub fn init_with_mut(ctr: impl FnOnce(&mut S)) -> Self {
        let mut secret = Self::default();
        secret.with_secret_mut(|s| ctr(s));
        secret
    }
}

impl<S: Zeroize + Default> Default for SecretBox<S> {
    fn default() -> Self {
        Self {
            canary_front: Self::front_canary_expected(),
            inner_secret: S::default(),
            canary_back: Self::back_canary_expected(),
        }
    }
}

impl<S: Zeroize> Debug for SecretBox<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SecretBox<{}>([REDACTED])", any::type_name::<S>())
    }
}

// Clone is strictly removed. 
// A high-security SecretBox should never blindly duplicate memory allocations.
// If duplication is required, the user must explicitly construct a new SecretBox.

impl<S: Zeroize> SecretBox<S> {
    /// Performs a constant-time comparison of two SecretBoxes.
    /// Requires `&mut` on both to enforce exclusive access and check canaries.
    pub fn constant_time_eq(&mut self, other: &mut Self) -> subtle::Choice 
    where 
        S: subtle::ConstantTimeEq 
    {
        self.check_canaries();
        other.check_canaries();
        self.inner_secret.ct_eq(&other.inner_secret)
    }
}