lib-q-aead 0.0.11

//! Secure memory handling and zeroization
//!
//! This module provides secure memory management functions including:
//! - Automatic zeroization of sensitive data
//! - Secure memory allocation and deallocation
//! - Memory barrier operations
//! - Secure memory copying and comparison

use core::ptr;

/// Secure memory zeroization
///
/// Securely zeros a memory region to prevent sensitive data from remaining
/// in memory after use. This function uses compiler barriers to prevent
/// optimization that might skip the zeroing.
///
/// # Arguments
/// * `data` - Memory region to zero
///
/// # Security
/// This function uses compiler barriers to ensure the zeroing operation
/// is not optimized away by the compiler.
///
/// # Safety
/// Overwriting `*data` with all-zero bytes must produce a valid value of
/// `T`. This is **not** true for most types: references, `Box`, `NonNull`,
/// `NonZero*`, most `enum`s, and any niche-optimized or
/// `#[repr(transparent)]` wrapper over such types have an all-zero bit
/// pattern that is not a legal value, so zeroing them is instant undefined
/// behaviour. The caller must ensure `T` has no validity invariant that
/// excludes the all-zero bit pattern (e.g. `T` is `u8`/`[u8; N]`, a plain
/// `#[repr(C)]`/`#[repr(Rust)]` aggregate of such types with no niches, or a
/// type that documents all-zero as a valid representation). Prefer
/// [`secure_zero_slice`] for byte buffers, which carries no such
/// requirement.
///
/// # Regression test (F5)
/// This function used to be a safe `pub fn`, which let safe code zero an
/// arbitrary `T` — including types (like `&str`) whose all-zero bit pattern
/// is not a legal value, which is instant undefined behaviour with no
/// `unsafe` in sight. Making it `unsafe fn` closes that hole at the type
/// system level: calling it without an `unsafe` block is now a compile
/// error, which this `compile_fail` doctest pins down so the fix cannot
/// silently regress back to a safe fn.
///
/// ```compile_fail
/// let mut r: &'static str = "k";
/// // No `unsafe` block: must fail to compile now that `secure_zero` is
/// // `unsafe fn`. Before the F5 fix this compiled (and was UB at runtime).
/// lib_q_aead::security::memory::secure_zero(&mut r);
/// ```
pub unsafe fn secure_zero<T>(data: &mut T) {
    let size = size_of_val(data);
    let ptr = data as *mut T as *mut u8;

    unsafe {
        ptr::write_bytes(ptr, 0, size);
    }
    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure zeroization of a slice
///
/// Securely zeros a slice of memory to prevent sensitive data from remaining
/// in memory after use.
///
/// # Arguments
/// * `data` - Slice to zero
///
/// # Security
/// This function uses compiler barriers to ensure the zeroing operation
/// is not optimized away by the compiler.
pub fn secure_zero_slice(data: &mut [u8]) {
    for byte in data.iter_mut() {
        *byte = 0;
    }

    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory copy
///
/// Securely copies memory from source to destination, ensuring that
/// sensitive data is properly handled.
///
/// # Arguments
/// * `dst` - Destination memory
/// * `src` - Source memory
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
///
/// # Safety
/// This byte-copies `size_of_val(src)` bytes over `*dst` without running
/// `T`'s destructor on the value it overwrites and without preventing `src`
/// from later being independently dropped. The caller must ensure:
/// - `T: Copy` (or the value in `*dst` is otherwise known not to own a
///   resource — e.g. heap memory, a file handle — that would be double-freed
///   once both `*dst` and `*src` are eventually dropped), and
/// - the raw byte copy is a valid way to duplicate `T` (true for `Copy`
///   types; not generally true for types with padding-sensitive invariants).
///
/// Violating either bullet is undefined behaviour or a double free.
pub unsafe fn secure_copy<T>(dst: &mut T, src: &T) {
    let size = size_of_val(src);
    let dst_ptr = dst as *mut T as *mut u8;
    let src_ptr = src as *const T as *const u8;

    unsafe {
        ptr::copy_nonoverlapping(src_ptr, dst_ptr, size);
    }
    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory copy for slices
///
/// Securely copies memory from source slice to destination slice.
///
/// # Arguments
/// * `dst` - Destination slice
/// * `src` - Source slice
///
/// # Panics
/// Panics if the slices have different lengths.
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
pub fn secure_copy_slice(dst: &mut [u8], src: &[u8]) {
    assert_eq!(dst.len(), src.len());

    for (d, s) in dst.iter_mut().zip(src.iter()) {
        *d = *s;
    }

    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory move
///
/// Securely moves memory from source to destination, zeroing the source
/// after the move to prevent data leakage.
///
/// # Arguments
/// * `dst` - Destination memory
/// * `src` - Source memory
///
/// # Security
/// This function securely moves data and zeroes the source to prevent
/// sensitive data from remaining in memory.
///
/// # Safety
/// Same requirements as [`secure_copy`] and [`secure_zero`]: `T` must be
/// safely byte-copyable (in practice, `T: Copy`) and must have a valid
/// all-zero representation.
pub unsafe fn secure_move<T>(dst: &mut T, src: &mut T) {
    // SAFETY: caller upholds the same obligations documented on this
    // function, which are exactly `secure_copy`'s and `secure_zero`'s.
    unsafe {
        secure_copy(dst, src);
        secure_zero(src);
    }
}

/// Secure memory move for slices
///
/// Securely moves memory from source slice to destination slice, zeroing
/// the source after the move.
///
/// # Arguments
/// * `dst` - Destination slice
/// * `src` - Source slice
///
/// # Panics
/// Panics if the slices have different lengths.
///
/// # Security
/// This function securely moves data and zeroes the source to prevent
/// sensitive data from remaining in memory.
pub fn secure_move_slice(dst: &mut [u8], src: &mut [u8]) {
    secure_copy_slice(dst, src);
    secure_zero_slice(src);
}

/// Secure memory comparison
///
/// Securely compares two memory regions in constant time to prevent
/// timing attacks.
///
/// # Arguments
/// * `a` - First memory region
/// * `b` - Second memory region
///
/// # Returns
/// * `true` if the regions are equal, `false` otherwise
///
/// # Security
/// This function performs the comparison in constant time to prevent
/// timing attacks.
///
/// # Safety
/// This reads every byte of `*a` and `*b`, including any padding bytes
/// inserted by the compiler between/after fields. Padding is not guaranteed
/// to be initialized (it is `undef` for a `#[repr(Rust)]`/`#[repr(C)]`
/// struct unless every byte was explicitly written), so reading it is
/// undefined behaviour, and even where it happens to be initialized its
/// value is unspecified, which can make this function return `false` for
/// two structurally-equal values. The caller must ensure `T` has no padding
/// (e.g. `T` is `u8`, `[u8; N]`, or another type whose every byte is a
/// defined, initialized field with no gaps). Prefer
/// [`secure_compare_slice`] for byte buffers, which carries no such
/// requirement.
pub unsafe fn secure_compare<T>(a: &T, b: &T) -> bool {
    let size = size_of_val(a);
    let a_ptr = a as *const T as *const u8;
    let b_ptr = b as *const T as *const u8;

    let mut result = 0u8;

    unsafe {
        for i in 0..size {
            result |= *a_ptr.add(i) ^ *b_ptr.add(i);
        }
    }

    result == 0
}

/// Secure memory comparison for slices
///
/// Securely compares two slices in constant time to prevent timing attacks.
///
/// # Arguments
/// * `a` - First slice
/// * `b` - Second slice
///
/// # Returns
/// * `true` if the slices are equal, `false` otherwise
///
/// # Security
/// This function performs the comparison in constant time to prevent
/// timing attacks.
pub fn secure_compare_slice(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }

    let mut result = 0u8;
    for (x, y) in a.iter().zip(b.iter()) {
        result |= x ^ y;
    }

    result == 0
}

/// Secure memory allocation with enhanced security features
///
/// Allocates memory securely with proper alignment, zeroing, and protection.
/// Implements secure memory allocation best practices including:
/// - Proper alignment for security-sensitive data
/// - Memory zeroing to prevent data leakage
/// - Compiler barriers to prevent optimization
/// - Memory protection where available
///
/// # Arguments
/// * `size` - Size of memory to allocate
/// * `alignment` - Memory alignment (defaults to cache line size for security)
///
/// # Returns
/// * `Some(ptr)` if allocation succeeds, `None` otherwise
///
/// # Security
/// This function allocates memory securely and zeros it to prevent
/// data leakage from previous allocations. Uses cache-line alignment
/// to prevent side-channel attacks through cache timing.
#[cfg(feature = "alloc")]
pub fn secure_alloc(size: usize) -> Option<*mut u8> {
    secure_alloc_aligned(size, 64) // Default to cache line alignment
}

/// Secure memory allocation with custom alignment
#[cfg(feature = "alloc")]
pub fn secure_alloc_aligned(size: usize, alignment: usize) -> Option<*mut u8> {
    use alloc::alloc::{
        Layout,
        alloc,
    };

    if size == 0 {
        return None;
    }

    // Ensure alignment is a power of 2
    let alignment = if alignment == 0 || !alignment.is_power_of_two() {
        64 // Default to cache line alignment
    } else {
        alignment
    };

    let layout = Layout::from_size_align(size, alignment).ok()?;
    let ptr = unsafe { alloc(layout) };

    if ptr.is_null() {
        return None;
    }

    // Zero the allocated memory with secure zeroing
    secure_zero_raw(ptr, size);

    // Memory barrier to ensure zeroing completes before use
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);

    Some(ptr)
}

/// Secure zeroing of raw memory
#[cfg(feature = "alloc")]
fn secure_zero_raw(ptr: *mut u8, size: usize) {
    if ptr.is_null() || size == 0 {
        return;
    }

    unsafe {
        // Use volatile writes to prevent compiler optimization
        let mut current = ptr;
        for _ in 0..size {
            ptr::write_volatile(current, 0);
            current = current.add(1);
        }
    }

    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory deallocation with enhanced security
///
/// Deallocates memory securely by zeroing it before deallocation.
/// Implements secure deallocation best practices including:
/// - Secure zeroing before deallocation
/// - Multiple passes of zeroing for sensitive data
/// - Compiler barriers to prevent optimization
/// - Proper layout reconstruction for deallocation
///
/// # Arguments
/// * `ptr` - Pointer to memory to deallocate
/// * `size` - Size of memory to deallocate
/// * `alignment` - Alignment that was used for the corresponding allocation
///
/// There used to be a zero-argument-alignment `secure_dealloc` convenience
/// wrapper here that silently deallocated with a hardcoded alignment of 64,
/// regardless of what alignment the memory was actually allocated with
/// (`secure_alloc_aligned` accepts an arbitrary caller-chosen alignment).
/// Deallocating with a `Layout` whose alignment differs from the one used at
/// allocation time is undefined behaviour, so that wrapper has been removed:
/// callers must state the alignment explicitly and it must match the value
/// passed to [`secure_alloc_aligned`] (or, for memory obtained via
/// [`secure_alloc`], the value `64`).
///
/// # Safety
/// This function is unsafe because it:
/// - Takes a raw pointer that must be valid for the given size
/// - The pointer must have been allocated with the same allocator
/// - The size must match the size used for allocation
/// - The alignment must match the alignment used for allocation
///
/// # Security
/// This function securely deallocates memory by zeroing it first
/// to prevent data leakage.
#[cfg(feature = "alloc")]
pub unsafe fn secure_dealloc(ptr: *mut u8, size: usize, alignment: usize) {
    unsafe { secure_dealloc_aligned(ptr, size, alignment) }
}

/// Secure memory deallocation with custom alignment
///
/// # Safety
///
/// - `ptr` must be a valid pointer returned by a previous allocation
/// - `size` must be the same size that was used for the original allocation
/// - `alignment` must be the same alignment that was used for the original allocation
/// - The memory must not be accessed after this function returns
#[cfg(feature = "alloc")]
pub unsafe fn secure_dealloc_aligned(ptr: *mut u8, size: usize, alignment: usize) {
    if ptr.is_null() || size == 0 {
        return;
    }

    // Ensure alignment is a power of 2
    let alignment = if alignment == 0 || !alignment.is_power_of_two() {
        64 // Default to cache line alignment
    } else {
        alignment
    };

    // Secure zeroing with multiple passes for sensitive data
    secure_zero_raw(ptr, size);

    // Additional pass with pattern to ensure zeroing
    unsafe {
        let mut current = ptr;
        for _ in 0..size {
            ptr::write_volatile(current, 0xFF);
            current = current.add(1);
        }
    }

    // Final zeroing pass
    secure_zero_raw(ptr, size);

    // Memory barrier to ensure all zeroing completes
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);

    use alloc::alloc::{
        Layout,
        dealloc,
    };
    let layout = Layout::from_size_align(size, alignment).unwrap();
    unsafe {
        dealloc(ptr, layout);
    }
}

/// Memory barrier
///
/// Inserts a memory barrier to prevent reordering of memory operations.
/// This is useful for ensuring that sensitive operations complete
/// before other operations begin.
///
/// # Security
/// This function prevents memory reordering that could lead to
/// timing attacks or other side-channel vulnerabilities.
pub fn memory_barrier() {
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory fill
///
/// Securely fills a memory region with a specific value.
///
/// # Arguments
/// * `data` - Memory region to fill
/// * `value` - Value to fill with
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
///
/// # Safety
/// Overwriting every byte of `*data` with `value` must produce a valid value
/// of `T` for all `value` in `0..=255` the caller might pass. This is **not**
/// true for most types: references, `Box`, `NonNull`, `NonZero*`, most
/// `enum`s, and any niche-optimized or `#[repr(transparent)]` wrapper over
/// such types have byte patterns (in particular the all-zero pattern, but
/// also most non-zero ones) that are not legal values, so filling them is
/// instant undefined behaviour. The caller must ensure `T` has no validity
/// invariant that excludes an arbitrary repeated-byte pattern (e.g. `T` is
/// `u8`/`[u8; N]`, or a plain `#[repr(C)]`/`#[repr(Rust)]` aggregate of such
/// types with no niches). Prefer [`secure_fill_slice`] for byte buffers,
/// which carries no such requirement.
///
/// # Regression test (F5-class)
/// `secure_fill` had the same unbounded-generic-safe-fn shape that was fixed
/// on `secure_zero`/`secure_copy`/`secure_compare`/`secure_move` (see F5 in
/// the API-soundness audit). Making it `unsafe fn` closes the hole at the
/// type system level: calling it without an `unsafe` block is now a compile
/// error, pinned down by this `compile_fail` doctest so the fix cannot
/// silently regress back to a safe fn.
///
/// ```compile_fail
/// let mut r: &'static str = "k";
/// // No `unsafe` block: must fail to compile now that `secure_fill` is
/// // `unsafe fn`. Before the fix this compiled (and was UB at runtime for
/// // most `value`s, since an arbitrary byte pattern is not a valid `&str`).
/// lib_q_aead::security::memory::secure_fill(&mut r, 0x41);
/// ```
pub unsafe fn secure_fill<T>(data: &mut T, value: u8) {
    let size = size_of_val(data);
    let ptr = data as *mut T as *mut u8;

    unsafe {
        ptr::write_bytes(ptr, value, size);
        // Compiler barrier to prevent optimization
        core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
    }
}

/// Secure memory fill for slices
///
/// Securely fills a slice with a specific value.
///
/// # Arguments
/// * `data` - Slice to fill
/// * `value` - Value to fill with
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
pub fn secure_fill_slice(data: &mut [u8], value: u8) {
    for byte in data.iter_mut() {
        *byte = value;
    }

    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Secure memory XOR
///
/// Securely XORs two memory regions and stores the result in the first.
///
/// # Arguments
/// * `a` - First memory region (modified in place)
/// * `b` - Second memory region
///
/// # Panics
/// Panics if the memory regions have different sizes.
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
///
/// # Safety
/// This XORs `*a` with `*b` byte-by-byte and writes the result back into
/// `*a`, including any padding bytes inserted by the compiler between/after
/// fields. Two independent problems make this unsound for a general `T`:
/// - Padding is not guaranteed to be initialized (`undef` for a
///   `#[repr(Rust)]`/`#[repr(C)]` struct unless every byte is an explicit
///   field), so reading it through `*b_ptr.add(i)` is undefined behaviour.
/// - The XOR result is an arbitrary derived bit pattern that need not be a
///   valid value of `T` (references, `NonZero*`, most `enum`s, niche-
///   optimized types, ...), so writing it back through `*a_ptr.add(i)` can
///   produce an instantly-invalid value.
///
/// The caller must ensure `T` has no padding and no validity invariant that
/// excludes an arbitrary bit pattern (e.g. `T` is `u8`, `[u8; N]`, or another
/// type whose every byte is a defined, initialized field with no gaps and no
/// niches). Prefer [`secure_xor_slice`] for byte buffers, which carries no
/// such requirement.
///
/// # Regression test (F5-class)
/// `secure_xor` had the same unbounded-generic-safe-fn shape that was fixed
/// on `secure_zero`/`secure_copy`/`secure_compare`/`secure_move` (see F5 in
/// the API-soundness audit). Making it `unsafe fn` closes the hole at the
/// type system level: calling it without an `unsafe` block is now a compile
/// error, pinned down by this `compile_fail` doctest so the fix cannot
/// silently regress back to a safe fn.
///
/// ```compile_fail
/// let mut r: &'static str = "k";
/// let s: &'static str = "k";
/// // No `unsafe` block: must fail to compile now that `secure_xor` is
/// // `unsafe fn`. Before the fix this compiled (and was UB at runtime).
/// lib_q_aead::security::memory::secure_xor(&mut r, &s);
/// ```
pub unsafe fn secure_xor<T>(a: &mut T, b: &T) {
    let size = size_of_val(a);
    assert_eq!(size, size_of_val(b));

    let a_ptr = a as *mut T as *mut u8;
    let b_ptr = b as *const T as *const u8;

    unsafe {
        for i in 0..size {
            *a_ptr.add(i) ^= *b_ptr.add(i);
        }
        // Compiler barrier to prevent optimization
        core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
    }
}

/// Secure memory XOR for slices
///
/// Securely XORs two slices and stores the result in the first.
///
/// # Arguments
/// * `a` - First slice (modified in place)
/// * `b` - Second slice
///
/// # Panics
/// Panics if the slices have different lengths.
///
/// # Security
/// This function uses secure memory operations to prevent data leakage.
pub fn secure_xor_slice(a: &mut [u8], b: &[u8]) {
    assert_eq!(a.len(), b.len());

    for (x, y) in a.iter_mut().zip(b.iter()) {
        *x ^= *y;
    }

    // Compiler barrier to prevent optimization
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_secure_zero() {
        let mut data = [1, 2, 3, 4, 5];
        // SAFETY: `[i32; 5]` has an all-zero valid representation.
        unsafe { secure_zero(&mut data) };
        assert_eq!(data, [0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_secure_zero_slice() {
        let mut data = [1, 2, 3, 4, 5];
        secure_zero_slice(&mut data);
        assert_eq!(data, [0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_secure_copy() {
        let src = [1, 2, 3, 4, 5];
        let mut dst = [0; 5];
        // SAFETY: `[i32; 5]` is `Copy` and has no padding-sensitive invariant.
        unsafe { secure_copy(&mut dst, &src) };
        assert_eq!(dst, src);
    }

    #[test]
    fn test_secure_copy_slice() {
        let src = [1, 2, 3, 4, 5];
        let mut dst = [0; 5];
        secure_copy_slice(&mut dst, &src);
        assert_eq!(dst, src);
    }

    #[test]
    fn test_secure_move() {
        let mut src = [1, 2, 3, 4, 5];
        let mut dst = [0; 5];
        // SAFETY: `[i32; 5]` is `Copy` and has an all-zero valid representation.
        unsafe { secure_move(&mut dst, &mut src) };
        assert_eq!(dst, [1, 2, 3, 4, 5]);
        assert_eq!(src, [0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_secure_move_slice() {
        let mut src = [1, 2, 3, 4, 5];
        let mut dst = [0; 5];
        secure_move_slice(&mut dst, &mut src);
        assert_eq!(dst, [1, 2, 3, 4, 5]);
        assert_eq!(src, [0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_secure_compare() {
        let a = [1, 2, 3, 4, 5];
        let b = [1, 2, 3, 4, 5];
        let c = [1, 2, 3, 4, 6];

        // SAFETY: `[i32; 5]` has no padding.
        unsafe {
            assert!(secure_compare(&a, &b));
            assert!(!secure_compare(&a, &c));
        }
    }

    #[test]
    fn test_secure_compare_slice() {
        let a = [1, 2, 3, 4, 5];
        let b = [1, 2, 3, 4, 5];
        let c = [1, 2, 3, 4, 6];

        assert!(secure_compare_slice(&a, &b));
        assert!(!secure_compare_slice(&a, &c));
    }

    #[test]
    fn test_secure_fill() {
        let mut data = [0u8; 5];
        // SAFETY: `[u8; 5]` accepts any repeated-byte pattern as a valid value.
        unsafe { secure_fill(&mut data, 42) };
        assert_eq!(data, [42, 42, 42, 42, 42]);
    }

    #[test]
    fn test_secure_fill_slice() {
        let mut data = [0; 5];
        secure_fill_slice(&mut data, 42);
        assert_eq!(data, [42, 42, 42, 42, 42]);
    }

    #[test]
    fn test_secure_xor() {
        let mut a = [0b1010, 0b1100, 0b1111];
        let b = [0b1100, 0b1010, 0b0000];
        // SAFETY: `[i32; 3]` has no padding and accepts an arbitrary bit pattern.
        unsafe { secure_xor(&mut a, &b) };
        assert_eq!(a, [0b0110, 0b0110, 0b1111]);
    }

    #[test]
    fn test_secure_xor_slice() {
        let mut a = [0b1010, 0b1100, 0b1111];
        let b = [0b1100, 0b1010, 0b0000];
        secure_xor_slice(&mut a, &b);
        assert_eq!(a, [0b0110, 0b0110, 0b1111]);
    }

    #[test]
    fn test_memory_barrier() {
        // This test just ensures the function doesn't panic
        memory_barrier();
    }
}