1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

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());
    }
}