signer/src/secure_mem.rs

//! Secure memory buffer — mlock'd, zeroized on free.
//!
//! Port of `secure_mem.c`. Sensitive buffers (mnemonic, private keys)
//! live in mlock'd RAM and are zeroized on drop.

use std::alloc::{alloc, dealloc, Layout};
use std::sync::atomic::{AtomicBool, Ordering};
use zeroize::Zeroize;

/// Global flag permitting unlocked operation (when mlock fails).
static ALLOW_UNLOCKED: AtomicBool = AtomicBool::new(false);

/// Set the global flag permitting unlocked operation.
///
/// Call at startup when running in containers or environments with
/// limited `RLIMIT_MEMLOCK`. When set, `mlock` failures produce a
/// warning but do not abort allocation.
pub fn allow_unlocked() {
    ALLOW_UNLOCKED.store(true, Ordering::SeqCst);
}

/// Whether unlocked operation is currently permitted.
pub fn is_unlocked_allowed() -> bool {
    ALLOW_UNLOCKED.load(Ordering::SeqCst)
}

/// Secure memory buffer — mlock'd, zeroized on drop.
///
/// Holds sensitive material (mnemonic phrases, private keys). The
/// memory is locked with `mlock(2)` to prevent swap-out, and
/// zeroized with `explicit_bzero` semantics on free.
pub struct SecureBuf {
    ptr: *mut u8,
    size: usize,
    locked: bool,
}

unsafe impl Send for SecureBuf {}
unsafe impl Sync for SecureBuf {}

impl SecureBuf {
    /// Allocate a secure buffer of `size` bytes.
    ///
    /// Returns `MemoryFailed` if allocation or mlock fails (unless
    /// `allow_unlocked()` was called).
    pub fn alloc(size: usize) -> Result<Self, crate::NsignerError> {
        if size == 0 {
            return Err(crate::NsignerError::InvalidInput);
        }

        let layout = Layout::from_size_align(size, 1)
            .map_err(|_| crate::NsignerError::MemoryFailed)?;

        let ptr = unsafe { alloc(layout) };
        if ptr.is_null() {
            return Err(crate::NsignerError::MemoryFailed);
        }

        // Zero-initialize
        unsafe { std::ptr::write_bytes(ptr, 0, size) };

        // Attempt mlock
        let locked = unsafe { libc::mlock(ptr as *const libc::c_void, size) } == 0;
        if !locked && !is_unlocked_allowed() {
            // mlock failed and unlocked mode not permitted — fail hard
            unsafe { dealloc(ptr, layout) };
            return Err(crate::NsignerError::MemoryFailed);
        }

        Ok(SecureBuf { ptr, size, locked })
    }

    /// Usable size in bytes.
    pub fn size(&self) -> usize {
        self.size
    }

    /// Whether mlock succeeded.
    pub fn is_locked(&self) -> bool {
        self.locked
    }

    /// Read access to the buffer contents.
    pub fn as_slice(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.ptr, self.size) }
    }

    /// Write access to the buffer contents.
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
    }

    /// Copy data into the buffer (truncates to buffer size).
    pub fn copy_from(&mut self, src: &[u8]) {
        let len = src.len().min(self.size);
        self.as_mut_slice()[..len].copy_from_slice(&src[..len]);
    }

    /// Copy data into the buffer from a slice (alias for compatibility).
    pub fn copy_from_slice(&mut self, src: &[u8]) {
        self.copy_from(src);
    }

    /// Zeroize the buffer contents in place.
    pub fn clear(&mut self) {
        self.as_mut_slice().zeroize();
    }

    /// Resize the buffer to `new_size`, preserving the prefix that fits.
    ///
    /// If `new_size` is 0, returns `InvalidInput`. If allocation of the
    /// new buffer fails, the original buffer is left intact and an error
    /// is returned.
    pub fn resize(&mut self, new_size: usize) -> Result<(), crate::NsignerError> {
        if new_size == 0 {
            return Err(crate::NsignerError::InvalidInput);
        }
        if new_size == self.size {
            return Ok(());
        }

        let mut new_buf = SecureBuf::alloc(new_size)?;
        let copy_len = self.size.min(new_size);
        new_buf.as_mut_slice()[..copy_len].copy_from_slice(&self.as_slice()[..copy_len]);

        // Swap internals
        let old_ptr = self.ptr;
        let old_size = self.size;
        let old_locked = self.locked;

        self.ptr = new_buf.ptr;
        self.size = new_buf.size;
        self.locked = new_buf.locked;

        // Prevent new_buf's Drop from running on the moved-out pointer
        new_buf.ptr = std::ptr::null_mut();
        new_buf.size = 0;
        new_buf.locked = false;

        // Free old buffer
        if !old_ptr.is_null() && old_size > 0 {
            unsafe { std::ptr::write_bytes(old_ptr, 0, old_size) };
            if old_locked {
                unsafe { libc::munlock(old_ptr as *const libc::c_void, old_size) };
            }
            let layout = Layout::from_size_align(old_size, 1).unwrap();
            unsafe { dealloc(old_ptr, layout) };
        }

        Ok(())
    }
}

impl Drop for SecureBuf {
    fn drop(&mut self) {
        if !self.ptr.is_null() && self.size > 0 {
            // Zeroize
            unsafe { std::ptr::write_bytes(self.ptr, 0, self.size) };

            // munlock if locked
            if self.locked {
                unsafe { libc::munlock(self.ptr as *const libc::c_void, self.size) };
            }

            // Dealloc
            let layout = Layout::from_size_align(self.size, 1).unwrap();
            unsafe { dealloc(self.ptr, layout) };
        }
        self.ptr = std::ptr::null_mut();
        self.size = 0;
    }
}

/// Zeroize `len` bytes at `ptr` in a way the compiler cannot optimize away.
pub fn secure_memzero(buf: &mut [u8]) {
    buf.zeroize();
}

/// Constant-time comparison of two byte slices.
///
/// Returns `true` if the slices are equal. The comparison runs in time
/// proportional to the shorter slice's length (callers should ensure
/// equal lengths for full constant-time properties).
pub fn secure_compare(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}

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

    #[test]
    fn test_alloc_free() {
        let buf = SecureBuf::alloc(32).unwrap();
        assert_eq!(buf.size(), 32);
        // Drop runs zeroize + munlock + dealloc
    }

    #[test]
    fn test_copy_from() {
        let mut buf = SecureBuf::alloc(8).unwrap();
        let data = [1u8, 2, 3, 4, 5, 6, 7, 8];
        buf.copy_from(&data);
        assert_eq!(buf.as_slice(), &data[..]);
    }

    #[test]
    fn test_copy_from_truncates() {
        let mut buf = SecureBuf::alloc(4).unwrap();
        let data = [1u8, 2, 3, 4, 5, 6, 7, 8];
        buf.copy_from(&data);
        assert_eq!(buf.as_slice(), &[1, 2, 3, 4]);
    }

    #[test]
    fn test_zero_size_rejected() {
        assert!(SecureBuf::alloc(0).is_err());
    }

    #[test]
    fn test_clear_zeroizes() {
        let mut buf = SecureBuf::alloc(16).unwrap();
        buf.copy_from(&[0xFF; 16]);
        buf.clear();
        assert!(buf.as_slice().iter().all(|&b| b == 0));
    }

    #[test]
    fn test_resize_grows_preserving_prefix() {
        let mut buf = SecureBuf::alloc(4).unwrap();
        buf.copy_from(&[10, 20, 30, 40]);
        buf.resize(8).unwrap();
        assert_eq!(buf.size(), 8);
        assert_eq!(&buf.as_slice()[..4], &[10, 20, 30, 40]);
        assert_eq!(&buf.as_slice()[4..], &[0, 0, 0, 0]);
    }

    #[test]
    fn test_resize_shrinks_preserving_prefix() {
        let mut buf = SecureBuf::alloc(8).unwrap();
        buf.copy_from(&[10, 20, 30, 40, 50, 60, 70, 80]);
        buf.resize(3).unwrap();
        assert_eq!(buf.size(), 3);
        assert_eq!(buf.as_slice(), &[10, 20, 30]);
    }

    #[test]
    fn test_resize_zero_rejected() {
        let mut buf = SecureBuf::alloc(4).unwrap();
        assert!(buf.resize(0).is_err());
    }

    #[test]
    fn test_secure_memzero() {
        let mut data = [0xABu8; 32];
        secure_memzero(&mut data);
        assert!(data.iter().all(|&b| b == 0));
    }

    #[test]
    fn test_secure_compare_equal() {
        assert!(secure_compare(&[1, 2, 3], &[1, 2, 3]));
    }

    #[test]
    fn test_secure_compare_unequal() {
        assert!(!secure_compare(&[1, 2, 3], &[1, 2, 4]));
    }

    #[test]
    fn test_secure_compare_different_lengths() {
        assert!(!secure_compare(&[1, 2, 3], &[1, 2]));
    }

    #[test]
    fn test_allow_unlocked_flag() {
        let prev = is_unlocked_allowed();
        allow_unlocked();
        assert!(is_unlocked_allowed());
        // Restore for other tests
        ALLOW_UNLOCKED.store(prev, Ordering::SeqCst);
    }
}