Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement BufOutputReader; add rand_core impls for BufOutputReader #362

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ mmap = ["std", "dep:memmap2"]
# Implement the zeroize::Zeroize trait for types in this crate.
zeroize = ["dep:zeroize", "arrayvec/zeroize"]

# Implement the rand_core Rng traits for OutputReader, allowing using it as
# a rand::Rng.
rand = ["dep:rand_core"]

# This crate implements traits from the RustCrypto project, exposed here as the
# "traits-preview" feature. However, these traits aren't stable, and they're
# expected to change in incompatible ways before they reach 1.0. For that
Expand Down Expand Up @@ -88,8 +92,8 @@ no_avx512 = []
no_neon = []

[package.metadata.docs.rs]
# Document the rayon/mmap methods and the Serialize/Deserialize/Zeroize impls on docs.rs.
features = ["mmap", "rayon", "serde", "zeroize"]
# Document the rayon/mmap methods and the Serialize/Deserialize/Zeroize/RngCore impls on docs.rs.
features = ["mmap", "rayon", "serde", "zeroize", "rand"]

[dependencies]
arrayref = "0.3.5"
Expand All @@ -101,6 +105,7 @@ memmap2 = { version = "0.7.1", optional = true }
rayon = { version = "1.2.1", optional = true }
serde = { version = "1.0", default-features = false, features = ["derive"], optional = true }
zeroize = { version = "1", default-features = false, features = ["zeroize_derive"], optional = true }
rand_core = { version = "0.6", default-features = false, optional = true }

[dev-dependencies]
hmac = "0.12.0"
Expand Down
194 changes: 194 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1643,3 +1643,197 @@ impl std::io::Seek for OutputReader {
Ok(self.position())
}
}

/// A buffering wrapper around [`OutputReader`].
///
/// This fills some of the simpler niches of a [`std::io::BufReader`] for no_std
/// and rng use-cases that don't need a full [`std::io::BufReader`]. If you
/// need the [`std::io`] traits with buffering, you're probably better off with
/// a full [`std::io::BufReader`] wrapper around [`OutputReader`].
///
/// With the `rand` feature, this struct implements [`rand_core::RngCore`],
/// [`rand_core::SeedableRng`], and [`rand_core::CryptoRng`], allowing this
/// type to be used as a full [`rand::Rng`]. A [`Rng`] type alias is given as a
/// convenient suggested buffer size for Rng use.
///
/// [`std::io`]: https://doc.rust-lang.org/std/io/index.html
/// [`std::io::BufReader`]: https://doc.rust-lang.org/std/io/struct.BufReader.html
/// [`OutputReader`]: struct.OutputReader.html
/// [`rand_core::RngCore`]: https://rust-random.github.io/rand/rand_core/trait.RngCore.html
/// [`rand_core::SeedableRng`]: https://rust-random.github.io/rand/rand_core/trait.SeedableRng.html
/// [`rand_core::CryptoRng`]: https://rust-random.github.io/rand/rand_core/trait.CryptoRng.html
/// [`rand::Rng`]: https://docs.rs/rand/latest/rand/trait.Rng.html
/// [`Rng`]: type.Rng.html
#[derive(Clone, Debug)]
pub struct BufOutputReader<const N: usize = 64> {
reader: OutputReader,
buffer: [u8; N],

/// The amount of buffer that has been read already.
offset: usize,
}

impl<const N: usize> BufOutputReader<N> {
#[inline]
pub fn new(reader: OutputReader) -> Self {
reader.into()
}

/// The position in the output stream, minus the remaining characters in
/// the buffer.
#[inline]
pub fn position(&self) -> u64 {
let buffered = (N - self.offset) as u64;
self.reader.position() - buffered
}

/// Drop what's remaining in the buffer and give a mutable reference to the
/// inner reader, so it can be seeked or otherwise manipulated.
#[inline]
pub fn output_reader(&mut self) -> &mut OutputReader {
self.offset = N;
&mut self.reader
}

/// Efficiently fill the destination buffer, calling the underlying
/// [`OutputReader::fill`] as few times as possible.
///
/// [`OutputReader::fill`]: struct.OutputReader.html#method.fill
pub fn fill(&mut self, mut dest: &mut [u8]) {
if dest.is_empty() {
return;
}

let buffer_remaining = N - self.offset;

if dest.len() <= buffer_remaining {
// There are enough bytes left in the buffer to consume without
// reading.
let end = self.offset + dest.len();
dest.copy_from_slice(&self.buffer[self.offset..end]);
self.offset = end;
} else {
// First empty the buffer.
if buffer_remaining > 0 {
dest[..buffer_remaining].copy_from_slice(&self.buffer[self.offset..N]);
let copied = N - self.offset;
dest = &mut dest[copied..];
}

let buffers = dest.len() / N;
let remainder = dest.len() % N;

// Copy full-sized chunks directly to the destination, bypassing
// the buffer.
if buffers > 0 {
let buffers_bytes = buffers * N;
self.reader.fill(&mut dest[..buffers_bytes]);
dest = &mut dest[buffers_bytes..];
}

// Fill the buffer for the remainder, if there is any.
if remainder > 0 {
self.reader.fill(&mut self.buffer);
dest.copy_from_slice(&self.buffer[..remainder]);
self.offset = remainder;
} else {
// We have emptied the remaining buffer, so mark this empty.
self.offset = N;
}
}
}
}

impl<const N: usize> From<OutputReader> for BufOutputReader<N> {
fn from(value: OutputReader) -> Self {
Self {
reader: value,
buffer: [0u8; N],

// Start buffer unfilled.
offset: N,
}
}
}

#[cfg(feature = "rand")]
impl<const N: usize> rand_core::SeedableRng for BufOutputReader<N> {
type Seed = [u8; 32];

#[inline]
fn from_seed(seed: Self::Seed) -> Self {
Hasher::new_keyed(&seed).finalize_xof().into()
}
}

#[cfg(feature = "rand")]
impl<const N: usize> rand_core::RngCore for BufOutputReader<N> {
#[inline]
fn next_u32(&mut self) -> u32 {
rand_core::impls::next_u32_via_fill(self)
Taywee marked this conversation as resolved.
Show resolved Hide resolved
}

#[inline]
fn next_u64(&mut self) -> u64 {
rand_core::impls::next_u64_via_fill(self)
}

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

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

#[cfg(feature = "rand")]
impl<const N: usize> rand_core::block::BlockRngCore for BufOutputReader<N>
where
[u8; N]: Default,
{
type Item = u8;
type Results = [u8; N];

fn generate(&mut self, results: &mut Self::Results) {
self.fill(results);
}
}

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

#[cfg(feature = "rand")]
/// A convenience type alias for the recommended Rng buffer size.
///
/// # Examples
///
/// ```
/// # use rand::{Rng as _, SeedableRng as _};
/// # fn main() {
/// // Hash input and convert the output stream to an rng.
/// let mut hasher = blake3::Hasher::new();
/// hasher.update(b"foo");
/// hasher.update(b"bar");
/// hasher.update(b"baz");
/// let mut rng: blake3::Rng = hasher.finalize_xof().into();
/// let output: u64 = rng.gen();
/// assert_eq!(output, 0xfb61f3c9e0fe9ac0u64);
///
/// // Alternately, seed it as a rand::SeedableRng.
/// let mut rng = blake3::Rng::from_seed(*b"0123456789abcdefghijklmnopqrstuv");
/// let output: u64 = rng.gen();
/// assert_eq!(output, 0x9958c58595366357u64);
///
/// // In the real world, you will probably not use a static seed, but seed from
/// // OsRng or something of the sort.
/// let mut seed = [0u8; 32];
/// rand::rngs::OsRng.fill(&mut seed);
/// let mut rng = blake3::Rng::from_seed(seed);
/// let _output: u64 = rng.gen();
/// # }
/// ```
pub type Rng = BufOutputReader<64>;
42 changes: 42 additions & 0 deletions src/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -820,3 +820,45 @@ fn test_serde() {
let hash2: crate::Hash = serde_json::from_str(&json).unwrap();
assert_eq!(hash, hash2);
}

#[test]
#[cfg(feature = "rand")]
fn test_rand_core() {
let mut seeded = crate::Rng::from_seed(*b"0123456789abcdefghijklmnopqrstuv");
let mut buf = [0u8; 64];
seeded.fill_bytes(&mut buf);
// Verified using: printf 0123456789abcdefghijklmnopqrstuv | b3sum -l 76 --keyed <(true)
assert_eq!(
&buf,
b"\
\x57\x63\x36\x95\x85\xc5\x58\x99\x4a\x3e\xe0\x27\x78\x87\x94\x1f\
\xf0\xf8\xbd\x3a\xca\x96\xfa\x00\xdb\xb8\x25\x07\x2c\x47\x67\xf1\
\x69\xd0\xf2\x11\x68\xff\x75\x74\x4c\x1c\x48\x8f\xee\x7a\x01\x78\
\x52\xcf\x04\x5d\xc2\x9e\xa1\x0e\x09\x63\x76\x18\xc3\x5f\xf6\x10\
",
);

// defers to rand_core::impls, which interpret bytes little-endian.
assert_eq!(seeded.gen::<u32>(), 0xc6a18732);
assert_eq!(seeded.gen::<u64>(), 0x705c00977b0d7be0);

// Test partial consumption, to be sure buffering doesn't cause problems

let mut seeded = crate::Rng::from_seed(*b"0123456789abcdefghijklmnopqrstuv");
let mut buf = [0u8; 63];
seeded.fill_bytes(&mut buf);
// Verified using: printf 0123456789abcdefghijklmnopqrstuv | b3sum -l 76 --keyed <(true)
assert_eq!(
&buf,
b"\
\x57\x63\x36\x95\x85\xc5\x58\x99\x4a\x3e\xe0\x27\x78\x87\x94\x1f\
\xf0\xf8\xbd\x3a\xca\x96\xfa\x00\xdb\xb8\x25\x07\x2c\x47\x67\xf1\
\x69\xd0\xf2\x11\x68\xff\x75\x74\x4c\x1c\x48\x8f\xee\x7a\x01\x78\
\x52\xcf\x04\x5d\xc2\x9e\xa1\x0e\x09\x63\x76\x18\xc3\x5f\xf6\
",
);

// defers to rand_core::impls, which interpret bytes little-endian.
assert_eq!(seeded.gen::<u32>(), 0xa1873210);
assert_eq!(seeded.gen::<u64>(), 0x5c00977b0d7be0c6);
}
Loading