71fb99a265
Change-Id: I99ae0e27b4db4799db8af7cd6b9cc8d7f09227de Reviewed-on: https://cl.tvl.fyi/c/depot/+/11366 Tested-by: BuildkiteCI Reviewed-by: picnoir picnoir <picnoir@alternativebit.fr>
89 lines
2.1 KiB
Rust
89 lines
2.1 KiB
Rust
use pin_project_lite::pin_project;
|
|
use tokio::io::AsyncRead;
|
|
|
|
pin_project! {
|
|
/// Wraps an existing AsyncRead, and allows querying for the digest of all
|
|
/// data read "through" it.
|
|
/// The hash function is configurable by type parameter.
|
|
pub struct HashingReader<R, H>
|
|
where
|
|
R: AsyncRead,
|
|
H: digest::Digest,
|
|
{
|
|
#[pin]
|
|
inner: R,
|
|
hasher: H,
|
|
}
|
|
}
|
|
|
|
pub type B3HashingReader<R> = HashingReader<R, blake3::Hasher>;
|
|
|
|
impl<R, H> HashingReader<R, H>
|
|
where
|
|
R: AsyncRead,
|
|
H: digest::Digest,
|
|
{
|
|
pub fn from(r: R) -> Self {
|
|
Self {
|
|
inner: r,
|
|
hasher: H::new(),
|
|
}
|
|
}
|
|
|
|
/// Return the digest.
|
|
pub fn digest(self) -> digest::Output<H> {
|
|
self.hasher.finalize()
|
|
}
|
|
}
|
|
|
|
impl<R, H> tokio::io::AsyncRead for HashingReader<R, H>
|
|
where
|
|
R: AsyncRead,
|
|
H: digest::Digest,
|
|
{
|
|
fn poll_read(
|
|
self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
buf: &mut tokio::io::ReadBuf<'_>,
|
|
) -> std::task::Poll<std::io::Result<()>> {
|
|
let buf_filled_len_before = buf.filled().len();
|
|
|
|
let this = self.project();
|
|
let ret = this.inner.poll_read(cx, buf);
|
|
|
|
// write everything new filled into the hasher.
|
|
this.hasher.update(&buf.filled()[buf_filled_len_before..]);
|
|
|
|
ret
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::io::Cursor;
|
|
|
|
use rstest::rstest;
|
|
|
|
use crate::fixtures::BLOB_A;
|
|
use crate::fixtures::BLOB_A_DIGEST;
|
|
use crate::fixtures::BLOB_B;
|
|
use crate::fixtures::BLOB_B_DIGEST;
|
|
use crate::fixtures::EMPTY_BLOB_DIGEST;
|
|
use crate::{B3Digest, B3HashingReader};
|
|
|
|
#[rstest]
|
|
#[case::blob_a(&BLOB_A, &BLOB_A_DIGEST)]
|
|
#[case::blob_b(&BLOB_B, &BLOB_B_DIGEST)]
|
|
#[case::empty_blob(&[], &EMPTY_BLOB_DIGEST)]
|
|
#[tokio::test]
|
|
async fn test_b3_hashing_reader(#[case] data: &[u8], #[case] b3_digest: &B3Digest) {
|
|
let r = Cursor::new(data);
|
|
let mut hr = B3HashingReader::from(r);
|
|
|
|
tokio::io::copy(&mut hr, &mut tokio::io::sink())
|
|
.await
|
|
.expect("read must succeed");
|
|
|
|
assert_eq!(*b3_digest, hr.digest().into());
|
|
}
|
|
}
|