quasor 6.1.2

// src/quasor.rs

use crate::constants::*;
use crate::sponge::Sponge;
use argon2::{self, Argon2};
use subtle::ConstantTimeEq;
use zeroize::{Zeroize, ZeroizeOnDrop};

// Constants
const REKEY_INTERVAL: usize = 64 * 1024; // 64 KiB as per SPEC v6.1
const PARALLEL_THRESHOLD: usize = 1024 * 512;
const PROCESSING_CHUNK_SIZE: usize = 1024;

#[derive(Debug, PartialEq, Eq)]
pub enum Error {
    InvalidTag,
    Argon2Error(String),
}

// Securely handle the master key, zeroing it on drop
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct Quasor {
    master_key: [u8; 32],
}

impl Quasor {
    /// Creates a new Quasor instance, deriving a master key from a password
    /// using the memory-hard Argon2 function with secure, default parameters.
    pub fn new(password: &[u8], salt: &[u8]) -> Result<Self, Error> {
        let mut master_key = [0u8; 32];
        let argon2 = Argon2::default();
        argon2
            .hash_password_into(password, salt, &mut master_key)
            .map_err(|e| Error::Argon2Error(e.to_string()))?;
        Ok(Self { master_key })
    }

    /// Creates a new Quasor instance directly from a raw key.
    ///
    /// # Warning
    /// This function bypasses the slow, memory-hard Argon2 key derivation function.
    /// It should ONLY be used for testing purposes where performance is critical.
    /// **DO NOT USE IN PRODUCTION CODE.**
    pub fn from_raw_key(key: [u8; 32]) -> Self {
        Self { master_key: key }
    }

    /// Encrypts data using the derived master key.
    pub fn encrypt(&self, plaintext: &[u8], ad: &[u8]) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
        let nonce = self.derive_nonce(plaintext, ad);
        let mut enc_state = EncryptState::new(&self.master_key, &nonce, ad);
        let ciphertext = enc_state.process(plaintext);
        let tag = enc_state.finalize();
        (ciphertext, tag.to_vec(), nonce.to_vec())
    }

    /// Decrypts data using the derived master key.
    pub fn decrypt(
        &self,
        nonce: &[u8],
        ciphertext: &[u8],
        ad: &[u8],
        tag: &[u8],
    ) -> Result<Vec<u8>, Error> {
        let mut dec_state = DecryptState::new(&self.master_key, nonce, ad);
        let plaintext = dec_state.process(ciphertext);
        dec_state.verify_tag(tag)?;

        // Final verification of the SIV nonce
        let derived_nonce = self.derive_nonce(&plaintext, ad);
        if !bool::from(derived_nonce.as_slice().ct_eq(nonce)) {
            return Err(Error::InvalidTag);
        }

        Ok(plaintext)
    }

    /// Derives a nonce from the key and message content using a secure,
    /// length-prefixed serialization to prevent ambiguity.
    /// NOTE: Made public for benchmarking purposes.
    pub fn derive_nonce(&self, plaintext: &[u8], ad: &[u8]) -> [u8; 16] {
        let mut hasher = blake3::Hasher::new_keyed(&self.master_key);

        // Use unambiguous, length-prefixed serialization
        hasher.update(&(ad.len() as u64).to_le_bytes());
        hasher.update(ad);
        hasher.update(&(plaintext.len() as u64).to_le_bytes());

        if plaintext.len() > PARALLEL_THRESHOLD {
            hasher.update_rayon(plaintext);
        } else {
            hasher.update(plaintext);
        }

        let mut output = [0u8; 16];
        let mut output_reader = hasher.finalize_xof();
        output_reader.fill(&mut output);
        output
    }
}

pub struct EncryptState {
    sponge: Sponge,
    rekey_counter: u64,
}

impl EncryptState {
    pub fn new(key: &[u8], nonce: &[u8], ad: &[u8]) -> Self {
        let mut sponge = Sponge::new();

        // Domain separation for initialization
        sponge.absorb(DOMAIN_INIT);
        sponge.absorb(key);
        sponge.absorb(nonce);
        sponge.absorb(ad);

        Self {
            sponge,
            rekey_counter: 0,
        }
    }

    pub fn process(&mut self, plaintext: &[u8]) -> Vec<u8> {
        let mut ciphertext = Vec::with_capacity(plaintext.len());
        let mut bytes_processed_since_rekey = 0;

        // Domain separation for encryption
        self.sponge.absorb(DOMAIN_ENCRYPT);

        // Process plaintext in chunks, using separate sponges to avoid duplex issues
        for chunk in plaintext.chunks(PROCESSING_CHUNK_SIZE) {
            // Fork the sponge state for squeezing keystream
            let mut squeeze_sponge = self.sponge.fork();
            let keystream = squeeze_sponge.squeeze(chunk.len());

            // XOR with plaintext to produce ciphertext
            let mut cipher_chunk = Vec::with_capacity(chunk.len());
            for (p, k) in chunk.iter().zip(keystream.iter()) {
                cipher_chunk.push(p ^ k);
            }

            // Continue with the original sponge for absorbing plaintext
            self.sponge.absorb(chunk);
            ciphertext.extend_from_slice(&cipher_chunk);

            bytes_processed_since_rekey += chunk.len();

            if bytes_processed_since_rekey >= REKEY_INTERVAL {
                self.perform_rekey();
                bytes_processed_since_rekey = 0;
            }
        }

        ciphertext
    }

    fn perform_rekey(&mut self) {
        // Domain separation for rekeying
        self.sponge.absorb(DOMAIN_REKEY);
        self.sponge.absorb(&self.rekey_counter.to_le_bytes());

        // Fork sponge for squeezing ephemeral key to avoid duplex transition issues
        let mut squeeze_sponge = self.sponge.fork();
        let new_key = squeeze_sponge.squeeze(MASTER_KEY_SIZE);

        // Continue with original sponge for absorbing the ephemeral key
        self.sponge.absorb(&new_key);

        self.rekey_counter += 1;
    }

    pub fn finalize(&mut self) -> [u8; TAG_SIZE] {
        // Domain separation for authentication
        self.sponge.absorb(DOMAIN_AUTH);

        let tag_vec = self.sponge.squeeze(TAG_SIZE);
        let mut tag = [0u8; TAG_SIZE];
        tag.copy_from_slice(&tag_vec);
        tag
    }
}

pub struct DecryptState {
    sponge: Sponge,
    rekey_counter: u64,
}

impl DecryptState {
    pub fn new(key: &[u8], nonce: &[u8], ad: &[u8]) -> Self {
        let mut sponge = Sponge::new();

        // Domain separation for initialization
        sponge.absorb(DOMAIN_INIT);
        sponge.absorb(key);
        sponge.absorb(nonce);
        sponge.absorb(ad);

        Self {
            sponge,
            rekey_counter: 0,
        }
    }

    pub fn process(&mut self, ciphertext: &[u8]) -> Vec<u8> {
        let mut plaintext = Vec::with_capacity(ciphertext.len());
        let mut bytes_processed_since_rekey = 0;

        // Domain separation for encryption
        self.sponge.absorb(DOMAIN_ENCRYPT);

        // Process ciphertext in chunks, using separate sponges to avoid duplex issues
        for chunk in ciphertext.chunks(PROCESSING_CHUNK_SIZE) {
            // Fork the sponge state for squeezing keystream
            let mut squeeze_sponge = self.sponge.fork();
            let keystream = squeeze_sponge.squeeze(chunk.len());

            // XOR with ciphertext to recover plaintext
            let mut plain_chunk = Vec::with_capacity(chunk.len());
            for (c, k) in chunk.iter().zip(keystream.iter()) {
                plain_chunk.push(c ^ k);
            }

            // Continue with the original sponge for absorbing plaintext
            self.sponge.absorb(&plain_chunk);
            plaintext.extend_from_slice(&plain_chunk);

            bytes_processed_since_rekey += chunk.len();

            if bytes_processed_since_rekey >= REKEY_INTERVAL {
                self.perform_rekey();
                bytes_processed_since_rekey = 0;
            }
        }

        plaintext
    }

    fn perform_rekey(&mut self) {
        // Domain separation for rekeying
        self.sponge.absorb(DOMAIN_REKEY);
        self.sponge.absorb(&self.rekey_counter.to_le_bytes());

        // Fork sponge for squeezing ephemeral key to avoid duplex transition issues
        let mut squeeze_sponge = self.sponge.fork();
        let new_key = squeeze_sponge.squeeze(MASTER_KEY_SIZE);

        // Continue with original sponge for absorbing the ephemeral key
        self.sponge.absorb(&new_key);

        self.rekey_counter += 1;
    }

    pub fn verify_tag(&mut self, tag: &[u8]) -> Result<(), Error> {
        // Domain separation for authentication
        self.sponge.absorb(DOMAIN_AUTH);

        let expected_tag_vec = self.sponge.squeeze(TAG_SIZE);

        if bool::from(expected_tag_vec.as_slice().ct_eq(tag)) {
            Ok(())
        } else {
            Err(Error::InvalidTag)
        }
    }
}

// --- UNIT TESTS ---
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_forward_secrecy_rekeying_unit() {
        let key = [99; 32];
        let quasor = Quasor::from_raw_key(key);
        let ad = b"forward_secrecy_test";

        let block1_size = REKEY_INTERVAL - 100;
        let block2_size = 200;
        let total_size = block1_size + block2_size;
        let plaintext = vec![1u8; total_size];

        let nonce = quasor.derive_nonce(&plaintext, ad);
        let mut enc_state = EncryptState::new(&quasor.master_key, &nonce, ad);

        let c_block1 = enc_state.process(&plaintext[0..block1_size]);
        enc_state.process(&plaintext[block1_size..total_size]);

        let state_after_rekey = enc_state.sponge.clone();

        let mut attack_dec_state = DecryptState {
            sponge: state_after_rekey,
            rekey_counter: 0,
        };
        let recovered_p_block1 = attack_dec_state.process(&c_block1);

        assert_ne!(
            recovered_p_block1,
            &plaintext[0..block1_size],
            "Forward Secrecy FAILED: Post-rekey state was able to decrypt pre-rekey data!"
        );
    }
}