Skip to content

Commit

Permalink
Implement BlockRngCore for ChaCha
Browse files Browse the repository at this point in the history
  • Loading branch information
@pitdicker
pitdicker committed Mar 3, 2018
1 parent 0a6ac86 commit 2b18577
Showing 1 changed file with 103 additions and 88 deletions.
191 changes: 103 additions & 88 deletions src/prng/chacha.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@
//! The ChaCha random number generator.
use core::fmt;
use {RngCore, CryptoRng, SeedableRng};
use {impls, le};
use {RngCore, BlockRngCore, CryptoRng, SeedableRng, Error};
use impls::BlockRng;
use le;

const SEED_WORDS: usize = 8; // 8 words for the 256-bit key
const STATE_WORDS: usize = 16;
Expand Down Expand Up @@ -62,45 +63,45 @@ const STATE_WORDS: usize = 16;
/// http://cr.yp.to/papers.html#xsalsa)
///
/// [`set_rounds`]: #method.set_counter
#[derive(Clone)]
pub struct ChaChaRng {
buffer: [u32; STATE_WORDS], // Internal buffer of output
state: [u32; STATE_WORDS], // Initial state
index: usize, // Index into state
rounds: usize,
}
#[derive(Clone, Debug)]
pub struct ChaChaRng(BlockRng<ChaChaCore>);

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for ChaChaRng {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "ChaChaRng {{}}")
impl RngCore for ChaChaRng {
#[inline]
fn next_u32(&mut self) -> u32 {
self.0.next_u32()
}
}

macro_rules! quarter_round{
($a: expr, $b: expr, $c: expr, $d: expr) => {{
$a = $a.wrapping_add($b); $d ^= $a; $d = $d.rotate_left(16);
$c = $c.wrapping_add($d); $b ^= $c; $b = $b.rotate_left(12);
$a = $a.wrapping_add($b); $d ^= $a; $d = $d.rotate_left( 8);
$c = $c.wrapping_add($d); $b ^= $c; $b = $b.rotate_left( 7);
}}
#[inline]
fn next_u64(&mut self) -> u64 {
self.0.next_u64()
}

#[inline]
fn fill_bytes(&mut self, dest: &mut [u8]) {
self.0.fill_bytes(dest)
}

#[inline]
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
self.0.try_fill_bytes(dest)
}
}

macro_rules! double_round{
($x: expr) => {{
// Column round
quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
// Diagonal round
quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
}}
impl SeedableRng for ChaChaRng {
type Seed = <ChaChaCore as SeedableRng>::Seed;

fn from_seed(seed: Self::Seed) -> Self {
ChaChaRng(BlockRng::<ChaChaCore>::from_seed(seed))
}

fn from_rng<R: RngCore>(rng: &mut R) -> Result<Self, Error> {
BlockRng::<ChaChaCore>::from_rng(rng).map(|rng| ChaChaRng(rng))
}
}

impl CryptoRng for ChaChaRng {}

impl ChaChaRng {
/// Create an ChaCha random number generator using the default
/// fixed key of 8 zero words.
Expand Down Expand Up @@ -153,11 +154,8 @@ impl ChaChaRng {
/// assert_eq!(rng1.next_u32(), rng2.next_u32());
/// ```
pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
self.state[12] = counter_low as u32;
self.state[13] = (counter_low >> 32) as u32;
self.state[14] = counter_high as u32;
self.state[15] = (counter_high >> 32) as u32;
self.index = STATE_WORDS; // force recomputation on next use
self.0.core.set_counter(counter_low, counter_high);
self.0.index = STATE_WORDS; // force recomputation on next use
}

/// Sets the number of rounds to run the ChaCha core algorithm per block to
Expand All @@ -180,13 +178,52 @@ impl ChaChaRng {
/// assert_eq!(rng.next_u32(), 0x2fef003e);
/// ```
pub fn set_rounds(&mut self, rounds: usize) {
assert!([4usize, 8, 12, 16, 20].iter().any(|x| *x == rounds));
self.rounds = rounds;
self.index = STATE_WORDS; // force recomputation on next use
self.0.core.set_rounds(rounds);
self.0.index = STATE_WORDS; // force recomputation on next use
}
}

/// Refill the internal output buffer (`self.buffer`)
fn update(&mut self) {
#[derive(Clone)]
pub struct ChaChaCore {
state: [u32; STATE_WORDS],
rounds: usize,
}

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for ChaChaCore {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "ChaChaCore {{}}")
}
}

macro_rules! quarter_round{
($a: expr, $b: expr, $c: expr, $d: expr) => {{
$a = $a.wrapping_add($b); $d ^= $a; $d = $d.rotate_left(16);
$c = $c.wrapping_add($d); $b ^= $c; $b = $b.rotate_left(12);
$a = $a.wrapping_add($b); $d ^= $a; $d = $d.rotate_left( 8);
$c = $c.wrapping_add($d); $b ^= $c; $b = $b.rotate_left( 7);
}}
}

macro_rules! double_round{
($x: expr) => {{
// Column round
quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
// Diagonal round
quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
}}
}

impl BlockRngCore<u32> for ChaChaCore {
type Results = [u32; STATE_WORDS];

fn generate(&mut self, results: &mut Self::Results) -> Result<(), Error> {
// For some reason extracting this part into a separate function
// improves performance by 50%.
fn core(results: &mut [u32; STATE_WORDS],
Expand All @@ -202,71 +239,49 @@ impl ChaChaRng {
}
}

core(&mut self.buffer, &self.state, self.rounds);
self.index = 0;
core(results, &self.state, self.rounds);

// update 128-bit counter
self.state[12] = self.state[12].wrapping_add(1);
if self.state[12] != 0 { return };
if self.state[12] != 0 { return Ok(()) };
self.state[13] = self.state[13].wrapping_add(1);
if self.state[13] != 0 { return };
if self.state[13] != 0 { return Ok(()) };
self.state[14] = self.state[14].wrapping_add(1);
if self.state[14] != 0 { return };
if self.state[14] != 0 { return Ok(()) };
self.state[15] = self.state[15].wrapping_add(1);
Ok(())
}
}

impl RngCore for ChaChaRng {
#[inline]
fn next_u32(&mut self) -> u32 {
// Using a local variable for `index`, and checking the size avoids a
// bounds check later on.
let mut index = self.index as usize;
if index >= STATE_WORDS {
self.update();
index = 0;
}

let value = self.buffer[index];
self.index += 1;
value
}

fn next_u64(&mut self) -> u64 {
impls::next_u64_via_u32(self)
impl ChaChaCore {
/// Sets the internal 128-bit ChaCha counter to a user-provided value. This
/// permits jumping arbitrarily ahead (or backwards) in the pseudorandom
/// stream.
pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
self.state[12] = counter_low as u32;
self.state[13] = (counter_low >> 32) as u32;
self.state[14] = counter_high as u32;
self.state[15] = (counter_high >> 32) as u32;
}


fn fill_bytes(&mut self, dest: &mut [u8]) {
let mut read_len = 0;
while read_len < dest.len() {
if self.index >= self.buffer.len() {
self.update();
}

let (consumed_u32, filled_u8) =
impls::fill_via_u32_chunks(&self.buffer[self.index..],
&mut dest[read_len..]);

self.index += consumed_u32;
read_len += filled_u8;
}
/// Sets the number of rounds to run the ChaCha core algorithm per block to
/// generate.
pub fn set_rounds(&mut self, rounds: usize) {
assert!([4usize, 8, 12, 16, 20].iter().any(|x| *x == rounds));
self.rounds = rounds;
}
}

impl CryptoRng for ChaChaRng {}

impl SeedableRng for ChaChaRng {
impl SeedableRng for ChaChaCore {
type Seed = [u8; SEED_WORDS*4];
fn from_seed(seed: Self::Seed) -> Self {
let mut seed_le = [0u32; SEED_WORDS];
le::read_u32_into(&seed, &mut seed_le);
ChaChaRng {
buffer: [0; STATE_WORDS],
Self {
state: [0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, // constants
seed_le[0], seed_le[1], seed_le[2], seed_le[3], // seed
seed_le[4], seed_le[5], seed_le[6], seed_le[7], // seed
0, 0, 0, 0], // counter
index: STATE_WORDS, // generate on first use
rounds: 20,
}
}
Expand Down

0 comments on commit 2b18577

Please sign in to comment.