snix_store/nar/
hashing_reader.rs

1use std::{
2    io::Result,
3    pin::Pin,
4    task::{Context, Poll, ready},
5};
6
7use md5::{Digest, digest::DynDigest};
8use nix_compat::nixhash::{HashAlgo, NixHash};
9use pin_project_lite::pin_project;
10use tokio::io::{AsyncRead, ReadBuf};
11
12pin_project! {
13    /// AsyncRead implementation with a type-erased hasher.
14    ///
15    /// After it's read the bytes from the underlying reader, it can
16    /// produce the NixHash value corresponding to the Digest that it's been
17    /// constructed with.
18    ///
19    /// Because we are type-erasing the underlying Digest, it uses dynamic dispatch
20    /// and boxing. While it may seem like it could be slow, in practice it's used
21    /// in IO-bound workloads so the slowdown should be negligible.
22    ///
23    /// On the other hand it greatly improves ergonomics of using different hashing
24    /// algorithms and retrieving the corresponding NixHash values.
25    pub struct HashingReader<R> {
26        #[pin]
27        reader: R,
28        digest: Box<dyn ToHash>,
29    }
30}
31
32/// Utility trait that simplifies digesting different hashes.
33///
34/// The main benefit is that each corresponding impl produces its corresponding
35/// NixHash value as opposed to a lower level byte slice.
36trait ToHash: DynDigest + Send {
37    fn consume(self: Box<Self>) -> NixHash;
38}
39
40impl ToHash for sha1::Sha1 {
41    fn consume(self: Box<Self>) -> NixHash {
42        NixHash::Sha1(self.finalize().to_vec().try_into().expect("Snix bug"))
43    }
44}
45
46impl ToHash for sha2::Sha256 {
47    fn consume(self: Box<Self>) -> NixHash {
48        NixHash::Sha256(self.finalize().to_vec().try_into().expect("Snix bug"))
49    }
50}
51
52impl ToHash for sha2::Sha512 {
53    fn consume(self: Box<Self>) -> NixHash {
54        NixHash::Sha512(Box::new(
55            self.finalize().to_vec().try_into().expect("Snix bug"),
56        ))
57    }
58}
59
60impl ToHash for md5::Md5 {
61    fn consume(self: Box<Self>) -> NixHash {
62        NixHash::Md5(self.finalize().to_vec().try_into().expect("Snix bug"))
63    }
64}
65
66impl<R> HashingReader<R> {
67    /// Given a NixHash, creates a HashingReader that uses the same hashing algorithm.
68    pub fn new_with_algo(algo: HashAlgo, reader: R) -> Self {
69        match algo {
70            HashAlgo::Md5 => HashingReader::new::<md5::Md5>(reader),
71            HashAlgo::Sha1 => HashingReader::new::<sha1::Sha1>(reader),
72            HashAlgo::Sha256 => HashingReader::new::<sha2::Sha256>(reader),
73            HashAlgo::Sha512 => HashingReader::new::<sha2::Sha512>(reader),
74        }
75    }
76    fn new<D: ToHash + Digest + 'static>(reader: R) -> Self {
77        HashingReader {
78            reader,
79            digest: Box::new(D::new()),
80        }
81    }
82
83    /// Returns the [`NixHash`] of the data that's been read from this reader.
84    pub fn consume(self) -> NixHash {
85        self.digest.consume()
86    }
87}
88
89impl<R: AsyncRead> AsyncRead for HashingReader<R> {
90    fn poll_read(
91        self: Pin<&mut Self>,
92        cx: &mut Context<'_>,
93        buf: &mut ReadBuf<'_>,
94    ) -> Poll<Result<()>> {
95        let me = self.project();
96        let filled_length = buf.filled().len();
97        ready!(me.reader.poll_read(cx, buf))?;
98        me.digest.update(&buf.filled()[filled_length..]);
99        Poll::Ready(Ok(()))
100    }
101}