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
90
91
92
93
94
95
96
97
98
99
100
101

use std::{
    io::Result,
    pin::Pin,
    task::{ready, Context, Poll},
};

use md5::{digest::DynDigest, Digest};
use nix_compat::nixhash::{HashAlgo, NixHash};
use pin_project_lite::pin_project;
use tokio::io::{AsyncRead, ReadBuf};

pin_project! {
    /// AsyncRead implementation with a type-erased hasher.
    ///
    /// After it's read the bytes from the underlying reader, it can
    /// produce the NixHash value corresponding to the Digest that it's been
    /// constructed with.
    ///
    /// Because we are type-erasing the underlying Digest, it uses dynamic dispatch
    /// and boxing. While it may seem like it could be slow, in practice it's used
    /// in IO-bound workloads so the slowdown should be negligible.
    ///
    /// On the other hand it greatly improves ergonomics of using different hashing
    /// algorithms and retrieving the corresponding NixHash values.
    pub struct HashingReader<R> {
        #[pin]
        reader: R,
        digest: Box<dyn ToHash>,
    }
}

/// Utility trait that simplifies digesting different hashes.
///
/// The main benefit is that each corresponding impl produces its corresponding
/// NixHash value as opposed to a lower level byte slice.
trait ToHash: DynDigest + Send {
    fn consume(self: Box<Self>) -> NixHash;
}

impl ToHash for sha1::Sha1 {
    fn consume(self: Box<Self>) -> NixHash {
        NixHash::Sha1(self.finalize().to_vec().try_into().expect("Tvix bug"))
    }
}

impl ToHash for sha2::Sha256 {
    fn consume(self: Box<Self>) -> NixHash {
        NixHash::Sha256(self.finalize().to_vec().try_into().expect("Tvix bug"))
    }
}

impl ToHash for sha2::Sha512 {
    fn consume(self: Box<Self>) -> NixHash {
        NixHash::Sha512(Box::new(
            self.finalize().to_vec().try_into().expect("Tvix bug"),
        ))
    }
}

impl ToHash for md5::Md5 {
    fn consume(self: Box<Self>) -> NixHash {
        NixHash::Md5(self.finalize().to_vec().try_into().expect("Tvix bug"))
    }
}

impl<R> HashingReader<R> {
    /// Given a NixHash, creates a HashingReader that uses the same hashing algorithm.
    pub fn new_with_algo(algo: HashAlgo, reader: R) -> Self {
        match algo {
            HashAlgo::Md5 => HashingReader::new::<md5::Md5>(reader),
            HashAlgo::Sha1 => HashingReader::new::<sha1::Sha1>(reader),
            HashAlgo::Sha256 => HashingReader::new::<sha2::Sha256>(reader),
            HashAlgo::Sha512 => HashingReader::new::<sha2::Sha512>(reader),
        }
    }
    fn new<D: ToHash + Digest + 'static>(reader: R) -> Self {
        HashingReader {
            reader,
            digest: Box::new(D::new()),
        }
    }

    /// Returns the [`NixHash`] of the data that's been read from this reader.
    pub fn consume(self) -> NixHash {
        self.digest.consume()
    }
}

impl<R: AsyncRead> AsyncRead for HashingReader<R> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<()>> {
        let me = self.project();
        let filled_length = buf.filled().len();
        ready!(me.reader.poll_read(cx, buf))?;
        me.digest.update(&buf.filled()[filled_length..]);
        Poll::Ready(Ok(()))
    }
}