certified_vars/
hashtree.rs

1// This file is copied from ic-certified-map which was released under Apache V2.
2// Some modifications are made to improve the code quality.
3use serde::{ser::SerializeSeq, Serialize, Serializer};
4use serde_bytes::Bytes;
5use sha2::{Digest, Sha256};
6use std::borrow::Cow;
7
8/// SHA-256 hash bytes.
9pub type Hash = [u8; 32];
10
11#[derive(Debug, Eq, PartialEq)]
12pub struct ForkInner<'a>(pub HashTree<'a>, pub HashTree<'a>);
13
14impl<'a> ForkInner<'a> {
15    pub fn left(&self) -> &HashTree<'a> {
16        &self.0
17    }
18
19    pub fn right(&self) -> &HashTree<'a> {
20        &self.1
21    }
22}
23
24/// HashTree as defined in the interfaces spec.
25/// https://sdk.dfinity.org/docs/interface-spec/index.html#_certificate
26#[derive(Debug, Eq, PartialEq)]
27pub enum HashTree<'a> {
28    Empty,
29    Fork(Box<ForkInner<'a>>),
30    Labeled(Cow<'a, [u8]>, Box<HashTree<'a>>),
31    Leaf(Cow<'a, [u8]>),
32    Pruned(Hash),
33}
34
35pub fn fork<'a>(l: HashTree<'a>, r: HashTree<'a>) -> HashTree<'a> {
36    HashTree::Fork(Box::new(ForkInner(l, r)))
37}
38
39pub fn labeled<'a>(l: &'a [u8], t: HashTree<'a>) -> HashTree<'a> {
40    HashTree::Labeled(Cow::Borrowed(l), Box::new(t))
41}
42
43pub fn fork_hash(l: &Hash, r: &Hash) -> Hash {
44    let mut h = domain_sep("ic-hashtree-fork");
45    h.update(&l[..]);
46    h.update(&r[..]);
47    h.finalize().into()
48}
49
50pub fn leaf_hash(data: &[u8]) -> Hash {
51    let mut h = domain_sep("ic-hashtree-leaf");
52    h.update(data);
53    h.finalize().into()
54}
55
56pub fn labeled_hash(label: &[u8], content_hash: &Hash) -> Hash {
57    let mut h = domain_sep("ic-hashtree-labeled");
58    h.update(label);
59    h.update(&content_hash[..]);
60    h.finalize().into()
61}
62
63impl<'a> HashTree<'a> {
64    pub fn reconstruct(&self) -> Hash {
65        match self {
66            Self::Empty => domain_sep("ic-hashtree-empty").finalize().into(),
67            Self::Fork(f) => fork_hash(&f.0.reconstruct(), &f.1.reconstruct()),
68            Self::Labeled(l, t) => {
69                let thash = t.reconstruct();
70                labeled_hash(l, &thash)
71            }
72            Self::Leaf(data) => leaf_hash(data),
73            Self::Pruned(h) => *h,
74        }
75    }
76
77    /// Collect and return all of the labels in this HashTree.
78    ///
79    /// This method is intended for testing purposes.
80    pub fn get_labels<'b: 'a>(&'b self) -> Vec<&'b [u8]> {
81        fn go<'a>(keys: &mut Vec<&'a [u8]>, tree: &'a HashTree<'a>) {
82            match tree {
83                HashTree::Labeled(key, value) => {
84                    keys.push(key);
85                    go(keys, value);
86                }
87                HashTree::Fork(lr) => {
88                    go(keys, lr.left());
89                    go(keys, lr.right());
90                }
91                _ => (),
92            }
93        }
94
95        let mut keys = Vec::new();
96        go(&mut keys, self);
97        keys
98    }
99
100    /// Collect and return all of the values in this HashTree.
101    ///
102    /// This method is intended for testing purposes.
103    pub fn get_leaf_values<'b: 'a>(&'b self) -> Vec<&'b [u8]> {
104        fn go<'a>(values: &mut Vec<&'a [u8]>, tree: &'a HashTree<'a>) {
105            match tree {
106                HashTree::Leaf(value) => {
107                    values.push(value);
108                }
109                HashTree::Fork(lr) => {
110                    go(values, lr.left());
111                    go(values, lr.right());
112                }
113                HashTree::Labeled(_, t) => {
114                    go(values, &*t);
115                }
116                _ => (),
117            }
118        }
119
120        let mut values = Vec::new();
121        go(&mut values, self);
122        values
123    }
124}
125
126impl Serialize for HashTree<'_> {
127    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
128    where
129        S: Serializer,
130    {
131        match self {
132            HashTree::Empty => {
133                let mut seq = serializer.serialize_seq(Some(1))?;
134                seq.serialize_element(&0u8)?;
135                seq.end()
136            }
137            HashTree::Fork(p) => {
138                let mut seq = serializer.serialize_seq(Some(3))?;
139                seq.serialize_element(&1u8)?;
140                seq.serialize_element(&p.0)?;
141                seq.serialize_element(&p.1)?;
142                seq.end()
143            }
144            HashTree::Labeled(label, tree) => {
145                let mut seq = serializer.serialize_seq(Some(3))?;
146                seq.serialize_element(&2u8)?;
147                seq.serialize_element(Bytes::new(label))?;
148                seq.serialize_element(&tree)?;
149                seq.end()
150            }
151            HashTree::Leaf(leaf_bytes) => {
152                let mut seq = serializer.serialize_seq(Some(2))?;
153                seq.serialize_element(&3u8)?;
154                seq.serialize_element(Bytes::new(leaf_bytes))?;
155                seq.end()
156            }
157            HashTree::Pruned(digest) => {
158                let mut seq = serializer.serialize_seq(Some(2))?;
159                seq.serialize_element(&4u8)?;
160                seq.serialize_element(Bytes::new(&digest[..]))?;
161                seq.end()
162            }
163        }
164    }
165}
166
167fn domain_sep(s: &str) -> sha2::Sha256 {
168    let buf: [u8; 1] = [s.len() as u8];
169    let mut h = Sha256::new();
170    h.update(&buf[..]);
171    h.update(s.as_bytes());
172    h
173}
174
175#[cfg(test)]
176mod tests {
177    use super::{
178        fork, labeled,
179        HashTree::{Empty, Leaf},
180    };
181    use std::borrow::Cow;
182
183    //─┬─┬╴"a" ─┬─┬╴"x" ─╴"hello"
184    // │ │      │ └╴Empty
185    // │ │      └╴  "y" ─╴"world"
186    // │ └╴"b" ──╴"good"
187    // └─┬╴"c" ──╴Empty
188    //   └╴"d" ──╴"morning"
189    #[test]
190    fn test_public_spec_example() {
191        let t = fork(
192            fork(
193                labeled(
194                    b"a",
195                    fork(
196                        fork(labeled(b"x", Leaf(Cow::Borrowed(b"hello"))), Empty),
197                        labeled(b"y", Leaf(Cow::Borrowed(b"world"))),
198                    ),
199                ),
200                labeled(b"b", Leaf(Cow::Borrowed(b"good"))),
201            ),
202            fork(
203                labeled(b"c", Empty),
204                labeled(b"d", Leaf(Cow::Borrowed(b"morning"))),
205            ),
206        );
207
208        assert_eq!(
209            hex::encode(&t.reconstruct()[..]),
210            "eb5c5b2195e62d996b84c9bcc8259d19a83786a2f59e0878cec84c811f669aa0".to_string()
211        );
212
213        assert_eq!(
214            hex::encode(serde_cbor::to_vec(&t).unwrap()),
215            "8301830183024161830183018302417882034568656c6c6f810083024179820345776f726c6483024162820344676f6f648301830241638100830241648203476d6f726e696e67".to_string());
216    }
217}
certified_vars/hashtree.rs

certified_vars/
hashtree.rs
