purecrypt 0.1.0-alpha

#[cfg(feature = "rand")]
use rand_core::{CryptoRng, RngCore, SeedableRng};

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

use super::core::IETFChaChaCore;
use super::{NONCE_LEN, types::StreamId};
use crate::chacha::{Constants, Seed, consts::*};

const DEFAULT_STREAM_ID: [u8; NONCE_LEN] = [0; NONCE_LEN];

#[derive(Clone)]
#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
pub struct IETFChaChaRng<const ROUNDS: usize> {
    core: IETFChaChaCore<ROUNDS>,
    buffer: [u8; OUTPUT_LEN],
    buffer_pos: usize,
}

impl<const ROUNDS: usize> IETFChaChaRng<ROUNDS> {
    pub fn new<K, N>(seed: K, stream_id: N) -> Self
    where
        K: Into<Seed>,
        N: Into<StreamId>,
    {
        let key = seed.into().into_key();
        let nonce = stream_id.into().into_nonce();

        let buffer = [0; OUTPUT_LEN];
        let core = IETFChaChaCore::new(&key, &nonce);
        let buffer_pos = buffer.len();

        Self {
            core,
            buffer,
            buffer_pos,
        }
    }

    pub fn from_seed<S>(seed: S) -> Self
    where
        S: Into<Seed>,
    {
        Self::new(seed, StreamId::new(DEFAULT_STREAM_ID))
    }

    #[inline]
    pub fn get_stream_id(&self) -> StreamId {
        self.core.get_nonce().into_stream_id()
    }

    #[inline]
    pub fn set_stream_id<S>(&mut self, stream_id: S)
    where
        S: Into<StreamId>,
    {
        let nonce = stream_id.into().into_nonce();

        self.core.set_nonce(&nonce);
    }

    #[inline]
    pub fn get_counter(&self) -> u32 {
        self.core.get_counter()
    }

    pub fn set_counter(&mut self, counter: u32) {
        self.core.set_counter(counter);
    }

    #[inline]
    pub fn get_seed(&self) -> Seed {
        self.core.get_key().into_seed()
    }

    #[inline]
    pub fn set_seed<S>(&mut self, seed: S)
    where
        S: Into<[u8; KEY_LEN]>,
    {
        self.core.set_key(seed.into());
    }

    #[inline]
    pub fn get_constants(&self) -> Constants {
        self.core.get_constants()
    }

    pub fn set_constants<C>(&mut self, constants: C)
    where
        C: Into<Constants>,
    {
        self.core.set_constants(&constants.into());
    }

    pub fn fill_bytes(&mut self, mut dst: &mut [u8]) {
        // use the remaining buffer
        if self.buffer_pos < 64 {
            let take = dst.len().min(64 - self.buffer_pos);
            dst[..take].copy_from_slice(&self.buffer[self.buffer_pos..self.buffer_pos + take]);
            self.buffer_pos += take;
            dst = &mut dst[take..];
        }

        // filling in the main part of the dst
        while dst.len() >= 64 {
            self.refill();
            dst[..64].copy_from_slice(&self.buffer);
            dst = &mut dst[64..];
        }

        // filling in the tail
        if !dst.is_empty() {
            self.refill();
            let n = dst.len();
            dst.copy_from_slice(&self.buffer[..n]);
            self.buffer_pos = n;
        }
    }

    #[inline(always)]
    fn refill(&mut self) {
        self.core.generate_block(&mut self.buffer);
        self.buffer_pos = 0;
    }
}

#[cfg(feature = "rand")]
impl<const ROUNDS: usize> RngCore for IETFChaChaRng<ROUNDS> {
    fn next_u32(&mut self) -> u32 {
        let mut buf = [0; (u32::BITS / 8) as usize];
        self.fill_bytes(&mut buf);

        u32::from_le_bytes(buf)
    }

    fn next_u64(&mut self) -> u64 {
        let mut buf = [0; (u64::BITS / 8) as usize];
        self.fill_bytes(&mut buf);

        u64::from_le_bytes(buf)
    }

    fn fill_bytes(&mut self, dst: &mut [u8]) {
        self.fill_bytes(dst);
    }
}

#[cfg(feature = "rand")]
impl<const ROUNDS: usize> CryptoRng for IETFChaChaRng<ROUNDS> {}

#[cfg(feature = "rand")]
impl<const ROUNDS: usize> SeedableRng for IETFChaChaRng<ROUNDS> {
    type Seed = Seed;

    fn from_seed(seed: Self::Seed) -> Self {
        Self::from_seed(seed)
    }
}