chacha20 0.10.0

#![allow(clippy::cast_possible_truncation, reason = "needs triage")]
#![allow(clippy::undocumented_unsafe_blocks, reason = "TODO")]

use core::fmt;

use rand_core::{
    Infallible, SeedableRng, TryCryptoRng, TryRng,
    block::{BlockRng, Generator},
};

#[cfg(feature = "zeroize")]
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::{
    ChaChaCore, R8, R12, R20, Rounds, backends,
    variants::{Legacy, Variant},
};

use cfg_if::cfg_if;

/// Seed value used to initialize ChaCha-based RNGs.
pub type Seed = [u8; 32];

/// Serialized RNG state.
pub type SerializedRngState = [u8; 49];

/// Number of 32-bit words per ChaCha block (fixed by algorithm definition).
pub(crate) const BLOCK_WORDS: u8 = 16;

/// Number of blocks generated by RNG core.
const BUF_BLOCKS: u8 = 4;
/// Buffer size in words used by buffered RNG.
const BUFFER_SIZE: usize = (BLOCK_WORDS * BUF_BLOCKS) as usize;

impl<R: Rounds, V: Variant> SeedableRng for ChaChaCore<R, V> {
    type Seed = Seed;

    #[inline]
    fn from_seed(seed: Self::Seed) -> Self {
        ChaChaCore::new_internal(&seed, &[0u8; 8])
    }
}

impl<R: Rounds, V: Variant> Generator for ChaChaCore<R, V> {
    type Output = [u32; BUFFER_SIZE];

    /// Generates 4 blocks in parallel with avx2 & neon, but merely fills
    /// 4 blocks with sse2 & soft
    fn generate(&mut self, buffer: &mut [u32; BUFFER_SIZE]) {
        cfg_if! {
            if #[cfg(chacha20_backend = "soft")] {
                backends::soft::Backend(self).gen_ks_blocks(buffer);
            } else if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
                cfg_if! {
                    // AVX-512 doesn't support RNG, so use AVX-2 instead
                    if #[cfg(any(chacha20_backend = "avx2", chacha20_backend = "avx512"))] {
                        unsafe {
                            backends::avx2::rng_inner::<R, V>(self, buffer);
                        }
                    } else if #[cfg(chacha20_backend = "sse2")] {
                        unsafe {
                            backends::sse2::rng_inner::<R, V>(self, buffer);
                        }
                    } else {
                        #[cfg(chacha20_avx512)]
                        let (_avx512_token, avx2_token, sse2_token) = self.tokens;
                        #[cfg(not(chacha20_avx512))]
                        let (avx2_token, sse2_token) = self.tokens;

                        if avx2_token.get() {
                            unsafe {
                                backends::avx2::rng_inner::<R, V>(self, buffer);
                            }
                        } else if sse2_token.get() {
                            unsafe {
                                backends::sse2::rng_inner::<R, V>(self, buffer);
                            }
                        } else {
                            backends::soft::Backend(self).gen_ks_blocks(buffer);
                        }
                    }
                }
            } else if #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] {
                // SAFETY: we have used conditional compilation to ensure NEON is available
                unsafe {
                    backends::neon::rng_inner::<R, V>(self, buffer);
                }
            } else {
                backends::soft::Backend(self).gen_ks_blocks(buffer);
            }
        }
    }

    // `Drop` impl of `BlockRng` calls this method and passes reference to
    // its internal buffer in `output`. So we zeroize its contents here.
    #[cfg(feature = "zeroize")]
    fn drop(&mut self, output: &mut Self::Output) {
        output.zeroize();
    }
}

macro_rules! impl_chacha_rng {
    ($Rng:ident, $rounds:ident) => {
        /// A cryptographically secure random number generator that uses the ChaCha stream cipher.
        ///
        /// See the [crate docs][crate] for more information about the underlying stream cipher.
        ///
        /// This RNG implementation uses a 64-bit counter and 64-bit stream identifier (a.k.a nonce).
        /// A 64-bit counter over 64-byte (16 word) blocks allows 1 ZiB of output before cycling,
        /// and the stream identifier allows 2<sup>64</sup> unique streams of output per seed.
        /// Both counter and stream are initialized to zero but may be set via the [`set_word_pos`]
        /// and [`set_stream`] methods.
        ///
        /// [`set_word_pos`]: Self::set_word_pos
        /// [`set_stream`]: Self::set_stream
        ///
        /// # Example
        ///
        /// ```rust
        #[doc = concat!("use chacha20::", stringify!($Rng), ";")]
        /// use rand_core::{SeedableRng, Rng};
        ///
        /// let seed = [42u8; 32];
        #[doc = concat!("let mut rng = ", stringify!($Rng), "::from_seed(seed);")]
        ///
        /// let random_u32 = rng.next_u32();
        /// let random_u64 = rng.next_u64();
        ///
        /// let mut random_bytes = [0u8; 3];
        /// rng.fill_bytes(&mut random_bytes);
        /// ```
        ///
        /// See the [`rand`](https://docs.rs/rand/) crate for more advanced RNG functionality.
        pub struct $Rng {
            core: BlockRng<ChaChaCore<$rounds, Legacy>>,
        }

        impl SeedableRng for $Rng {
            type Seed = Seed;

            #[inline]
            fn from_seed(seed: Self::Seed) -> Self {
                let core = ChaChaCore::new_internal(&seed, &[0u8; 8]);
                Self {
                    core: BlockRng::new(core),
                }
            }
        }

        impl TryRng for $Rng {
            type Error = Infallible;

            #[inline]
            fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
                Ok(self.core.next_word())
            }
            #[inline]
            fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
                Ok(self.core.next_u64_from_u32())
            }
            #[inline]
            fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
                self.core.fill_bytes(dest);
                Ok(())
            }
        }

        impl TryCryptoRng for $Rng {}

        #[cfg(feature = "zeroize")]
        impl ZeroizeOnDrop for $Rng {}

        // We use custom implementation of `PartialEq` because RNG states
        // may buffer different parts of the same keystream, while keeping
        // buffer cursor pointing towards the same keystream point.
        impl PartialEq<$Rng> for $Rng {
            fn eq(&self, rhs: &$Rng) -> bool {
                (self.get_seed() == rhs.get_seed())
                    && (self.get_stream() == rhs.get_stream())
                    && (self.get_word_pos() == rhs.get_word_pos())
            }
        }

        impl Eq for $Rng {}

        // Custom Debug implementation that does not expose the internal state
        impl fmt::Debug for $Rng {
            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
                write!(f, concat!(stringify!($Rng), " {{ ... }}"))
            }
        }

        impl $Rng {
            /// Get the offset from the start of the stream, in 32-bit words.
            ///
            /// Since the generated blocks are 64 words (2<sup>6</sup>) long and the
            /// counter is 64-bits, the offset is a 68-bit number. Sub-word offsets are
            /// not supported, hence the result can simply be multiplied by 4 to get a
            /// byte-offset.
            #[inline]
            #[must_use]
            pub fn get_word_pos(&self) -> u128 {
                let mut block_counter = (u64::from(self.core.core.state[13]) << 32)
                    | u64::from(self.core.core.state[12]);
                if self.core.word_offset() != 0 {
                    block_counter = block_counter.wrapping_sub(u64::from(BUF_BLOCKS));
                }
                let word_pos = u128::from(block_counter) * u128::from(BLOCK_WORDS)
                    + self.core.word_offset() as u128;
                // eliminate bits above the 68th bit
                word_pos & ((1 << 68) - 1)
            }

            /// Set the offset from the start of the stream, in 32-bit words.
            ///
            /// **This value will be erased when calling `set_stream()`,
            /// so call `set_stream()` before calling `set_word_pos()`**
            /// if you intend on using both of them together.
            ///
            /// As with `get_word_pos`, we use a 68-bit number. Since the generator
            /// simply cycles at the end of its period (1 ZiB), we ignore the upper
            /// 60 bits.
            #[inline]
            pub fn set_word_pos(&mut self, word_offset: u128) {
                let index = (word_offset % u128::from(BLOCK_WORDS)) as usize;
                let counter = word_offset / u128::from(BLOCK_WORDS);
                //self.set_block_pos(counter as u64);
                self.core.core.state[12] = counter as u32;
                self.core.core.state[13] = (counter >> 32) as u32;
                self.core.reset_and_skip(index);
            }

            /// Sets the block pos and resets the RNG's index.
            ///
            /// **This value will be erased when calling `set_stream()`,
            /// so call `set_stream()` before calling `set_block_pos()`**
            /// if you intend on using both of them together.
            ///
            /// The word pos will be equal to `block_pos * 16 words per block`.
            #[inline]
            #[allow(unused)]
            pub fn set_block_pos(&mut self, block_pos: u64) {
                self.core.reset_and_skip(0);
                self.core.core.set_block_pos(block_pos);
            }

            /// Get the block pos.
            #[inline]
            #[allow(unused)]
            #[must_use]
            pub fn get_block_pos(&self) -> u64 {
                let counter = self.core.core.get_block_pos();
                let offset = self.core.word_offset();
                if offset != 0 {
                    counter - u64::from(BUF_BLOCKS) + offset as u64 / 16
                } else {
                    counter
                }
            }

            /// Set the stream ID and reset the `word_pos` to 0.
            #[inline]
            pub fn set_stream(&mut self, stream: u64) {
                self.core.core.state[14] = stream as u32;
                self.core.core.state[15] = (stream >> 32) as u32;
                self.set_block_pos(0);
            }

            /// Get the stream number (nonce).
            #[inline]
            #[must_use]
            pub fn get_stream(&self) -> u64 {
                let mut result = [0u8; 8];
                result[..4].copy_from_slice(&self.core.core.state[14].to_le_bytes());
                result[4..].copy_from_slice(&self.core.core.state[15].to_le_bytes());
                u64::from_le_bytes(result)
            }

            /// Get the RNG seed.
            #[inline]
            #[must_use]
            pub fn get_seed(&self) -> [u8; 32] {
                let seed = &self.core.core.state[4..12];
                let mut result = [0u8; 32];
                for (src, dst) in seed.iter().zip(result.chunks_exact_mut(4)) {
                    dst.copy_from_slice(&src.to_le_bytes())
                }
                result
            }

            /// Serialize RNG state.
            ///
            /// # Warning
            /// Leaking serialized RNG state to an attacker defeats security properties
            /// provided by the RNG.
            #[inline]
            pub fn serialize_state(&self) -> SerializedRngState {
                let seed = self.get_seed();
                let stream = self.get_stream().to_le_bytes();
                let word_pos = self.get_word_pos().to_le_bytes();

                let mut res = [0u8; 49];
                let (seed_dst, res_rem) = res.split_at_mut(32);
                let (stream_dst, word_pos_dst) = res_rem.split_at_mut(8);

                seed_dst.copy_from_slice(&seed);
                stream_dst.copy_from_slice(&stream);
                word_pos_dst.copy_from_slice(&word_pos[..9]);

                debug_assert_eq!(&word_pos[9..], &[0u8; 7]);

                res
            }

            /// Deserialize RNG state.
            #[inline]
            pub fn deserialize_state(state: &SerializedRngState) -> Self {
                let (seed, state_rem) = state.split_at(32);
                let (stream, word_pos_raw) = state_rem.split_at(8);

                let seed: &[u8; 32] = seed.try_into().expect("seed.len() is equal to 32");
                let stream: &[u8; 8] = stream.try_into().expect("stream.len() is equal to 8");

                // Note that we use only 68 bits from `word_pos_raw`, i.e. 4 remaining bits
                // get ignored and should be equal to zero in practice.
                let mut word_pos_buf = [0u8; 16];
                word_pos_buf[..9].copy_from_slice(word_pos_raw);
                let word_pos = u128::from_le_bytes(word_pos_buf);

                let core = ChaChaCore::new_internal(seed, stream);
                let mut res = Self {
                    core: BlockRng::new(core),
                };

                res.set_word_pos(word_pos);
                res
            }
        }
    };
}

impl_chacha_rng!(ChaCha8Rng, R8);
impl_chacha_rng!(ChaCha12Rng, R12);
impl_chacha_rng!(ChaCha20Rng, R20);