uv_extract/
hash.rs

1use blake2::digest::consts::U32;
2use sha2::Digest;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5use tokio::io::{AsyncReadExt, ReadBuf};
6
7use uv_pypi_types::{HashAlgorithm, HashDigest};
8
9#[derive(Debug)]
10pub enum Hasher {
11    Md5(md5::Md5),
12    Sha256(sha2::Sha256),
13    Sha384(sha2::Sha384),
14    Sha512(sha2::Sha512),
15    Blake2b(blake2::Blake2b<U32>),
16}
17
18impl Hasher {
19    pub fn update(&mut self, data: &[u8]) {
20        match self {
21            Self::Md5(hasher) => hasher.update(data),
22            Self::Sha256(hasher) => hasher.update(data),
23            Self::Sha384(hasher) => hasher.update(data),
24            Self::Sha512(hasher) => hasher.update(data),
25            Self::Blake2b(hasher) => hasher.update(data),
26        }
27    }
28}
29
30impl From<HashAlgorithm> for Hasher {
31    fn from(algorithm: HashAlgorithm) -> Self {
32        match algorithm {
33            HashAlgorithm::Md5 => Self::Md5(md5::Md5::new()),
34            HashAlgorithm::Sha256 => Self::Sha256(sha2::Sha256::new()),
35            HashAlgorithm::Sha384 => Self::Sha384(sha2::Sha384::new()),
36            HashAlgorithm::Sha512 => Self::Sha512(sha2::Sha512::new()),
37            HashAlgorithm::Blake2b => Self::Blake2b(blake2::Blake2b::new()),
38        }
39    }
40}
41
42impl From<Hasher> for HashDigest {
43    fn from(hasher: Hasher) -> Self {
44        match hasher {
45            Hasher::Md5(hasher) => Self {
46                algorithm: HashAlgorithm::Md5,
47                digest: format!("{:x}", hasher.finalize()).into(),
48            },
49            Hasher::Sha256(hasher) => Self {
50                algorithm: HashAlgorithm::Sha256,
51                digest: format!("{:x}", hasher.finalize()).into(),
52            },
53            Hasher::Sha384(hasher) => Self {
54                algorithm: HashAlgorithm::Sha384,
55                digest: format!("{:x}", hasher.finalize()).into(),
56            },
57            Hasher::Sha512(hasher) => Self {
58                algorithm: HashAlgorithm::Sha512,
59                digest: format!("{:x}", hasher.finalize()).into(),
60            },
61            Hasher::Blake2b(hasher) => Self {
62                algorithm: HashAlgorithm::Blake2b,
63                digest: format!("{:x}", hasher.finalize()).into(),
64            },
65        }
66    }
67}
68
69pub struct HashReader<'a, R> {
70    reader: R,
71    hashers: &'a mut [Hasher],
72}
73
74impl<'a, R> HashReader<'a, R>
75where
76    R: tokio::io::AsyncRead + Unpin,
77{
78    pub fn new(reader: R, hashers: &'a mut [Hasher]) -> Self {
79        HashReader { reader, hashers }
80    }
81
82    /// Exhaust the underlying reader.
83    pub async fn finish(&mut self) -> Result<(), std::io::Error> {
84        while self.read(&mut vec![0; 8192]).await? > 0 {}
85
86        Ok(())
87    }
88}
89
90impl<R> tokio::io::AsyncRead for HashReader<'_, R>
91where
92    R: tokio::io::AsyncRead + Unpin,
93{
94    fn poll_read(
95        mut self: Pin<&mut Self>,
96        cx: &mut Context<'_>,
97        buf: &mut ReadBuf<'_>,
98    ) -> Poll<std::io::Result<()>> {
99        let reader = Pin::new(&mut self.reader);
100        match reader.poll_read(cx, buf) {
101            Poll::Ready(Ok(())) => {
102                for hasher in self.hashers.iter_mut() {
103                    hasher.update(buf.filled());
104                }
105                Poll::Ready(Ok(()))
106            }
107            other => other,
108        }
109    }
110}
uv_extract/hash.rs

uv_extract/
hash.rs
