quasor 6.1.6

// src/quasor.rs

use crate::builder::QuasorBuilder;
use crate::constants::*;
use crate::sponge::Sponge;
use subtle::ConstantTimeEq;
use thiserror::Error;
use zeroize::{Zeroize, ZeroizeOnDrop};

// Constants
// This must match the value in the tests, which is now 64 * 1024.
const REKEY_INTERVAL: usize = 64 * 1024;
const PARALLEL_THRESHOLD: usize = 1024 * 512;
const PROCESSING_CHUNK_SIZE: usize = 1024;

#[derive(Error, Debug, PartialEq, Eq)]
pub enum Error {
    #[error("Invalid authentication tag")]
    InvalidTag,

    #[error("Invalid state: The operation was called in an incorrect order")]
    InvalidState,

    #[error("Argon2 error: {0}")]
    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 using the default Argon2 parameters.
    pub fn new(password: &[u8], salt: &[u8]) -> Result<Self, Error> {
        QuasorBuilder::new(password, salt).build()
    }

    /// Internal constructor for the builder.
    pub(crate) fn new_with_key(master_key: [u8; 32]) -> Self {
        Self { master_key }
    }

    /// Returns a builder for creating a Quasor instance with custom Argon2 parameters.
    pub fn builder<'a>(password: &'a [u8], salt: &'a [u8]) -> QuasorBuilder<'a> {
        QuasorBuilder::new(password, salt)
    }

    /// Creates a new Quasor instance directly from a raw key.
    pub fn from_raw_key(key: [u8; 32]) -> Self {
        Self { master_key: key }
    }

    /// Returns the master key.
    pub fn master_key(&self) -> &[u8; 32] {
        &self.master_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)?;

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

        Ok(plaintext)
    }

    /// Creates a streaming encryption instance.
    pub fn create_encrypt_stream<'a>(
        &'a self,
        plaintext: &'a [u8],
        ad: &'a [u8],
    ) -> QuasorEncryptStream<'a> {
        QuasorEncryptStream::new(&self.master_key, plaintext, ad)
    }

    /// Creates a streaming decryption instance.
    pub fn create_decrypt_stream<'a>(
        &'a self,
        nonce: &'a [u8],
        ad: &'a [u8],
        tag: &'a [u8],
        plaintext_len: u64,
    ) -> QuasorDecryptStream<'a> {
        QuasorDecryptStream::new(&self.master_key, nonce, ad, tag, plaintext_len)
    }

    /// Derives a nonce from the key and message content.
    pub fn derive_nonce(&self, plaintext: &[u8], ad: &[u8]) -> [u8; 16] {
        let mut hasher = blake3::Hasher::new_keyed(&self.master_key);
        // This part is correct, it prefixes AD with its length
        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
    }
}

// --- STREAMING STRUCTS ---

pub struct QuasorEncryptStream<'a> {
    master_key: &'a [u8; 32],
    plaintext: &'a [u8],
    ad: &'a [u8],
    enc_state: Option<EncryptState>,
    nonce: Option<[u8; 16]>,
    tag: Option<[u8; TAG_SIZE]>,
}

impl<'a> QuasorEncryptStream<'a> {
    fn new(master_key: &'a [u8; 32], plaintext: &'a [u8], ad: &'a [u8]) -> Self {
        Self {
            master_key,
            plaintext,
            ad,
            enc_state: None,
            nonce: None,
            tag: None,
        }
    }

    pub fn init(&mut self) {
        let mut hasher = blake3::Hasher::new_keyed(self.master_key);
        hasher.update(&(self.ad.len() as u64).to_le_bytes());
        hasher.update(self.ad);
        hasher.update(&(self.plaintext.len() as u64).to_le_bytes());
        if self.plaintext.len() > PARALLEL_THRESHOLD {
            hasher.update_rayon(self.plaintext);
        } else {
            hasher.update(self.plaintext);
        }
        let mut output = [0u8; 16];
        let mut output_reader = hasher.finalize_xof();
        output_reader.fill(&mut output);
        self.nonce = Some(output);
        self.enc_state = Some(EncryptState::new(self.master_key, &output, self.ad));
    }

    pub fn get_nonce(&self) -> Result<&[u8; 16], Error> {
        self.nonce.as_ref().ok_or(Error::InvalidState)
    }

    pub fn process(&mut self, plaintext_chunk: &[u8]) -> Result<Vec<u8>, Error> {
        let state = self.enc_state.as_mut().ok_or(Error::InvalidState)?;
        Ok(state.process(plaintext_chunk))
    }

    pub fn finalize(&mut self) -> Result<[u8; TAG_SIZE], Error> {
        if self.tag.is_some() {
            return Ok(self.tag.unwrap());
        }
        let state = self.enc_state.as_mut().ok_or(Error::InvalidState)?;
        let tag = state.finalize();
        self.tag = Some(tag);
        Ok(tag)
    }
}

pub struct QuasorDecryptStream<'a> {
    dec_state: DecryptState,
    nonce: &'a [u8],
    tag: &'a [u8],
    nonce_hasher: blake3::Hasher,
}

impl<'a> QuasorDecryptStream<'a> {
    fn new(
        master_key: &'a [u8; 32],
        nonce: &'a [u8],
        ad: &'a [u8],
        tag: &'a [u8],
        plaintext_len: u64,
    ) -> Self {
        let mut nonce_hasher = blake3::Hasher::new_keyed(master_key);
        nonce_hasher.update(&(ad.len() as u64).to_le_bytes());
        nonce_hasher.update(ad);
        nonce_hasher.update(&plaintext_len.to_le_bytes());
        Self {
            dec_state: DecryptState::new(master_key, nonce, ad),
            nonce,
            tag,
            nonce_hasher,
        }
    }

    pub fn process(&mut self, ciphertext_chunk: &[u8]) -> Vec<u8> {
        let plaintext_chunk = self.dec_state.process(ciphertext_chunk);
        self.nonce_hasher.update(&plaintext_chunk);
        plaintext_chunk
    }

    pub fn finalize(&mut self) -> Result<(), Error> {
        self.dec_state.verify_tag(self.tag)?;
        let mut derived_nonce = [0u8; 16];
        let mut output_reader = self.nonce_hasher.finalize_xof();
        output_reader.fill(&mut derived_nonce);
        if !bool::from(derived_nonce.as_slice().ct_eq(self.nonce)) {
            return Err(Error::InvalidTag);
        }
        Ok(())
    }
}

// --- ENCRYPT/DECRYPT STATE STRUCTS ---

pub struct EncryptState {
    sponge: Sponge,
    rekey_counter: u64,
    bytes_processed_since_rekey: usize,
    encryption_started: bool,
}

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

        // *** THE FIX: Add the length-prefix for the Associated Data ***
        sponge.absorb(&(ad.len() as u64).to_le_bytes());
        sponge.absorb(ad);

        Self {
            sponge,
            rekey_counter: 0,
            bytes_processed_since_rekey: 0,
            encryption_started: false,
        }
    }

    pub fn process(&mut self, plaintext: &[u8]) -> Vec<u8> {
        let mut ciphertext = Vec::with_capacity(plaintext.len());
        if !self.encryption_started {
            self.sponge.absorb(DOMAIN_ENCRYPT);
            self.encryption_started = true;
        }

        for chunk in plaintext.chunks(PROCESSING_CHUNK_SIZE) {
            // Original correct logic: squeeze then absorb.
            let keystream = self.sponge.squeeze(chunk.len());
            let mut cipher_chunk = Vec::with_capacity(chunk.len());
            for (p, k) in chunk.iter().zip(keystream.iter()) {
                cipher_chunk.push(p ^ k);
            }

            self.sponge.absorb(chunk);
            ciphertext.extend_from_slice(&cipher_chunk);
            self.bytes_processed_since_rekey += chunk.len();

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

    fn perform_rekey(&mut self) {
        self.sponge.absorb(DOMAIN_REKEY);
        self.sponge.absorb(&self.rekey_counter.to_le_bytes());
        let new_key = self.sponge.squeeze(MASTER_KEY_SIZE);
        self.sponge.absorb(&new_key);
        self.rekey_counter += 1;
    }

    pub fn finalize(&mut self) -> [u8; TAG_SIZE] {
        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,
    bytes_processed_since_rekey: usize,
    encryption_started: bool,
}

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

        // *** THE FIX: Add the length-prefix for the Associated Data ***
        sponge.absorb(&(ad.len() as u64).to_le_bytes());
        sponge.absorb(ad);

        Self {
            sponge,
            rekey_counter: 0,
            bytes_processed_since_rekey: 0,
            encryption_started: false,
        }
    }

    pub fn process(&mut self, ciphertext: &[u8]) -> Vec<u8> {
        let mut plaintext = Vec::with_capacity(ciphertext.len());
        if !self.encryption_started {
            self.sponge.absorb(DOMAIN_ENCRYPT);
            self.encryption_started = true;
        }

        for chunk in ciphertext.chunks(PROCESSING_CHUNK_SIZE) {
            // Original correct logic: squeeze then absorb.
            let keystream = self.sponge.squeeze(chunk.len());
            let mut plain_chunk = Vec::with_capacity(chunk.len());
            for (c, k) in chunk.iter().zip(keystream.iter()) {
                plain_chunk.push(c ^ k);
            }

            self.sponge.absorb(&plain_chunk);
            plaintext.extend_from_slice(&plain_chunk);
            self.bytes_processed_since_rekey += chunk.len();

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

    fn perform_rekey(&mut self) {
        self.sponge.absorb(DOMAIN_REKEY);
        self.sponge.absorb(&self.rekey_counter.to_le_bytes());
        let new_key = self.sponge.squeeze(MASTER_KEY_SIZE);
        self.sponge.absorb(&new_key);
        self.rekey_counter += 1;
    }

    pub fn verify_tag(&mut self, tag: &[u8]) -> Result<(), Error> {
        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)
        }
    }
}