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}