aws_tree_hash/
lib.rs

1use sha2::{Digest, Sha256};
2
3const BLOCK_SIZE: usize = 1024 * 1024;
4
5/// Calculate the Amazon SHA256 tree hash as described on [Checksum Calculation page](https://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html) of the AWS S3 Glacier developer guide.
6/// The function is meant to be used on the complete data in the case of an upload in a single request and on each part in case of a multi-part upload.
7///
8/// # Example
9/// ```rust
10/// # use hex::FromHex;
11/// #
12/// # fn main() {
13/// assert_eq!(
14///     aws_tree_hash::calculate_tree_hash(&[0; 7680000]),
15///     Vec::from_hex("7a43777ddc7a0326d36b15bc482e6c7736e1c2e9d80a647e8c301646f6a4785c").unwrap()
16/// );
17/// # }
18/// ```
19pub fn calculate_tree_hash(data: &[u8]) -> Vec<u8> {
20    let mut start = 0;
21    let mut end = start + BLOCK_SIZE;
22    let mut hashes = Vec::new();
23
24    while end < data.len() {
25        let mut hasher = Sha256::new();
26
27        hasher.update(&data[start..end]);
28        hashes.push(hasher.finalize().to_vec());
29
30        start = end;
31        end += BLOCK_SIZE;
32    }
33
34    let mut hasher = Sha256::new();
35
36    hasher.update(&data[start..data.len()]);
37    hashes.push(hasher.finalize().to_vec());
38
39    combine_hashes(hashes)
40}
41
42/// Combine the tree hashes from multiple parts (i.e. multiple invocations of [`calculate_tree_hash`]) into the overall tree hash.
43///
44/// [`calculate_tree_hash`]: fn.calculate_tree_hash.html
45///
46/// # Example
47/// ```rust
48/// # use hex::FromHex;
49/// #
50/// # fn main() {
51/// let data = &[0u8; 7680000];
52/// let tree_hash_1 = aws_tree_hash::calculate_tree_hash(&data[..2097152]);
53/// let tree_hash_2 = aws_tree_hash::calculate_tree_hash(&data[2097152..4194304]);
54/// let tree_hash_3 = aws_tree_hash::calculate_tree_hash(&data[4194304..6291456]);
55/// let tree_hash_4 = aws_tree_hash::calculate_tree_hash(&data[6291456..]);
56///
57/// let res = aws_tree_hash::combine_hashes(vec![tree_hash_1, tree_hash_2, tree_hash_3, tree_hash_4]);
58///
59/// assert_eq!(
60///     res,
61///     Vec::from_hex("7a43777ddc7a0326d36b15bc482e6c7736e1c2e9d80a647e8c301646f6a4785c")
62///         .unwrap()
63/// );
64/// # }
65/// ```
66pub fn combine_hashes(hashes: Vec<Vec<u8>>) -> Vec<u8> {
67    let mut res = combine_hashes_once(hashes);
68
69    while res.len() > 1 {
70        res = combine_hashes_once(res);
71    }
72
73    res[0].clone()
74}
75
76/// Combine multiple tree hashes one time.
77/// The length of the result vector (n_res) relates to the length tof the `hashes` vector (n_hashes) as follows: n_res = n_hashes / 2 + n_hashes % 2).
78fn combine_hashes_once(hashes: Vec<Vec<u8>>) -> Vec<Vec<u8>> {
79    let mut iter = hashes.iter();
80    let mut vec = Vec::new();
81
82    loop {
83        match (iter.next(), iter.next()) {
84            (Some(e1), Some(e2)) => {
85                let mut combined = e1.to_owned();
86
87                combined.extend(e2);
88
89                let mut hasher = Sha256::new();
90
91                hasher.update(combined);
92                vec.push(hasher.finalize().to_vec());
93            }
94            (Some(e1), None) => vec.push(e1.to_owned()),
95            (_, _) => break,
96        }
97    }
98
99    vec
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use hex::FromHex;
106
107    #[test]
108    fn null_data() {
109        assert_eq!(
110            calculate_tree_hash(&[0; 7680000]),
111            Vec::from_hex("7a43777ddc7a0326d36b15bc482e6c7736e1c2e9d80a647e8c301646f6a4785c")
112                .unwrap()
113        );
114    }
115
116    #[test]
117    fn combine() {
118        let data = &[0u8; 7680000];
119        let tree_hash_1 = calculate_tree_hash(&data[..2097152]);
120        let tree_hash_2 = calculate_tree_hash(&data[2097152..4194304]);
121        let tree_hash_3 = calculate_tree_hash(&data[4194304..6291456]);
122        let tree_hash_4 = calculate_tree_hash(&data[6291456..]);
123
124        let res = combine_hashes(vec![tree_hash_1, tree_hash_2, tree_hash_3, tree_hash_4]);
125
126        assert_eq!(
127            res,
128            Vec::from_hex("7a43777ddc7a0326d36b15bc482e6c7736e1c2e9d80a647e8c301646f6a4785c")
129                .unwrap()
130        );
131    }
132}