lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! # Cryptographic Operations for LCPFS
//!
//! This module provides authenticated encryption and cryptographic hashing
//! for filesystem-level encryption in LCPFS.
//!
//! ## Algorithms
//!
//! - **ChaCha20-Poly1305**: AEAD cipher for block encryption/decryption
//! - **SHAKE256**: Extendable-output function (XOF) for variable-length hashing
//!
//! ## Security Guarantees
//!
//! - **Authenticated Encryption**: ChaCha20-Poly1305 provides both confidentiality
//!   and integrity. Tampering is detected during decryption.
//! - **Nonce Management**: 12-byte nonces are prepended to ciphertext for stateless
//!   decryption.
//! - **Key Zeroization**: Keys are automatically zeroed when dropped using the
//!   `zeroize` crate.
//!
//! ## Usage
//!
//! ```rust,ignore
//! use lcpfs::lcpfs_crypto::{encrypt_block, decrypt_block, shake256};
//!
//! // Encrypt data
//! let key = [0u8; 32]; // Use proper key derivation in production
//! let plaintext = b"sensitive filesystem data";
//! let ciphertext = encrypt_block(plaintext, &key).unwrap();
//!
//! // Decrypt data
//! let decrypted = decrypt_block(&ciphertext, &key).unwrap();
//! assert_eq!(decrypted, plaintext);
//!
//! // Variable-length hash
//! let hash = shake256(b"data to hash", 64).unwrap();
//! assert_eq!(hash.len(), 64);
//! ```

use alloc::vec::Vec;
use chacha20poly1305_nostd::ChaCha20Poly1305;
use sha3::Shake256;
use sha3::digest::{ExtendableOutput, Update, XofReader};
use zeroize::Zeroize;

use crate::FsError;
use crate::arch;

// ═══════════════════════════════════════════════════════════════════════════════
// CONSTANTS
// ═══════════════════════════════════════════════════════════════════════════════

/// Size of ChaCha20-Poly1305 key in bytes (256 bits).
pub const KEY_SIZE: usize = 32;

/// Size of ChaCha20-Poly1305 nonce in bytes (96 bits).
pub const NONCE_SIZE: usize = 12;

/// Size of Poly1305 authentication tag in bytes (128 bits).
pub const TAG_SIZE: usize = 16;

// ═══════════════════════════════════════════════════════════════════════════════
// ENCRYPTION
// ═══════════════════════════════════════════════════════════════════════════════

/// Encrypt data using ChaCha20-Poly1305 AEAD.
///
/// This function provides authenticated encryption: the ciphertext cannot be
/// modified without detection. A random 12-byte nonce is generated and prepended
/// to the output.
///
/// # Arguments
///
/// * `data` - Plaintext data to encrypt
/// * `key` - 32-byte encryption key
///
/// # Returns
///
/// On success, returns `Vec<u8>` containing:
/// - Bytes 0..12: Random nonce
/// - Bytes 12..N: Ciphertext with authentication tag (16 bytes longer than plaintext)
///
/// Total output length: `NONCE_SIZE + data.len() + TAG_SIZE` (28 bytes overhead)
///
/// # Errors
///
/// - `FsError::InvalidArgument` - Key is not 32 bytes
/// - `FsError::EncryptionFailed` - Encryption failed (should not happen with valid inputs)
///
/// # Security
///
/// - Uses hardware entropy (RDRAND on x86_64, RNDR on AArch64) for nonce generation
/// - Falls back to platform entropy source if hardware RNG unavailable
/// - Nonce uniqueness is critical: never reuse a nonce with the same key
///
/// # Example
///
/// ```rust,ignore
/// use lcpfs::lcpfs_crypto::encrypt_block;
///
/// let key = [0x42u8; 32];
/// let plaintext = b"Hello, LCPFS!";
///
/// let ciphertext = encrypt_block(plaintext, &key).unwrap();
///
/// // Ciphertext is longer than plaintext (nonce + tag overhead)
/// assert!(ciphertext.len() > plaintext.len());
///
/// // First 12 bytes are the nonce
/// let nonce = &ciphertext[..12];
/// ```
pub fn encrypt_block(data: &[u8], key: &[u8]) -> Result<Vec<u8>, FsError> {
    // Validate key length
    if key.len() != KEY_SIZE {
        return Err(FsError::InvalidArgument {
            reason: "encryption key must be 32 bytes",
        });
    }

    // Create cipher instance
    let cipher = ChaCha20Poly1305::new(key).map_err(|_| FsError::InvalidArgument {
        reason: "invalid key for ChaCha20Poly1305",
    })?;

    // Generate random nonce using hardware entropy
    let mut nonce_bytes = [0u8; NONCE_SIZE];
    fill_hardware_entropy(&mut nonce_bytes)?;

    // Encrypt with authenticated encryption (no additional associated data)
    let ciphertext = cipher
        .encrypt(&nonce_bytes, data, None)
        .map_err(|_| FsError::EncryptionFailed)?;

    // Prepend nonce to ciphertext for stateless decryption
    let mut result = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&ciphertext);

    Ok(result)
}

// ═══════════════════════════════════════════════════════════════════════════════
// DECRYPTION
// ═══════════════════════════════════════════════════════════════════════════════

/// Decrypt data using ChaCha20-Poly1305 AEAD.
///
/// This function verifies the authentication tag before returning plaintext.
/// If the ciphertext has been tampered with, decryption will fail.
///
/// # Arguments
///
/// * `data` - Ciphertext from [`encrypt_block`], including prepended nonce
/// * `key` - 32-byte decryption key (must match encryption key)
///
/// # Returns
///
/// On success, returns `Vec<u8>` containing the original plaintext.
///
/// # Errors
///
/// - `FsError::InvalidArgument` - Key is not 32 bytes or data is too short
/// - `FsError::DecryptionFailed` - Authentication failed (wrong key or tampering)
///
/// # Security
///
/// - Constant-time authentication tag verification (via Poly1305)
/// - Fails fast if ciphertext is too short (prevents timing attacks)
///
/// # Example
///
/// ```rust,ignore
/// use lcpfs::lcpfs_crypto::{encrypt_block, decrypt_block};
///
/// let key = [0x42u8; 32];
/// let plaintext = b"Hello, LCPFS!";
///
/// let ciphertext = encrypt_block(plaintext, &key).unwrap();
/// let decrypted = decrypt_block(&ciphertext, &key).unwrap();
///
/// assert_eq!(decrypted, plaintext);
/// ```
///
/// # Tamper Detection
///
/// ```rust,ignore
/// let mut tampered = ciphertext.clone();
/// tampered[20] ^= 0xFF; // Flip a bit
///
/// // Decryption fails due to authentication tag mismatch
/// assert!(decrypt_block(&tampered, &key).is_err());
/// ```
pub fn decrypt_block(data: &[u8], key: &[u8]) -> Result<Vec<u8>, FsError> {
    // Validate key length
    if key.len() != KEY_SIZE {
        return Err(FsError::InvalidArgument {
            reason: "decryption key must be 32 bytes",
        });
    }

    // Minimum ciphertext length: nonce (12) + tag (16) = 28 bytes
    if data.len() < NONCE_SIZE + TAG_SIZE {
        return Err(FsError::InvalidArgument {
            reason: "ciphertext too short (minimum 28 bytes)",
        });
    }

    // Extract nonce from prepended bytes
    let nonce = &data[..NONCE_SIZE];
    let ciphertext = &data[NONCE_SIZE..];

    // Create cipher instance
    let cipher = ChaCha20Poly1305::new(key).map_err(|_| FsError::InvalidArgument {
        reason: "invalid key for ChaCha20Poly1305",
    })?;

    // Decrypt and verify authentication tag (no additional associated data)
    let plaintext = cipher
        .decrypt(nonce, ciphertext, None)
        .map_err(|_| FsError::DecryptionFailed)?;

    Ok(plaintext)
}

// ═══════════════════════════════════════════════════════════════════════════════
// HASHING
// ═══════════════════════════════════════════════════════════════════════════════

/// Compute SHAKE256 hash with variable output length.
///
/// SHAKE256 is an extendable-output function (XOF) based on Keccak. Unlike
/// fixed-output hash functions (SHA-256, BLAKE3), it can produce any output
/// length.
///
/// # Arguments
///
/// * `data` - Input data to hash
/// * `output_len` - Desired output length in bytes
///
/// # Returns
///
/// On success, returns `Vec<u8>` of length `output_len` containing the hash.
///
/// # Errors
///
/// - `FsError::InvalidArgument` - Output length is 0
///
/// # Security Properties
///
/// - **Collision Resistance**: ~min(output_len*4, 256) bits
/// - **Preimage Resistance**: ~min(output_len*8, 256) bits
/// - **Deterministic**: Same input always produces same output
///
/// # Use Cases
///
/// - Key derivation (when combined with salt)
/// - Variable-length fingerprints
/// - Mask generation functions
///
/// # Example
///
/// ```rust,ignore
/// use lcpfs::lcpfs_crypto::shake256;
///
/// // 32-byte hash
/// let hash32 = shake256(b"data", 32).unwrap();
/// assert_eq!(hash32.len(), 32);
///
/// // 64-byte hash
/// let hash64 = shake256(b"data", 64).unwrap();
/// assert_eq!(hash64.len(), 64);
///
/// // Deterministic: same input = same output
/// let hash32_again = shake256(b"data", 32).unwrap();
/// assert_eq!(hash32, hash32_again);
/// ```
pub fn shake256(data: &[u8], output_len: usize) -> Result<Vec<u8>, FsError> {
    if output_len == 0 {
        return Err(FsError::InvalidArgument {
            reason: "output length must be greater than 0",
        });
    }

    // Create SHAKE256 hasher and absorb input
    let mut hasher = Shake256::default();
    hasher.update(data);

    // Squeeze output
    let mut output = alloc::vec![0u8; output_len];
    let mut reader = hasher.finalize_xof();
    reader.read(&mut output);

    Ok(output)
}

// ═══════════════════════════════════════════════════════════════════════════════
// HELPER FUNCTIONS
// ═══════════════════════════════════════════════════════════════════════════════

/// Fill buffer with hardware entropy.
///
/// Attempts to use platform hardware RNG first, falling back to
/// the lcpfs_random module if hardware RNG is unavailable.
fn fill_hardware_entropy(buf: &mut [u8]) -> Result<(), FsError> {
    // Try hardware entropy first
    if arch::has_rdrand() && arch::fill_hardware_entropy(buf).is_ok() {
        return Ok(());
    }

    // Fallback to lcpfs_random module
    crate::crypto::random::fill_random(buf).map_err(|_| FsError::EncryptionFailed)
}

// ═══════════════════════════════════════════════════════════════════════════════
// ZEROIZING KEY WRAPPER
// ═══════════════════════════════════════════════════════════════════════════════

/// A 32-byte key that is automatically zeroed when dropped.
///
/// This wrapper ensures that cryptographic keys are securely erased from
/// memory when no longer needed, preventing key material from lingering
/// in memory.
///
/// # Example
///
/// ```rust,ignore
/// use lcpfs::lcpfs_crypto::ZeroizingKey;
///
/// let key = ZeroizingKey::new([0x42u8; 32]);
/// // Use key.as_bytes() for encryption...
/// // Key is automatically zeroed when `key` goes out of scope
/// ```
#[derive(Zeroize)]
#[zeroize(drop)]
pub struct ZeroizingKey {
    bytes: [u8; KEY_SIZE],
}

impl ZeroizingKey {
    /// Create a new zeroizing key from raw bytes.
    pub fn new(bytes: [u8; KEY_SIZE]) -> Self {
        Self { bytes }
    }

    /// Get the key bytes as a slice.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

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

    // ───────────────────────────────────────────────────────────────────────────
    // ENCRYPTION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_encrypt_actually_encrypts() {
        let plaintext = b"secret data that must be encrypted";
        let key = [0u8; 32];

        let ciphertext = encrypt_block(plaintext, &key).unwrap();

        // Ciphertext must NOT equal plaintext (data portion starts after nonce)
        assert_ne!(
            &ciphertext[NONCE_SIZE..],
            plaintext.as_slice(),
            "Encryption must change data - not an identity function!"
        );

        // Must be longer: nonce (12) + plaintext + tag (16)
        assert_eq!(
            ciphertext.len(),
            NONCE_SIZE + plaintext.len() + TAG_SIZE,
            "Ciphertext should be plaintext + 28 bytes overhead"
        );
    }

    #[test]
    fn test_encrypt_produces_different_ciphertext_each_time() {
        let plaintext = b"same plaintext";
        let key = [0x42u8; 32];

        let ciphertext1 = encrypt_block(plaintext, &key).unwrap();
        let ciphertext2 = encrypt_block(plaintext, &key).unwrap();

        // Different nonces should produce different ciphertext
        assert_ne!(
            ciphertext1, ciphertext2,
            "Same plaintext should produce different ciphertext due to random nonce"
        );
    }

    #[test]
    fn test_encrypt_invalid_key_length() {
        let plaintext = b"test";

        // Too short
        assert!(matches!(
            encrypt_block(plaintext, &[0u8; 16]),
            Err(FsError::InvalidArgument { .. })
        ));

        // Too long
        assert!(matches!(
            encrypt_block(plaintext, &[0u8; 64]),
            Err(FsError::InvalidArgument { .. })
        ));
    }

    // ───────────────────────────────────────────────────────────────────────────
    // DECRYPTION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_decrypt_reverses_encrypt() {
        let plaintext = b"secret data to round-trip";
        let key = [0x42u8; 32];

        let ciphertext = encrypt_block(plaintext, &key).unwrap();
        let decrypted = decrypt_block(&ciphertext, &key).unwrap();

        assert_eq!(
            decrypted.as_slice(),
            plaintext.as_slice(),
            "Decryption must recover original plaintext"
        );
    }

    #[test]
    fn test_decrypt_with_wrong_key_fails() {
        let plaintext = b"secret";
        let key1 = [0x42u8; 32];
        let key2 = [0x43u8; 32];

        let ciphertext = encrypt_block(plaintext, &key1).unwrap();

        // Wrong key should fail authentication
        assert!(
            matches!(
                decrypt_block(&ciphertext, &key2),
                Err(FsError::DecryptionFailed)
            ),
            "Decryption with wrong key must fail"
        );
    }

    #[test]
    fn test_decrypt_tampered_ciphertext_fails() {
        let plaintext = b"secret";
        let key = [0u8; 32];

        let mut ciphertext = encrypt_block(plaintext, &key).unwrap();

        // Tamper with ciphertext (flip a bit in the encrypted data, not nonce)
        ciphertext[NONCE_SIZE + 1] ^= 0xFF;

        // Tampered ciphertext should fail authentication
        assert!(
            matches!(
                decrypt_block(&ciphertext, &key),
                Err(FsError::DecryptionFailed)
            ),
            "Tampered ciphertext must fail authentication"
        );
    }

    #[test]
    fn test_decrypt_invalid_key_length() {
        // Fake ciphertext (minimum valid length)
        let fake_ciphertext = [0u8; 28];

        // Too short key
        assert!(matches!(
            decrypt_block(&fake_ciphertext, &[0u8; 16]),
            Err(FsError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn test_decrypt_ciphertext_too_short() {
        let key = [0u8; 32];

        // Less than minimum (nonce + tag = 28 bytes)
        assert!(matches!(
            decrypt_block(&[0u8; 27], &key),
            Err(FsError::InvalidArgument { .. })
        ));

        // Empty
        assert!(matches!(
            decrypt_block(&[], &key),
            Err(FsError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn test_encrypt_decrypt_empty_plaintext() {
        let plaintext = b"";
        let key = [0x42u8; 32];

        let ciphertext = encrypt_block(plaintext, &key).unwrap();

        // Should be exactly nonce + tag (no data)
        assert_eq!(ciphertext.len(), NONCE_SIZE + TAG_SIZE);

        let decrypted = decrypt_block(&ciphertext, &key).unwrap();
        assert!(decrypted.is_empty());
    }

    #[test]
    fn test_encrypt_decrypt_large_data() {
        let plaintext = alloc::vec![0xAB_u8; 1024 * 1024]; // 1 MiB
        let key = [0x42u8; 32];

        let ciphertext = encrypt_block(&plaintext, &key).unwrap();
        let decrypted = decrypt_block(&ciphertext, &key).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    // ───────────────────────────────────────────────────────────────────────────
    // SHAKE256 TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_shake256_not_zeros() {
        let data = b"test data for hashing";
        let hash = shake256(data, 64).unwrap();

        // Must NOT be all zeros
        assert_ne!(
            hash,
            alloc::vec![0u8; 64],
            "SHAKE256 must not return zeros - it must actually hash!"
        );
    }

    #[test]
    fn test_shake256_deterministic() {
        let data = b"deterministic test";

        let hash1 = shake256(data, 64).unwrap();
        let hash2 = shake256(data, 64).unwrap();

        assert_eq!(hash1, hash2, "SHAKE256 must be deterministic");
    }

    #[test]
    fn test_shake256_different_inputs_different_outputs() {
        let hash1 = shake256(b"input 1", 32).unwrap();
        let hash2 = shake256(b"input 2", 32).unwrap();

        assert_ne!(
            hash1, hash2,
            "Different inputs must produce different hashes"
        );
    }

    #[test]
    fn test_shake256_variable_output_length() {
        let data = b"test";

        let hash16 = shake256(data, 16).unwrap();
        let hash32 = shake256(data, 32).unwrap();
        let hash64 = shake256(data, 64).unwrap();
        let hash128 = shake256(data, 128).unwrap();

        assert_eq!(hash16.len(), 16);
        assert_eq!(hash32.len(), 32);
        assert_eq!(hash64.len(), 64);
        assert_eq!(hash128.len(), 128);

        // Shorter outputs should be prefixes of longer outputs (XOF property)
        assert_eq!(&hash32[..16], &hash16[..]);
        assert_eq!(&hash64[..32], &hash32[..]);
        assert_eq!(&hash128[..64], &hash64[..]);
    }

    #[test]
    fn test_shake256_zero_length_fails() {
        assert!(matches!(
            shake256(b"test", 0),
            Err(FsError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn test_shake256_empty_input() {
        // Empty input should still produce valid hash
        let hash = shake256(b"", 32).unwrap();
        assert_eq!(hash.len(), 32);
        assert_ne!(hash, alloc::vec![0u8; 32]);
    }

    #[test]
    fn test_shake256_known_test_vector() {
        // NIST test vector for SHAKE256("")
        // From https://csrc.nist.gov/Projects/cryptographic-algorithm-validation-program
        let hash = shake256(b"", 32).unwrap();

        // Expected output for SHAKE256("") with 32 bytes output
        let expected: [u8; 32] = [
            0x46, 0xb9, 0xdd, 0x2b, 0x0b, 0xa8, 0x8d, 0x13, 0x23, 0x3b, 0x3f, 0xeb, 0x74, 0x3e,
            0xeb, 0x24, 0x3f, 0xcd, 0x52, 0xea, 0x62, 0xb8, 0x1b, 0x82, 0xb5, 0x0c, 0x27, 0x64,
            0x6e, 0xd5, 0x76, 0x2f,
        ];

        assert_eq!(
            hash.as_slice(),
            &expected[..],
            "SHAKE256 must match NIST test vector"
        );
    }

    // ───────────────────────────────────────────────────────────────────────────
    // ZEROIZING KEY TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_zeroizing_key_creation() {
        let key_bytes = [0x42u8; 32];
        let key = ZeroizingKey::new(key_bytes);

        assert_eq!(key.as_bytes(), &key_bytes);
    }

    // ───────────────────────────────────────────────────────────────────────────
    // INTEGRATION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_encrypt_then_hash_workflow() {
        let plaintext = b"data to encrypt and hash";
        let key = [0x42u8; 32];

        // Encrypt
        let ciphertext = encrypt_block(plaintext, &key).unwrap();

        // Hash the ciphertext (for integrity tracking, storage dedup, etc.)
        let hash = shake256(&ciphertext, 32).unwrap();

        // Verify hash is not trivial
        assert_ne!(hash, alloc::vec![0u8; 32]);
        assert_eq!(hash.len(), 32);

        // Decrypt should still work
        let decrypted = decrypt_block(&ciphertext, &key).unwrap();
        assert_eq!(decrypted.as_slice(), plaintext.as_slice());
    }
}