zeros 17.3.2

/*
==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--

Zeros

Copyright (C) 2019-2025  Anonymous

There are several releases over multiple years,
they are listed as ranges, such as: "2019-2025".

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--
*/

//! # Chacha
//!
//! _Implemetation of [Chacha][site:chacha]._
//!
//! ## Notes
//!
//! - [Keys][type:Key] must be 256-bit, [IVs][type:IV] must be 128-bit.
//! - [`Write::write()`][fn:std/io/Write#write] implemetation of [`Chacha`][struct:Chacha]: see [`Chacha::encrypt()`][fn:Chacha#encrypt] for
//!   details.
//!
//! - Encrypting and decrypting take same steps:
//!
//!     + To encrypt, make new `Chacha` instance and write data into it.
//!     + To decrypt, make new `Chacha` instance and write (encrypted) data into it.
//!
//! - This implementation has been tested against [`openssl` crate][crate:openssl] (which uses [OpenSSL][site:OpenSSL] underneath).
//!
//! ## Examples
//!
//! Below is an example of encrypting/decrypting. You can use it for large data. For small amount of data, see [`Variant`][enum:Variant] for
//! some shortcuts.
//!
//! ```
//! use zeros::chacha::{self, Key, Variant};
//!
//! const KEY: Key = *b"00000000001111111111222222222233";
//! const IV: chacha::IV = *b"0123456789abcdef";
//! const DATA: &[u8] = b"test";
//!
//! let new_chacha20 = || Variant::Twenty.new(KEY, IV, Vec::with_capacity(DATA.len()));
//! let encrypted = {
//!     let mut chacha20 = new_chacha20();
//!     chacha20.encrypt(DATA)?;
//!     chacha20.finish()?
//! };
//! assert_eq!(DATA.len(), encrypted.len());
//! assert_ne!(DATA, encrypted);
//!
//! // Decrypting
//! assert_eq!(DATA, {
//!     let mut chacha20 = new_chacha20();
//!     chacha20.encrypt(encrypted)?;
//!     chacha20.finish()?
//! });
//!
//! # zeros::Result::Ok(())
//! ```
//!
//! ## References:
//!
//! - <https://cr.yp.to/chacha.html>
//! - <https://cr.yp.to/streamciphers/timings/estreambench/submissions/salsa20/chacha20/ref/chacha.c>
//! - <https://cr.yp.to/streamciphers/timings/estreambench/submissions/salsa20/chacha20/ref/ecrypt-sync.h>
//! - <https://cr.yp.to/snuffle/ecrypt-config.h>
//! - <https://cr.yp.to/snuffle/ecrypt-portable.h>
//!
//! [site:chacha]: https://cr.yp.to/chacha.html
//! [site:OpenSSL]: https://www.openssl.org
//!
//! [crate:openssl]: https://crates.io/crates/openssl
//!
//! [enum:Variant]: enum.Variant.html
//! [struct:Chacha]: struct.Chacha.html
//! [type:Key]: type.Key.html
//! [type:IV]: type.IV.html
//! [fn:Chacha#encrypt]: struct.Chacha.html#method.encrypt
//! [fn:Chacha#finish]: struct.Chacha.html#method.finish
//!
//! [fn:std/io/Write#write]: https://doc.rust-lang.org/std/io/trait.Write.html#tymethod.write

#![cfg(target_endian="little")]
#![doc(cfg(target_endian="little"))]

use {
    alloc::vec::Vec,
    crate::{Bytes, Result},
    self::word_array::WordArray,
};

#[cfg(not(feature="std"))]
use crate::io::Write;

#[cfg(feature="std")]
use {
    std::io::Write,
    crate::IoResult,
};

#[cfg(feature="simd")]
use core::simd::Simd;

#[cfg(test)]
use {
    core::mem,
    crate::keccak::Hash,
    openssl::symm::{Cipher, Crypter, Mode},
};

#[cfg(test)]
#[cfg(feature="std")]
use std::thread;

#[cfg(test)]
mod tests;

mod variant;

pub use self::variant::*;

#[cfg(not(feature="simd"))]
#[doc(cfg(not(feature="simd")))]
pub (crate) mod salsa20;

#[cfg(feature="simd")]
#[doc(cfg(feature="simd"))]
pub (crate) mod salsa20_simd;

pub (crate) mod word_array;

/// # Key (`256` bits)
pub type Key = [u8; KEY_SIZE_IN_BYTES];

/// # Key size in bytes
const KEY_SIZE_IN_BYTES: usize = 32;

/// # IV (`128` bits)
///
/// ## Notes
///
/// This is __`128-bit`__, which is different to [`poly1305::chacha20::IV`][type:poly1305/chacha20/IV] (`96` bits), or
/// [`poly1305::xchacha20::IV`][type:poly1305/xchacha20/IV] (`192` bits).
///
/// [type:poly1305/chacha20/IV]: ../poly1305/chacha20/type.IV.html
/// [type:poly1305/xchacha20/IV]: ../poly1305/xchacha20/type.IV.html
pub type IV = [u8; IV_SIZE_IN_BYTES];

/// # IV size in bytes
const IV_SIZE_IN_BYTES: usize = 16;

/// # Block size in bytes
pub (crate) const BLOCK_SIZE_IN_BYTES: usize = 64;

/// # Block
pub (crate) type Block = [u8; BLOCK_SIZE_IN_BYTES];

#[cfg(any(feature="simd", test))]
const BLOCK_OF_ZEROS: Block = [0; BLOCK_SIZE_IN_BYTES];

#[cfg(test)]
const KEY_HASH: Hash = Hash::Sha3_256;
#[cfg(test)]
const IV_HASH: Hash = Hash::Shake128;

/// # Chacha20
pub (crate) const TWENTY: Variant = Variant::Twenty;

/// # Chacha
///
/// You can make new instance via [`Variant`][enum:Variant].
///
/// [enum:Variant]: enum.Variant.html
#[derive(Debug)]
pub struct Chacha<W> where W: Write {
    variant: Variant,
    data: WordArray,
    buffer: Vec<u8>,
    output: W,
}

impl<W> Chacha<W> where W: Write {

    /// # Gets variant
    pub const fn variant(&self) -> &Variant {
        &self.variant
    }

    /// # Encrypts new bytes
    ///
    /// This function always returns your data's size, but it does *NOT* always write your whole data out. It's because encrypting works only on
    /// blocks of a fixed size. Exceeded data will be kept in an internal buffer. That buffer will be processed in
    /// [`finish()`](#method.finish).
    #[inline]
    pub fn encrypt<B>(&mut self, bytes: B) -> Result<usize> where B: AsRef<[u8]> {
        let mut bytes = bytes.as_ref();

        let result = bytes.len();
        if result == 0 {
            return Ok(result);
        }

        // Process buffer first
        if crate::io::fill_buffer(&mut self.buffer, BLOCK_SIZE_IN_BYTES, &mut bytes) {
            self.process_buffer_and_write_to_output(None)?;
            self.buffer.clear();
        }

        // Now the bytes
        {
            // Docs say chunks_exact() can be optimized better than chunks()
            let chunks = bytes.chunks_exact(BLOCK_SIZE_IN_BYTES);
            self.buffer.extend(chunks.remainder());
            for c in chunks {
                self.process_buffer_and_write_to_output(Some(c))?;
            }
        }

        Ok(result)
    }

    /// # Encrypts new data
    ///
    /// Returns length of input bytes. If overflow happens, `None` will be returned.
    ///
    /// See [`encrypt()`](#method.encrypt) for more details.
    pub fn encrypt_bytes<'a, const N: usize, B, B0>(&mut self, bytes: B) -> Result<Option<usize>>
    where B: Into<Bytes<'a, N, B0>>, B0: AsRef<[u8]> + 'a {
        let mut result = Some(usize::MIN);
        for bytes in bytes.into().as_slice() {
            let size = self.encrypt(bytes)?;
            if let Some(current) = result.as_mut() {
                match current.checked_add(size) {
                    Some(new) => *current = new,
                    None => result = None,
                };
            }
        }

        Ok(result)
    }

    /// # Processes buffer and writes to output
    ///
    /// - If buffer is `None`, self buffer will be used.
    /// - Buffer's size must be less than or equal to `BLOCK_SIZE_IN_BYTES`. Or an error will be returned.
    /// - Only the final buffer can have size less than `BLOCK_SIZE_IN_BYTES`. Or the encrypted data will be **BROKEN**. This is
    ///   responsibility of the caller (you).
    #[inline]
    fn process_buffer_and_write_to_output(&mut self, buffer: Option<&[u8]>) -> Result<()> {
        let buffer = buffer.unwrap_or(&self.buffer);

        let buffer_len = match buffer.len() {
            0 => return Ok(()),
            buffer_len @ 1..=BLOCK_SIZE_IN_BYTES => buffer_len,
            other => return Err(err!("Buffer is too large: {other} (max allowed: {max})", other=other, max=BLOCK_SIZE_IN_BYTES)),
        };

        #[cfg(feature="simd")]
        let mut output = salsa20_simd::words_to_bytes(&self.variant, &self.data);
        #[cfg(not(feature="simd"))]
        let mut output = salsa20::words_to_bytes(&self.variant, &self.data);
        {
            const TWELFTH: usize = 12;
            self.data[TWELFTH] = self.data[TWELFTH].wrapping_add(1);
            if self.data[TWELFTH] == 0 {
                const THIRTEENTH: usize = TWELFTH + 1;
                self.data[THIRTEENTH] = self.data[THIRTEENTH].wrapping_add(1);
            }
        }

        #[cfg(feature="simd")] {
            output = {
                let mut tmp = BLOCK_OF_ZEROS;
                tmp[..buffer_len].copy_from_slice(buffer);
                (Simd::from_array(output) ^ Simd::from_array(tmp)).to_array()
            };
        }
        #[cfg(not(feature="simd"))]
        (0..buffer_len).for_each(|i| output[i] ^= buffer[i]);

        let result = self.output.write_all(&output[..buffer_len]);
        #[cfg(feature="std")]
        let result = result.map_err(|e| from_io_err!(e));
        result
    }

    /// # Mutable Output
    #[cfg(feature="std")]
    #[inline(always)]
    pub (crate) const fn mut_output(&mut self) -> &mut W {
        &mut self.output
    }

    /// # Finishes and returns your output
    #[must_use]
    pub fn finish(mut self) -> Result<W> {
        self.process_buffer_and_write_to_output(None)?;
        drop(self.buffer);

        #[cfg(feature="std")]
        from_io_err!(self.output.flush())?;

        Ok(self.output)
    }

}

#[cfg(feature="std")]
#[doc(cfg(feature="std"))]
impl<W> Write for Chacha<W> where W: Write {

    fn write(&mut self, buffer: &[u8]) -> IoResult<usize> {
        Ok(self.encrypt(buffer)?)
    }

    fn flush(&mut self) -> IoResult<()> {
        self.output.flush()
    }

}

#[test]
fn tests() {
    assert_eq!(mem::size_of::<Block>(), mem::size_of::<WordArray>());

    assert_eq!(KEY_SIZE_IN_BYTES, 32);
    assert_eq!(mem::size_of::<Key>(), KEY_SIZE_IN_BYTES);

    assert_eq!(IV_SIZE_IN_BYTES, 16);
    assert_eq!(mem::size_of::<IV>(), IV_SIZE_IN_BYTES);
}

/// # Encrypts using OpenSSL's Chacha20
#[cfg(test)]
fn encrypt_using_open_ssl_chacha20<K, IV, B>(mode: Mode, key: K, iv: IV, bytes: B) -> Vec<u8>
where K: AsRef<[u8]>, IV: AsRef<[u8]>, B: AsRef<[u8]> {
    let bytes = bytes.as_ref();

    let mut crypter = Crypter::new(Cipher::chacha20(), mode, key.as_ref(), Some(iv.as_ref())).unwrap();
    let mut result = Vec::with_capacity(bytes.len());

    let mut output = BLOCK_OF_ZEROS;
    for c in bytes.chunks(BLOCK_SIZE_IN_BYTES) {
        let count = crypter.update(c, &mut output).unwrap();
        result.extend(&output[..count]);
    }
    {
        let count = crypter.finalize(&mut output).unwrap();
        result.extend(&output[..count]);
    }

    result
}

/// OpenSSL says that for IVs:
///
/// - The first 4 bytes is a counter.
/// - The last 12 bytes is nonce.
///
/// (Original C implementation does not mention counter anywhere. It mentions 'nonce' once, but vaguely.)
///
/// See EVP_chacha20(3ossl).
#[test]
fn cmp_to_data_generated_by_open_ssl() -> Result<()> {
    const DATA: &[u8] = &[
        0x50, 0xa6, 0xb7, 0xec, 0xb4, 0x2c, 0xc8, 0x9a, 0xbd, 0x5c, 0x40, 0xea, 0x5e, 0x30, 0x66, 0x31, 0xbf, 0x93, 0x49, 0x6a, 0xc5, 0x01,
        0x56, 0xb3, 0x6f, 0x51, 0x56, 0xe8, 0x56, 0x89, 0xe8, 0x56, 0x08, 0xe5, 0xb0, 0xe6, 0xa0, 0xd9, 0x3c, 0x1c, 0x6a, 0x8a, 0xd7, 0x12,
        0x09, 0xf0, 0xac, 0x9d, 0x57, 0xb0, 0x45, 0x33, 0x9d, 0x1f, 0x60, 0x1d, 0x34, 0xf8, 0xa7, 0xa0, 0x4e, 0x42, 0x64, 0xae,
        0x98, 0xda, 0xad, 0x4a, 0x94, 0x82, 0xcc, 0x9b, 0x84, 0x82, 0x1f, 0x50, 0x4f, 0xc0, 0x44, 0xba,
    ];

    let key = &KEY_HASH.hash("key");
    let iv = &IV_HASH.hash("IV");
    let encrypt = |data| TWENTY.encrypt(key, iv, data);
    let encrypt_using_open_ssl_chacha20 = |mode, data| encrypt_using_open_ssl_chacha20(mode, key, iv, data);

    let encrypted = encrypt(DATA)?;
    assert_eq!(DATA.len(), encrypted.len());
    assert_ne!(DATA, encrypted);
    assert_eq!(encrypted, encrypt_using_open_ssl_chacha20(Mode::Encrypt, DATA));

    assert_eq!(DATA, encrypt(&encrypted)?);
    assert_eq!(DATA, encrypt_using_open_ssl_chacha20(Mode::Decrypt, &encrypted));

    Ok(())
}

#[test]
#[cfg(feature="std")]
fn cmp_to_data_generated_by_open_ssl_using_threads() -> Result<()> {
    const INPUT_DATA: &[u8] = &[u8::MIN; BLOCK_SIZE_IN_BYTES + 1];

    (-1_i16..=999).map(|index| thread::spawn(move || {
        let key = &if index >= 0 {
            KEY_HASH.hash(index.to_be_bytes())
        } else {
            // Tests for overflowing additions in `salsa20::words_to_bytes()`
            vec![u8::MAX; KEY_SIZE_IN_BYTES]
        };
        let iv = &if index % 2 == 0 {
            vec![index as u8; IV_SIZE_IN_BYTES]
        } else {
            IV_HASH.hash(index.to_le_bytes())
        };
        assert_ne!(key, iv);

        let encrypt = |data| TWENTY.encrypt(key, iv, data);
        let encrypt_using_open_ssl_chacha20 = |mode, data| encrypt_using_open_ssl_chacha20(mode, key, iv, data);

        // Encrypt
        let encrypted_data = encrypt(INPUT_DATA)?;
        assert_eq!(INPUT_DATA.len(), encrypted_data.len());
        assert_eq!(encrypted_data, encrypt_using_open_ssl_chacha20(Mode::Encrypt, INPUT_DATA));

        // Decrypt
        assert_eq!(INPUT_DATA, encrypt(&encrypted_data)?);
        assert_eq!(INPUT_DATA, encrypt_using_open_ssl_chacha20(Mode::Decrypt, &encrypted_data));

        Result::Ok(())
    })).collect::<Vec<_>>().into_iter().for_each(|t| t.join().unwrap().unwrap());

    Ok(())
}