nydus_utils/
digest.rs

1// Copyright 2020 Ant Group. All rights reserved.
2// Copyright (C) 2020-2021 Alibaba Cloud. All rights reserved.
3//
4// SPDX-License-Identifier: Apache-2.0
5
6//! Fast message digest algorithms for Rafs and Nydus, including Blake3 and SHA256.
7
8use std::convert::TryFrom;
9use std::fmt;
10use std::io::{Error, Read};
11use std::str::FromStr;
12
13use sha2::digest::Digest;
14use sha2::Sha256;
15
16/// Size in bytes of chunk digest value.
17pub const RAFS_DIGEST_LENGTH: usize = 32;
18
19/// Type alias for digest data.
20pub type DigestData = [u8; RAFS_DIGEST_LENGTH];
21
22#[repr(u32)]
23#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
24pub enum Algorithm {
25    #[default]
26    Blake3 = 0,
27    Sha256 = 1,
28}
29
30impl fmt::Display for Algorithm {
31    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
32        write!(f, "{:?}", self)
33    }
34}
35
36impl FromStr for Algorithm {
37    type Err = Error;
38
39    fn from_str(s: &str) -> Result<Self, Self::Err> {
40        match s {
41            "blake3" => Ok(Self::Blake3),
42            "sha256" => Ok(Self::Sha256),
43            _ => Err(einval!("digest algorithm should be blake3 or sha256")),
44        }
45    }
46}
47
48impl TryFrom<u32> for Algorithm {
49    type Error = ();
50
51    fn try_from(value: u32) -> Result<Self, Self::Error> {
52        if value == Algorithm::Sha256 as u32 {
53            Ok(Algorithm::Sha256)
54        } else if value == Algorithm::Blake3 as u32 {
55            Ok(Algorithm::Blake3)
56        } else {
57            Err(())
58        }
59    }
60}
61
62impl TryFrom<u64> for Algorithm {
63    type Error = ();
64
65    fn try_from(value: u64) -> Result<Self, Self::Error> {
66        if value == Algorithm::Sha256 as u64 {
67            Ok(Algorithm::Sha256)
68        } else if value == Algorithm::Blake3 as u64 {
69            Ok(Algorithm::Blake3)
70        } else {
71            Err(())
72        }
73    }
74}
75
76pub trait DigestHasher {
77    fn digest_update(&mut self, buf: &[u8]);
78    fn digest_finalize(self) -> RafsDigest;
79}
80
81/// Fast message digest algorithm.
82///
83/// The size of Hasher struct is a little big, say
84/// blake3::Hasher: 1912 bytes
85/// Sha256: 112 bytes
86/// RafsDigestHasher: 1920
87///
88/// So we should avoid any unnecessary clone() operation. Add we prefer allocation on stack
89/// instead of allocation on heap.
90///
91/// If allocating memory for blake3::Hasher is preferred over using the stack, please try:
92/// Blake3(Box<blake3::Hasher>). But be careful, this will cause one extra memory allocation/free
93/// for each digest.
94#[derive(Clone, Debug)]
95pub enum RafsDigestHasher {
96    Blake3(Box<blake3::Hasher>),
97    Sha256(Sha256),
98}
99
100impl DigestHasher for RafsDigestHasher {
101    fn digest_update(&mut self, buf: &[u8]) {
102        match self {
103            RafsDigestHasher::Blake3(hasher) => {
104                hasher.update(buf);
105            }
106            RafsDigestHasher::Sha256(hasher) => {
107                hasher.update(buf);
108            }
109        }
110    }
111
112    fn digest_finalize(self) -> RafsDigest {
113        let data = match self {
114            RafsDigestHasher::Blake3(hasher) => hasher.finalize().into(),
115            RafsDigestHasher::Sha256(hasher) => hasher.finalize().into(),
116        };
117
118        RafsDigest { data }
119    }
120}
121
122impl DigestHasher for blake3::Hasher {
123    fn digest_update(&mut self, buf: &[u8]) {
124        self.update(buf);
125    }
126
127    fn digest_finalize(self) -> RafsDigest {
128        RafsDigest {
129            data: self.finalize().into(),
130        }
131    }
132}
133
134impl DigestHasher for Sha256 {
135    fn digest_update(&mut self, buf: &[u8]) {
136        self.update(buf);
137    }
138
139    fn digest_finalize(self) -> RafsDigest {
140        RafsDigest {
141            data: self.finalize().into(),
142        }
143    }
144}
145
146#[repr(C)]
147#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug, Default, Ord, PartialOrd)]
148pub struct RafsDigest {
149    pub data: DigestData,
150}
151
152impl RafsDigest {
153    pub fn from_buf(buf: &[u8], algorithm: Algorithm) -> Self {
154        let data: DigestData = match algorithm {
155            Algorithm::Blake3 => blake3::hash(buf).into(),
156            Algorithm::Sha256 => {
157                let mut hasher = Sha256::new();
158                hasher.update(buf);
159                hasher.finalize().into()
160            }
161        };
162
163        RafsDigest { data }
164    }
165
166    /// Compute message digest with the given algorithm by read data from the reader.
167    pub fn from_reader<R: Read>(reader: &mut R, algorithm: Algorithm) -> std::io::Result<Self> {
168        let mut digester = Self::hasher(algorithm);
169        let mut buf = vec![0u8; 16384];
170        loop {
171            let sz = reader.read(&mut buf)?;
172            if sz == 0 {
173                return Ok(digester.digest_finalize());
174            }
175            digester.digest_update(&buf[..sz]);
176        }
177    }
178
179    /// According to the format of sha256.
180    pub fn from_string(input: &str) -> Self {
181        let mut digest = RafsDigest::default();
182
183        for (i, byte) in input.as_bytes().chunks(2).enumerate() {
184            let hex_str = std::str::from_utf8(byte).unwrap();
185            digest.data[i] = u8::from_str_radix(hex_str, 16).unwrap();
186        }
187
188        digest
189    }
190
191    pub fn hasher(algorithm: Algorithm) -> RafsDigestHasher {
192        match algorithm {
193            Algorithm::Blake3 => RafsDigestHasher::Blake3(Box::new(blake3::Hasher::new())),
194            Algorithm::Sha256 => RafsDigestHasher::Sha256(Sha256::new()),
195        }
196    }
197}
198
199impl From<DigestData> for RafsDigest {
200    fn from(data: DigestData) -> Self {
201        Self { data }
202    }
203}
204
205impl From<&DigestData> for &RafsDigest {
206    fn from(data: &DigestData) -> Self {
207        unsafe { &*(data as *const DigestData as *const u8 as *const RafsDigest) }
208    }
209}
210
211impl AsRef<[u8]> for RafsDigest {
212    fn as_ref(&self) -> &[u8] {
213        &self.data
214    }
215}
216
217impl fmt::Display for RafsDigest {
218    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
219        for c in &self.data {
220            write!(f, "{:02x}", c)?;
221        }
222        Ok(())
223    }
224}
225
226impl From<RafsDigest> for String {
227    fn from(d: RafsDigest) -> Self {
228        format!("{}", d)
229    }
230}
231
232#[cfg(test)]
233mod test {
234    use super::*;
235
236    #[test]
237    fn test_algorithm() {
238        assert_eq!(Algorithm::from_str("blake3").unwrap(), Algorithm::Blake3);
239        assert_eq!(Algorithm::from_str("sha256").unwrap(), Algorithm::Sha256);
240        Algorithm::from_str("Blake3").unwrap_err();
241        Algorithm::from_str("SHA256").unwrap_err();
242    }
243
244    #[test]
245    fn test_hash_from_buf() {
246        let text = b"The quick brown fox jumps over the lazy dog";
247
248        let blake3 = RafsDigest::from_buf(text, Algorithm::Blake3);
249        let str: String = blake3.into();
250        assert_eq!(
251            str.as_bytes(),
252            b"2f1514181aadccd913abd94cfa592701a5686ab23f8df1dff1b74710febc6d4a"
253        );
254
255        let sha256 = RafsDigest::from_buf(text, Algorithm::Sha256);
256        let str: String = sha256.into();
257        assert_eq!(
258            str.as_bytes(),
259            b"d7a8fbb307d7809469ca9abcb0082e4f8d5651e46d3cdb762d02d0bf37c9e592"
260        );
261    }
262
263    #[test]
264    fn test_hasher() {
265        let text = b"The quick brown fox jumps ";
266        let text2 = b"over the lazy dog";
267
268        let mut hasher = RafsDigest::hasher(Algorithm::Blake3);
269        hasher.digest_update(text);
270        hasher.digest_update(text2);
271        let blake3 = hasher.digest_finalize();
272        let str: String = blake3.into();
273        assert_eq!(
274            str.as_bytes(),
275            b"2f1514181aadccd913abd94cfa592701a5686ab23f8df1dff1b74710febc6d4a"
276        );
277
278        let mut hasher = RafsDigest::hasher(Algorithm::Sha256);
279        hasher.digest_update(text);
280        hasher.digest_update(text2);
281        let sha256 = hasher.digest_finalize();
282        let str: String = sha256.into();
283        assert_eq!(
284            str.as_bytes(),
285            b"d7a8fbb307d7809469ca9abcb0082e4f8d5651e46d3cdb762d02d0bf37c9e592"
286        );
287    }
288
289    #[test]
290    fn test_try_from() {
291        assert!(Algorithm::try_from(Algorithm::Sha256 as u32).is_ok());
292        assert!(Algorithm::try_from(Algorithm::Blake3 as u32).is_ok());
293        assert!(Algorithm::try_from(0xffff_abcd as u32).is_err());
294
295        assert!(Algorithm::try_from(Algorithm::Sha256 as u64).is_ok());
296        assert!(Algorithm::try_from(Algorithm::Blake3 as u64).is_ok());
297        assert!(Algorithm::try_from(0xffff_abcd as u64).is_err());
298    }
299
300    #[test]
301    fn test_spec_hasher_new() {
302        let text = b"The quick brown fox jumps ";
303        let text2 = b"over the lazy dog";
304
305        let mut hasher: blake3::Hasher = blake3::Hasher::new();
306        hasher.digest_update(text);
307        hasher.digest_update(text2);
308        let blake3 = hasher.digest_finalize();
309        let str: String = blake3.into();
310        assert_eq!(
311            str.as_bytes(),
312            b"2f1514181aadccd913abd94cfa592701a5686ab23f8df1dff1b74710febc6d4a"
313        );
314
315        let mut hasher = RafsDigestHasher::Sha256(Sha256::new());
316        hasher.digest_update(text);
317        hasher.digest_update(text2);
318        let sha256 = hasher.digest_finalize();
319        let str: String = sha256.into();
320        assert_eq!(
321            str.as_bytes(),
322            b"d7a8fbb307d7809469ca9abcb0082e4f8d5651e46d3cdb762d02d0bf37c9e592"
323        );
324    }
325
326    #[test]
327    fn test_rafs_digest_try_from() {
328        let text = b"The quick brown fox jumps over the lazy dog";
329
330        let d1 = RafsDigest::from_buf(text, Algorithm::Blake3);
331        let d2 = RafsDigest::from(d1.data);
332        let s1: String = d1.into();
333        let s2: String = d2.into();
334        print!("{:?}", d1);
335        assert_eq!(s1, s2);
336        print!("{:?}, {:?}", Algorithm::Blake3, Algorithm::Sha256);
337    }
338}
nydus_utils/digest.rs

nydus_utils/
digest.rs
