uv_extract/
hash.rs

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

uv_extract/
hash.rs
