sec-mem 0.1.1

#![cfg(unix)]
#![forbid(unsafe_op_in_unsafe_fn)]

use libc::{
    madvise, mlock, mmap, mprotect, munlock, munmap, 
    MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_NONE, PROT_READ, PROT_WRITE,
};
use std::{
    fmt,
    mem::size_of,
    ptr::{self, NonNull},
    sync::atomic::{AtomicBool, AtomicUsize, Ordering, Ordering::SeqCst},
};
use subtle::{Choice, ConstantTimeEq};
use zeroize::{Zeroize, ZeroizeOnDrop};

// Platform-specific imports
#[cfg(target_os = "linux")]
use libc::{sysconf, _SC_LEVEL1_DCACHE_LINESIZE};

// --------------------------------------------------------------------------------
// Global Constants and Statics
// --------------------------------------------------------------------------------

/// Cache line size for the current platform (initialized once)
static CACHE_LINE_SIZE: std::sync::OnceLock<usize> = std::sync::OnceLock::new();

/// Global counter for tracking lock IDs
static GLOBAL_LOCK_COUNTER: AtomicUsize = AtomicUsize::new(0);

/// Platform-specific MAP_LOCKED flag
#[cfg(target_os = "linux")]
const MAP_LOCKED: i32 = libc::MAP_LOCKED;
#[cfg(not(target_os = "linux"))]
const MAP_LOCKED: i32 = 0;

// --------------------------------------------------------------------------------
// Helper Types
// --------------------------------------------------------------------------------

/// RAII wrapper for memory-mapped regions
struct MmapRegion {
    ptr: NonNull<libc::c_void>,
    size: usize,
    #[cfg(target_os = "linux")]
    locked: bool,
}

impl MmapRegion {
    /// Creates a new memory-mapped region with specified protection
    fn new(size: usize, prot: i32) -> Self {
        let mut flags = MAP_PRIVATE | MAP_ANONYMOUS;
        #[cfg(target_os = "linux")]
        let mut locked = false;
        
        // On Linux, try to use MAP_LOCKED for atomic allocation+locking
        #[cfg(target_os = "linux")]
        {
            flags |= MAP_LOCKED;
        }

        let ptr = unsafe { mmap(ptr::null_mut(), size, prot, flags, -1, 0) };

        // Fallback for Linux if MAP_LOCKED fails
        #[cfg(target_os = "linux")]
        if ptr == MAP_FAILED && (flags & MAP_LOCKED)  != 0 {
            let ptr = unsafe { mmap(ptr::null_mut(), size, prot, flags ^ MAP_LOCKED, -1, 0) };
            if ptr != MAP_FAILED {
                if unsafe { mlock(ptr, size) } == 0 {
                    locked = true;
                } else {
                    unsafe { munmap(ptr, size) };
                    panic!("Failed to lock memory region after fallback mmap");
                }
                return Self {
                    ptr: NonNull::new(ptr).expect("mmap returned non-null"),
                    size,
                    locked,
                };
            }
        }

        if ptr == MAP_FAILED {
            panic!("Failed to allocate memory region");
        }

        Self {
            ptr: NonNull::new(ptr).expect("mmap returned non-null"),
            size,
            #[cfg(target_os = "linux")]
            locked,
        }
    }

    /// Returns raw pointer to the mapped region
    #[allow(dead_code)]
    fn as_ptr(&self) -> *mut libc::c_void {
        self.ptr.as_ptr()
    }
}

impl Drop for MmapRegion {
    fn drop(&mut self) {
        if self.size > 0 {
            // Only zero memory if it was writable
            let prot = PROT_READ | PROT_WRITE;
            if secure_mprotect(self.ptr.as_ptr(), self.size, prot)  == 0 {
                unsafe { ptr::write_bytes(self.ptr.as_ptr() as *mut u8, 0, self.size) };
            }
            
            #[cfg(target_os = "linux")]
            if self.locked {
                unsafe { munlock(self.ptr.as_ptr(), self.size) };
            }
            
            let result = unsafe { munmap(self.ptr.as_ptr(), self.size) };
            if result != 0 {
                eprintln!("Failed to unmap memory region");
            }
        }
    }
}

/// RAII guard for temporary memory protection changes
struct ProtectionGuard {
    mapping: *mut libc::c_void,
    size: usize,
    original_prot: i32,
}

impl ProtectionGuard {
    /// Creates new guard and sets specified protection
    fn new(mapping: *mut libc::c_void, size: usize) -> Self {
        // Since SecMem always keeps memory at PROT_NONE when not accessed,
        // we can safely assume we should restore to PROT_NONE.
        let original_prot = PROT_NONE;
        
        let guard = Self {
            mapping,
            size,
            original_prot,
        };
        
        // Only change protection if we have something to protect
        if size > 0 {
            unsafe { set_pkey_rights(get_global_pkey(), 0); } // 0 = allow read/write
            if secure_mprotect(mapping, size, PROT_READ | PROT_WRITE) != 0 {
                panic!("Failed to set memory protection");
            }
        }
        
        guard
    }
}

impl Drop for ProtectionGuard {
    fn drop(&mut self) {
        if self.size > 0 {
            if secure_mprotect(self.mapping, self.size, self.original_prot) != 0 {
                eprintln!("Failed to restore memory protection");
            }
            unsafe { set_pkey_rights(get_global_pkey(), 3); } // 3 = disable access
        }
    }
}

// --------------------------------------------------------------------------------
// Global Hardening
// --------------------------------------------------------------------------------

/// Hardens the entire process against memory dumping and tracing (Linux-only).
/// This applies `prctl(PR_SET_DUMPABLE, 0)` to disable `ptrace` and core dumps globally.
pub fn harden_process() -> Result<(), &'static str> {
    #[cfg(target_os = "linux")]
    unsafe {
        // PR_SET_DUMPABLE = 4
        if libc::prctl(4, 0, 0, 0, 0) != 0 {
            return Err("Failed to set PR_SET_DUMPABLE");
        }
        Ok(())
    }
    #[cfg(not(target_os = "linux"))]
    {
        Ok(())
    }
}


// --------------------------------------------------------------------------------
// MPK (Memory Protection Keys) Module
// --------------------------------------------------------------------------------

#[cfg(target_os = "linux")]
fn get_global_pkey() -> i32 {
    static GLOBAL_PKEY: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
    *GLOBAL_PKEY.get_or_init(|| {
        unsafe {
            // Allocate a key with PKEY_DISABLE_ACCESS (1) | PKEY_DISABLE_WRITE (2) = 3
            let pkey = libc::syscall(330, 0, 3); // SYS_pkey_alloc for x86_64
            if pkey >= 0 {
                pkey as i32
            } else {
                -1
            }
        }
    })
}

#[cfg(not(target_os = "linux"))]
fn get_global_pkey() -> i32 {
    -1
}

#[cfg(target_arch = "x86_64")]
unsafe fn set_pkey_rights(pkey: i32, rights: u32) {
    if pkey < 0 { return; }
    let mut pkru: u32;
    unsafe { std::arch::asm!(
        "rdpkru",
        out("eax") pkru,
        in("ecx") 0,
        out("edx") _,
    ); }
    let shift = pkey * 2;
    pkru &= !(3 << shift);
    pkru |= (rights & 3) << shift;
    unsafe { std::arch::asm!(
        "wrpkru",
        in("eax") pkru,
        in("ecx") 0,
        in("edx") 0,
    ); }
}

#[cfg(not(target_arch = "x86_64"))]
unsafe fn set_pkey_rights(_pkey: i32, _rights: u32) {
    // Fallback for non-x86_64
}

fn secure_mprotect(addr: *mut libc::c_void, len: usize, prot: i32) -> i32 {
    let pkey = get_global_pkey();
    #[cfg(target_os = "linux")]
    if pkey >= 0 {
        let res = unsafe { libc::syscall(329, addr, len, prot, pkey) as i32 }; // SYS_pkey_mprotect
        if res == 0 {
            return 0;
        }
    }
    unsafe { mprotect(addr, len, prot) }
}

// --------------------------------------------------------------------------------
// Main SecMem Implementation
// --------------------------------------------------------------------------------

/// A hardware-accelerated secure memory container for isolating sensitive data from the OS.
///
/// `SecMem` uses raw Linux syscalls (`mmap`, `mprotect`, `mseal`, `mlock`) and Intel Memory Protection Keys
/// (MPK/`WRPKRU`) to physically isolate cryptographic keys or passwords from the rest of the process.
/// It strictly guards against buffer overflows via `mseal`ed guard pages, and actively prevents memory
/// dumping using `MADV_DONTDUMP`, `MADV_DONTFORK`, and `PR_SET_DUMPABLE(0)`.
///
/// When the `encryption` feature is enabled, data is permanently blinded at rest in RAM using ChaCha20. 
/// It utilizes OpenSSH-style Key Shielding, splitting a massive 16 KiB pre-key buffer across kernel `memfd_secret` 
/// mappings and dynamically deriving the true key using `blake3`. This mathematically neutralizes `/proc/self/mem` 
/// dumping, physical cold-boot attacks, and side-channel exploits like Spectre or Rowhammer.
///
/// # Examples
/// ```
/// use sec_mem::SecMem;
///
/// // Create a highly protected memory vault
/// let mut secure_key = SecMem::new([0xAAu8; 32]);
///
/// // Access is granted exclusively within this closure bounds
/// secure_key.access_mut(|key| {
///     key[0] = 0xBB;
///     // Data is processed here...
/// }); 
/// // Instantly upon closure exit, memory is zeroized, re-encrypted, and hardware-locked.
/// ```
pub struct SecMem<S: Zeroize> {
    // Memory mapping details
    mapping: NonNull<libc::c_void>,
    mapping_size: usize,

    // Secret data pointer
    secret_ptr: NonNull<S>,

    // Guard pages
    guard_pages: [NonNull<libc::c_void>; 2],
    guard_page_size: usize,

    // Debug checks
    #[cfg(debug_assertions)]
    canary: u64,

    // Access control
    locked: AtomicBool,
    #[allow(dead_code)]
    lock_id: usize,

    #[cfg(feature = "encryption")]
    nonce: [u8; 12],
    #[cfg(feature = "encryption")]
    is_encrypted: bool,
}

impl<S: Zeroize> SecMem<S> {
    /// Creates a new protected secret
    pub fn new(secret: S) -> Self {
        // Handle zero-sized types specially
        let size = size_of::<S>();
        if size == 0 {
            return Self {
                mapping: NonNull::dangling(),
                mapping_size: 0,
                secret_ptr: NonNull::dangling(),
                guard_pages: [NonNull::dangling(), NonNull::dangling()],
                guard_page_size: 0,
                #[cfg(debug_assertions)]
                canary: 0,
                locked: AtomicBool::new(false),
                lock_id: 0,
                #[cfg(feature = "encryption")]
                nonce: [0; 12],
                #[cfg(feature = "encryption")]
                is_encrypted: false,
            };
        }

        let lock_id = GLOBAL_LOCK_COUNTER.fetch_add(1, SeqCst);
        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as usize;

        // Calculate allocation size (round up to nearest page)
        let alloc_size = (size + page_size - 1) & !(page_size - 1);
        let total_size = alloc_size + 2 * page_size;

        // Allocate a single contiguous protected memory region
        let full_region = MmapRegion::new(total_size, PROT_NONE);

        let guard_before_ptr = full_region.ptr;
        let secret_region_ptr = unsafe { NonNull::new((full_region.ptr.as_ptr() as *mut u8).add(page_size) as *mut libc::c_void).unwrap() };
        let guard_after_ptr = unsafe { NonNull::new((full_region.ptr.as_ptr() as *mut u8).add(page_size + alloc_size) as *mut libc::c_void).unwrap() };

        // Make the middle region writable for initialization
        if secure_mprotect(secret_region_ptr.as_ptr(), alloc_size, PROT_READ | PROT_WRITE)  != 0 {
            panic!("Failed to set memory protection for secret region");
        }

        // Additional memory hardening applied to the entire region (including guard pages)
        // to prevent core dumps and inheriting memory across forks.
        #[cfg(target_os = "linux")]
        unsafe {
            if madvise(full_region.ptr.as_ptr(), total_size, libc::MADV_DONTDUMP) != 0 {
                println!("Failed to set MADV_DONTDUMP on memory region");
            }
            if madvise(full_region.ptr.as_ptr(), total_size, libc::MADV_DONTFORK) != 0 {
                println!("Failed to set MADV_DONTFORK on memory region");
            }
        }
        #[cfg(not(target_os = "linux"))]
        unsafe {
            if madvise(full_region.ptr.as_ptr(), total_size, libc::MADV_DONTDUMP) != 0 {
                println!("Failed to set MADV_DONTDUMP on memory region");
            }
        }

        if unsafe { mlock(secret_region_ptr.as_ptr(), alloc_size) } != 0 {
            println!("Failed to lock secret region in memory");
        }

        // Initialize secret
        let secret_ptr = secret_region_ptr.as_ptr() as *mut S;
        unsafe { ptr::write(secret_ptr, secret) };

        #[cfg(feature = "encryption")]
        let mut nonce = [0u8; 12];
        #[cfg(feature = "encryption")]
        let is_encrypted = true;

        #[cfg(feature = "encryption")]
        {
            get_random_bytes(&mut nonce);
            let secret_bytes = unsafe { std::slice::from_raw_parts_mut(secret_ptr as *mut u8, size)
            };
            encrypt_decrypt_memory(secret_bytes, &nonce);
        }

        // Lock down memory after initialization
        if secure_mprotect(secret_region_ptr.as_ptr(), alloc_size, PROT_NONE)  != 0 {
            unsafe { ptr::drop_in_place(secret_ptr) };
            panic!("Failed to set memory protection after secret initialization");
        }

        // Extract pointers before forgetting regions
        std::mem::forget(full_region);

        Self {
            mapping: secret_region_ptr,
            mapping_size: alloc_size,
            secret_ptr: NonNull::new(secret_ptr).unwrap(),
            guard_pages: [guard_before_ptr, guard_after_ptr],
            guard_page_size: page_size,
            #[cfg(debug_assertions)]
            canary: 0xDEADBEEFCAFEBABE,
            locked: AtomicBool::new(true),
            lock_id,
            #[cfg(feature = "encryption")]
            nonce,
            #[cfg(feature = "encryption")]
            is_encrypted,
        }
    }

    /// Access secret for reading
    pub fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&S) -> R,
    {
        self.check_canary();

        #[cfg(feature = "encryption")]
        {
            while !self.locked.swap(false, Ordering::SeqCst) {
                std::hint::spin_loop();
            }
        }

        struct AccessGuard<'a, S: Zeroize> {
            sec_mem: &'a SecMem<S>,
        }
        
        impl<'a, S: Zeroize> Drop for AccessGuard<'a, S> {
            fn drop(&mut self) {
                unsafe { set_pkey_rights(get_global_pkey(), 3); } // Hardware Lock ON

                #[cfg(feature = "encryption")]
                if self.sec_mem.is_encrypted {
                    unsafe { set_pkey_rights(get_global_pkey(), 0); } // Temporarily unlock for encryption
                    let secret_bytes = unsafe {
                        std::slice::from_raw_parts_mut(self.sec_mem.secret_ptr.as_ptr() as *mut u8, size_of::<S>())
                    };
                    encrypt_decrypt_memory(secret_bytes, &self.sec_mem.nonce);
                    unsafe { set_pkey_rights(get_global_pkey(), 3); } // Hardware Lock ON
                }

                unsafe {
                    flush_cache(
                        self.sec_mem.secret_ptr.as_ptr() as *const u8,
                        size_of::<S>()
                    );
                }

                #[cfg(feature = "encryption")]
                {
                    self.sec_mem.locked.store(true, Ordering::SeqCst);
                }
            }
        }

        self.set_protection(PROT_READ | PROT_WRITE); // Must be writable for decryption
        let _prot_guard = ProtectionGuard::new(self.mapping.as_ptr(), self.mapping_size);
        
        let _access_guard = AccessGuard { sec_mem: self };

        #[cfg(feature = "encryption")]
        if self.is_encrypted {
            let secret_bytes = unsafe {
                std::slice::from_raw_parts_mut(self.secret_ptr.as_ptr() as *mut u8, size_of::<S>())
            };
            encrypt_decrypt_memory(secret_bytes, &self.nonce);
        }

        unsafe {
            core::sync::atomic::compiler_fence(Ordering::SeqCst);
            let res = f(self.secret_ptr.as_ref());
            core::sync::atomic::compiler_fence(Ordering::SeqCst);
            res
        }
    }

    /// Access secret for mutation
    pub fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut S) -> R,
    {
        if !self.locked.swap(false, Ordering::SeqCst) {
            panic!("Attempted to create multiple mutable references to secret data");
        }

        struct MutGuard<'a, S: Zeroize> {
            box_ref: &'a mut SecMem<S>,
        }
        
        impl<'a, S: Zeroize> Drop for MutGuard<'a, S> {
            fn drop(&mut self) {
                #[cfg(feature = "encryption")]
                if self.box_ref.is_encrypted {
                    let secret_bytes = unsafe {
                        std::slice::from_raw_parts_mut(self.box_ref.secret_ptr.as_ptr() as *mut u8, size_of::<S>())
                    };
                    get_random_bytes(&mut self.box_ref.nonce);
                    encrypt_decrypt_memory(secret_bytes, &self.box_ref.nonce);
                }

                unsafe {
                    flush_cache(
                        self.box_ref.secret_ptr.as_ptr() as *const u8,
                        size_of::<S>()
                    );
                }
                self.box_ref.locked.store(true, Ordering::SeqCst);
            }
        }

        self.check_canary();
        self.set_protection(PROT_READ | PROT_WRITE);
        let _prot_guard = ProtectionGuard::new(self.mapping.as_ptr(), self.mapping_size);
        
        let mut _mut_guard = MutGuard { box_ref: self };

        #[cfg(feature = "encryption")]
        if _mut_guard.box_ref.is_encrypted {
            let secret_bytes = unsafe {
                std::slice::from_raw_parts_mut(_mut_guard.box_ref.secret_ptr.as_ptr() as *mut u8, size_of::<S>())
            };
            encrypt_decrypt_memory(secret_bytes, &_mut_guard.box_ref.nonce);
        }

        let result = unsafe { f(_mut_guard.box_ref.secret_ptr.as_mut()) };
        
        _mut_guard.box_ref.check_canary();
        result
    }

    /// Constant-time equality comparison
    pub fn constant_time_eq(&self, other: &Self) -> Choice
    where
        S: AsRef<[u8]>,
    {
        self.access(|s| {
            other.access(|o| {
                let s_bytes = s.as_ref();
                let o_bytes = o.as_ref();

                let len_equal = s_bytes.len().ct_eq(&o_bytes.len());

                let min_len = s_bytes.len().min(o_bytes.len());
                let content_equal = s_bytes[..min_len].ct_eq(&o_bytes[..min_len]);

                len_equal & content_equal
            })
        })
    }

    /// Explicitly applies Linux memory sealing (`mseal`) to the guard pages.
    /// 
    /// **WARNING:** `mseal` permanently locks the memory region's permissions. 
    /// Because Linux prevents `munmap` on sealed memory, calling this method 
    /// means the guard pages (and thus the allocation) will leak and remain in RAM 
    /// until the process terminates. It should only be used for global, long-lived 
    /// singletons (like root application keys).
    /// Requires Linux 6.10+. Fails silently on older kernels.
    pub fn seal_guard_pages(&self) {
        #[cfg(target_os = "linux")]
        unsafe {
            // SYS_mseal = 462 for x86_64, aarch64, etc.
            libc::syscall(462, self.guard_pages[0].as_ptr(), self.guard_page_size, 0);
            libc::syscall(462, self.guard_pages[1].as_ptr(), self.guard_page_size, 0);
        }
    }

    // -------------------------------------------------------------------------
    // Internal Helpers
    // -------------------------------------------------------------------------

    /// Debug check for memory corruption
    #[cfg(debug_assertions)]
    fn check_canary(&self) {
        if self.canary != 0xDEADBEEFCAFEBABE {
            panic!("Memory corruption detected (canary check failed)");
        }
    }

    #[cfg(not(debug_assertions))]
    fn check_canary(&self) {}

    /// Updates memory protection flags
    fn set_protection(&self, protection: i32) {
        if size_of::<S>() == 0 {
            return;
        }

        if secure_mprotect(self.mapping.as_ptr(), self.mapping_size, protection) != 0 {
            println!("Failed to set memory protection");
        }
    }
}

// --------------------------------------------------------------------------------
// Trait Implementations
// --------------------------------------------------------------------------------

impl<S: Zeroize> Drop for SecMem<S> {
    fn drop(&mut self) {
        let size = size_of::<S>();
        if size == 0 {
            return;
        }

        // Check for double-free
        if !self.locked.swap(false, Ordering::SeqCst) {
            #[cfg(debug_assertions)]
            panic!("Double-free detected for SecMem with lock_id {}", self.lock_id);
            
            #[cfg(not(debug_assertions))]
            {
                println!("[SECURITY CRITICAL] Double-free attempt detected");
                return;
            }
        }

        // Enable write access for secure cleanup
        if secure_mprotect(
            self.mapping.as_ptr(),
            self.mapping_size,
            PROT_READ | PROT_WRITE
        ) != 0 {
            let _ = unsafe { mlock(self.mapping.as_ptr(), self.mapping_size) };
            println!("[SECURITY CRITICAL] Failed to restore memory protection during drop");
            return;
        }

        #[cfg(feature = "encryption")]
        if self.is_encrypted {
            let secret_bytes = unsafe {
                std::slice::from_raw_parts_mut(self.secret_ptr.as_ptr() as *mut u8, size)
            };
            encrypt_decrypt_memory(secret_bytes, &self.nonce);
        }

        // Securely erase secret
        unsafe {
            let secret = self.secret_ptr.as_mut();
            secret.zeroize();
            
            flush_cache(self.secret_ptr.as_ptr() as *const u8, size);
            ptr::drop_in_place(secret);
            flush_cache(self.secret_ptr.as_ptr() as *const u8, size);
        }

        // Clean up memory
        if unsafe { munlock(self.mapping.as_ptr(), self.mapping_size) } != 0 {
            println!("Failed to unlock memory during drop");
        }

        unsafe {
            let total_size = self.mapping_size + 2 * self.guard_page_size;
            
            #[cfg(target_os = "linux")]
            libc::madvise(self.guard_pages[0].as_ptr(), total_size, libc::MADV_REMOVE);
            
            munmap(self.guard_pages[0].as_ptr(), total_size);
        }
    }
}

impl<S: Zeroize + Default> Default for SecMem<S> {
    fn default() -> Self {
        Self::new(S::default())
    }
}

impl<S: Zeroize + Clone> Clone for SecMem<S> {
    fn clone(&self) -> Self {
        self.access(|s| Self::new(s.clone()))
    }
}

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

impl<S: Zeroize> Zeroize for SecMem<S> {
    fn zeroize(&mut self) {
        let size = size_of::<S>();
        if size == 0 {
            return;
        }

        // Temporarily enable write access
        if secure_mprotect(self.mapping.as_ptr(), self.mapping_size, PROT_READ | PROT_WRITE) != 0 {
            println!("[SECURITY WARNING] Failed to enable write protection during zeroization");
            return;
        }

        #[cfg(feature = "encryption")]
        if self.is_encrypted {
            let secret_bytes = unsafe {
                std::slice::from_raw_parts_mut(self.secret_ptr.as_ptr() as *mut u8, size)
            };
            encrypt_decrypt_memory(secret_bytes, &self.nonce);
            self.is_encrypted = false; // Prevent re-decryption on Drop
        }

        // Securely erase contents
        unsafe { self.secret_ptr.as_mut().zeroize();
            flush_cache(self.secret_ptr.as_ptr() as *const u8, size);
        }

        // Restore protection
        if secure_mprotect(self.mapping.as_ptr(), self.mapping_size, PROT_NONE)  != 0 {
            println!("[SECURITY WARNING] Failed to restore protection after zeroization");
            let _ = unsafe { mlock(self.mapping.as_ptr(), self.mapping_size) };
        }
    }
}

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

unsafe impl<S: Zeroize + Send> Send for SecMem<S> {}
unsafe impl<S: Zeroize + Sync> Sync for SecMem<S> {}

// --------------------------------------------------------------------------------
// Custom Implementations
// --------------------------------------------------------------------------------


// --------------------------------------------------------------------------------
// Utility Functions
// --------------------------------------------------------------------------------

/// Gets platform-specific cache line size with proper error handling
fn get_cache_line_size() -> usize {
    CACHE_LINE_SIZE.get_or_init(|| {
        #[cfg(target_os = "linux")]
        {
            // Safe because we're reading a system constant
            match unsafe { sysconf(_SC_LEVEL1_DCACHE_LINESIZE) } {
                size if size > 0 => size as usize,
                _ => fallback_cache_line_size(),
            }
        }
        
        #[cfg(not(target_os = "linux"))]
        {
            fallback_cache_line_size()
        }
    }).to_owned()  // Returns a copy of the value
}

/// Default cache line size when detection fails
const fn fallback_cache_line_size() -> usize {
    // Common cache line sizes: 64 (x86), 128 (ARM), 64 (others)
    64
}

/// Safely flushes cache for the given memory region
///
/// # Safety
/// - `ptr` must be valid for reads of `len` bytes
/// - `ptr` must be properly aligned
/// - The memory region must not be concurrently modified
unsafe fn flush_cache(ptr: *const u8, len: usize) {
    if len == 0 {
        return;
    }

    // Validate pointer alignment
    let alignment = get_cache_line_size();
    if !(ptr as usize).is_multiple_of(alignment) {
        panic!("Pointer is not aligned to cache line boundary");
    }

    #[cfg(target_arch = "x86_64")]
    {
        use core::arch::x86_64::_mm_clflush;
        
        let mut addr = ptr as usize;
        let end_addr = addr.saturating_add(len);
        
        while addr < end_addr {
            unsafe { _mm_clflush(addr as *const _) };
            addr = addr.saturating_add(alignment);
        }
        core::sync::atomic::compiler_fence(Ordering::SeqCst);
    }

    #[cfg(target_arch = "x86")]
    {
        use core::arch::x86::_mm_clflush;
        
        let mut addr = ptr as usize;
        let end_addr = addr.saturating_add(len);
        
        while addr < end_addr {
            unsafe { _mm_clflush(addr as *const _) };
            addr = addr.saturating_add(alignment);
        }
        core::sync::atomic::compiler_fence(Ordering::SeqCst);
    }

    #[cfg(target_arch = "aarch64")]
    {
        let mut addr = ptr as usize;
        let end_addr = addr.saturating_add(len);
        
        while addr < end_addr {
            unsafe { core::arch::asm!("dc cvau, {}", in(reg) addr) };
            addr = addr.saturating_add(alignment);
        }
        unsafe {
            core::arch::asm!("dsb ish");
            core::arch::asm!("isb");
        }
    }

    #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
    {
        // Explicitly ignore parameters to prevent unused warnings
        let _ = (ptr, len);
        // No-op on unsupported architectures
    }
}

// --------------------------------------------------------------------------------
// Uniform Access Trait
// --------------------------------------------------------------------------------

/// Trait for types that can provide access to underlying data, whether protected or raw
pub trait SecretAccess<T: Zeroize + ?Sized> {
    /// Access the underlying data for reading
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&T) -> R;
    
    /// Access the underlying data for mutation
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R;
}

// Primary Implementations for [u8] (Most Common Case)

// SecMem<[u8; N]> as [u8]
impl<const N: usize> SecretAccess<[u8]> for SecMem<[u8; N]> {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        self.access(|array| f(array.as_slice()))
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        self.access_mut(|array| f(array.as_mut_slice()))
    }
}

// SecMem<Vec<u8>> as [u8]
impl SecretAccess<[u8]> for SecMem<Vec<u8>> {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        self.access(|vec| f(vec.as_slice()))
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        self.access_mut(|vec| f(vec.as_mut_slice()))
    }
}

// [u8; N] as [u8]
impl<const N: usize> SecretAccess<[u8]> for [u8; N] {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(self.as_slice())
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        f(self.as_mut_slice())
    }
}

// &[u8] as [u8]
impl SecretAccess<[u8]> for &[u8] {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(self)
    }
    
    fn access_mut<F, R>(&mut self, _f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        panic!("Cannot get mutable access through immutable reference")
    }
}

// &mut [u8] as [u8]
impl SecretAccess<[u8]> for &mut [u8] {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(self)
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        f(self)
    }
}

// Box<[u8]> as [u8]
impl SecretAccess<[u8]> for Box<[u8]> {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(self.as_ref())
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        f(self.as_mut())
    }
}

// --------------------------------------------------------------------------------
// Fallback Implementations for Exact Types (Remove conflicting ones)
// --------------------------------------------------------------------------------

// SecMem<T> as T (fallback for non-u8 types)
impl<T: Zeroize> SecretAccess<T> for SecMem<T> {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        self.access(f)
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        self.access_mut(f)
    }
}

// [T; N] as [T; N] (fallback for non-u8 arrays)
impl<T: Zeroize, const N: usize> SecretAccess<[T; N]> for [T; N] {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[T; N]) -> R,
    {
        f(self)
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [T; N]) -> R,
    {
        f(self)
    }
}

// [T] as [T] (fallback for non-u8 slices)
impl<T: Zeroize> SecretAccess<[T]> for [T] where [T]: Zeroize {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[T]) -> R,
    {
        f(self)
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [T]) -> R,
    {
        f(self)
    }
}

// Vec<T> as [T] (fallback for non-u8 vectors)
impl<T: Zeroize> SecretAccess<[T]> for Vec<T> where [T]: Zeroize {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&[T]) -> R,
    {
        f(self.as_slice())
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut [T]) -> R,
    {
        f(self.as_mut_slice())
    }
}

// Box<T> as T (fallback)
impl<T: Zeroize> SecretAccess<T> for Box<T> {
    fn access<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        f(self.as_ref())
    }
    
    fn access_mut<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        f(self.as_mut())
    }
}

// --------------------------------------------------------------------------------
// Encryption-at-Rest Module
// --------------------------------------------------------------------------------

#[cfg(feature = "encryption")]
struct KeyStorage {
    ptr: *mut u8,
    size: usize,
}

#[cfg(feature = "encryption")]
unsafe impl Sync for KeyStorage {}
#[cfg(feature = "encryption")]
unsafe impl Send for KeyStorage {}

#[cfg(feature = "encryption")]
impl KeyStorage {
    fn new_uninitialized() -> Self {
        #[cfg(feature = "key_shielding")]
        let size = 16384; // 16 KiB for Key Shielding against Spectre/Rowhammer
        #[cfg(not(feature = "key_shielding"))]
        let size = 32;

        let mut ptr = libc::MAP_FAILED;

        #[cfg(target_os = "linux")]
        unsafe {
            // SYS_memfd_secret = 447 (x86_64, aarch64)
            let fd = libc::syscall(447, 0); 
            if fd >= 0 {
                if libc::ftruncate(fd as i32, size as libc::off_t) == 0 {
                    ptr = libc::mmap(
                        std::ptr::null_mut(),
                        size,
                        libc::PROT_READ | libc::PROT_WRITE,
                        libc::MAP_SHARED,
                        fd as i32,
                        0,
                    );
                }
                libc::close(fd as i32);
            }
        }

        if ptr == libc::MAP_FAILED {
            unsafe {
                let flags = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS;
                ptr = libc::mmap(
                    std::ptr::null_mut(),
                    size,
                    libc::PROT_READ | libc::PROT_WRITE,
                    flags,
                    -1,
                    0,
                );
                if ptr == libc::MAP_FAILED {
                    panic!("Failed to allocate global encryption key segment");
                }
                libc::mlock(ptr, size);
                #[cfg(target_os = "linux")]
                libc::madvise(ptr, size, libc::MADV_DONTDUMP);
                #[cfg(target_os = "linux")]
                libc::madvise(ptr, size, libc::MADV_DONTFORK);
            }
        }

        Self {
            ptr: ptr as *mut u8,
            size,
        }
    }
}

#[cfg(feature = "encryption")]
struct SplitKeyStorage {
    mask: KeyStorage,
    blinded: KeyStorage,
}

#[cfg(feature = "encryption")]
impl SplitKeyStorage {
    fn new() -> Self {
        let mask = KeyStorage::new_uninitialized();
        let blinded = KeyStorage::new_uninitialized();

        #[cfg(feature = "key_shielding")]
        let (mut k, mut m, mut b, copy_size) = {
            (vec![0u8; 16384], vec![0u8; 16384], vec![0u8; 16384], 16384)
        };
        #[cfg(not(feature = "key_shielding"))]
        let (mut k, mut m, mut b, copy_size) = {
            (vec![0u8; 32], vec![0u8; 32], vec![0u8; 32], 32)
        };
        
        get_random_bytes(&mut k);
        get_random_bytes(&mut m);
        
        for i in 0..copy_size {
            b[i] = k[i] ^ m[i];
        }

        unsafe {
            std::ptr::copy_nonoverlapping(m.as_ptr(), mask.ptr, copy_size);
            std::ptr::copy_nonoverlapping(b.as_ptr(), blinded.ptr, copy_size);

            libc::mprotect(mask.ptr as *mut libc::c_void, mask.size, libc::PROT_READ);
            libc::mprotect(blinded.ptr as *mut libc::c_void, blinded.size, libc::PROT_READ);
        }

        zeroize::Zeroize::zeroize(&mut k);
        zeroize::Zeroize::zeroize(&mut m);
        zeroize::Zeroize::zeroize(&mut b);

        Self { mask, blinded }
    }
}

#[cfg(feature = "encryption")]
fn get_split_keys() -> &'static SplitKeyStorage {
    static GLOBAL_SPLIT_KEY: std::sync::OnceLock<SplitKeyStorage> = std::sync::OnceLock::new();
    GLOBAL_SPLIT_KEY.get_or_init(SplitKeyStorage::new)
}

#[cfg(feature = "encryption")]
fn get_random_bytes(buf: &mut [u8]) {
    #[cfg(target_os = "linux")]
    unsafe {
        let mut total = 0;
        while total < buf.len() {
            let res = libc::getrandom(
                buf.as_mut_ptr().add(total) as *mut libc::c_void,
                buf.len() - total,
                0,
            );
            if res > 0 {
                total += res as usize;
            } else {
                // Fallback to /dev/urandom
                use std::fs::File;
                use std::io::Read;
                let mut file = File::open("/dev/urandom").expect("Failed to open /dev/urandom");
                file.read_exact(buf).expect("Failed to read random bytes");
                return;
            }
        }
    }
    
    #[cfg(not(target_os = "linux"))]
    {
        use std::fs::File;
        use std::io::Read;
        let mut file = File::open("/dev/urandom").expect("Failed to open /dev/urandom");
        file.read_exact(buf).expect("Failed to read random bytes");
    }
}

#[cfg(feature = "encryption")]
fn encrypt_decrypt_memory(data: &mut [u8], nonce: &[u8; 12]) {
    use chacha20::cipher::{KeyIvInit, StreamCipher};
    use chacha20::ChaCha20;
    
    let storage = get_split_keys();
    let mut key = [0u8; 32];
    
    #[cfg(feature = "key_shielding")]
    let mut combined = vec![0u8; 16384];
    
    unsafe {
        let m_ptr = storage.mask.ptr;
        let b_ptr = storage.blinded.ptr;
        
        #[cfg(feature = "key_shielding")]
        {
            // Reconstruct the 16 KiB pre-key buffer directly into a local Vec
            for i in 0..16384 {
                combined[i] = *m_ptr.add(i) ^ *b_ptr.add(i);
            }
        }
        
        #[cfg(not(feature = "key_shielding"))]
        {
            // Reconstruct the true key K = M ^ B directly into a local stack register
            for i in 0..32 {
                key[i] = *m_ptr.add(i) ^ *b_ptr.add(i);
            }
        }
    }
    
    #[cfg(feature = "key_shielding")]
    {
        // Hash the 16 KiB buffer with Blake3 to derive the true 32-byte ChaCha20 key
        let mut hasher = blake3::Hasher::new();
        hasher.update(&combined);
        key.copy_from_slice(hasher.finalize().as_bytes());
    }
    
    let mut cipher = ChaCha20::new((&key).into(), nonce.into());
    cipher.apply_keystream(data);
    
    // Explicitly zeroize the reconstructed key and temporary pre-key buffer
    key.zeroize();
    #[cfg(feature = "key_shielding")]
    combined.zeroize();
}