oasis_core_runtime/storage/mkvs/tree/
marshal.rs

1use std::{cell::RefCell, mem::size_of, rc::Rc};
2
3use anyhow::Result;
4
5use crate::{
6    common::crypto::hash::Hash,
7    storage::mkvs::{
8        marshal::Marshal,
9        tree::{
10            Depth, DepthTrait, InternalNode, Key, LeafNode, Node, NodeBox, NodeKind, NodePointer,
11            TreeError, Value,
12        },
13    },
14};
15
16/// Size of the encoded value length.
17const VALUE_LENGTH_SIZE: usize = size_of::<u32>();
18
19impl NodeBox {
20    pub fn compact_marshal_binary(&self, v: u16) -> Result<Vec<u8>> {
21        match self {
22            NodeBox::Internal(ref n) => n.compact_marshal_binary(v),
23            NodeBox::Leaf(ref n) => n.compact_marshal_binary(),
24        }
25    }
26}
27
28impl InternalNode {
29    pub fn compact_marshal_binary(&self, v: u16) -> Result<Vec<u8>> {
30        match v {
31            0 => {
32                // In version 0 of compact serialization, leaf node is included in the internal
33                // nodes serialization.
34                let leaf_node_binary = if self.leaf_node.borrow().is_null() {
35                    vec![NodeKind::None as u8]
36                } else {
37                    noderef_as!(self.leaf_node.borrow().get_node(), Leaf).marshal_binary()?
38                };
39
40                let mut result: Vec<u8> = Vec::with_capacity(
41                    1 + size_of::<u16>() + self.label.len() + leaf_node_binary.len(),
42                );
43                result.push(NodeKind::Internal as u8);
44                result.append(&mut self.label_bit_length.marshal_binary()?);
45                result.extend_from_slice(&self.label);
46                result.extend_from_slice(leaf_node_binary.as_ref());
47
48                Ok(result)
49            }
50            1 => {
51                // In version 1 of compact serialization, leaf node is not included.
52                let mut result: Vec<u8> =
53                    Vec::with_capacity(1 + size_of::<u16>() + self.label.len() + 1);
54                result.push(NodeKind::Internal as u8);
55                result.append(&mut self.label_bit_length.marshal_binary()?);
56                result.extend_from_slice(&self.label);
57                result.push(NodeKind::None as u8);
58
59                Ok(result)
60            }
61            _ => panic!("unsupported compact serialization version: {:?}", v),
62        }
63    }
64}
65
66impl Marshal for NodeBox {
67    fn marshal_binary(&self) -> Result<Vec<u8>> {
68        match self {
69            NodeBox::Internal(ref n) => n.marshal_binary(),
70            NodeBox::Leaf(ref n) => n.marshal_binary(),
71        }
72    }
73
74    fn unmarshal_binary(&mut self, data: &[u8]) -> Result<usize> {
75        if data.is_empty() {
76            Err(TreeError::MalformedNode.into())
77        } else {
78            let mut kind = NodeKind::None;
79            kind.unmarshal_binary(data)?;
80            match kind {
81                NodeKind::Internal => {
82                    *self = NodeBox::Internal(InternalNode {
83                        ..Default::default()
84                    });
85                }
86                NodeKind::Leaf => {
87                    *self = NodeBox::Leaf(LeafNode {
88                        ..Default::default()
89                    });
90                }
91                _ => {
92                    return Err(TreeError::MalformedNode.into());
93                }
94            };
95            match self {
96                NodeBox::Internal(ref mut n) => n.unmarshal_binary(data),
97                NodeBox::Leaf(ref mut n) => n.unmarshal_binary(data),
98            }
99        }
100    }
101}
102
103impl Marshal for NodeKind {
104    fn marshal_binary(&self) -> Result<Vec<u8>> {
105        Ok(vec![*self as u8])
106    }
107
108    fn unmarshal_binary(&mut self, data: &[u8]) -> Result<usize> {
109        if data.is_empty() {
110            Err(TreeError::MalformedNode.into())
111        } else {
112            if data[0] == NodeKind::None as u8 {
113                *self = NodeKind::None;
114            } else if data[0] == NodeKind::Internal as u8 {
115                *self = NodeKind::Internal;
116            } else if data[0] == NodeKind::Leaf as u8 {
117                *self = NodeKind::Leaf;
118            } else {
119                return Err(TreeError::MalformedNode.into());
120            }
121            Ok(1)
122        }
123    }
124}
125
126impl Marshal for InternalNode {
127    fn marshal_binary(&self) -> Result<Vec<u8>> {
128        let leaf_node_binary = if self.leaf_node.borrow().is_null() {
129            vec![NodeKind::None as u8]
130        } else {
131            noderef_as!(self.leaf_node.borrow().get_node(), Leaf).marshal_binary()?
132        };
133
134        let mut result: Vec<u8> = Vec::with_capacity(
135            1 + size_of::<u16>() + self.label.len() + leaf_node_binary.len() + Hash::len() * 2,
136        );
137        result.push(NodeKind::Internal as u8);
138        result.append(&mut self.label_bit_length.marshal_binary()?);
139        result.extend_from_slice(&self.label);
140        result.extend_from_slice(leaf_node_binary.as_ref());
141        result.extend_from_slice(self.left.borrow().hash.as_ref());
142        result.extend_from_slice(self.right.borrow().hash.as_ref());
143
144        Ok(result)
145    }
146
147    fn unmarshal_binary(&mut self, data: &[u8]) -> Result<usize> {
148        let mut pos = 0;
149        if data.len() < 1 + size_of::<Depth>() + 1 || data[pos] != NodeKind::Internal as u8 {
150            return Err(TreeError::MalformedNode.into());
151        }
152        pos += 1;
153
154        pos += self.label_bit_length.unmarshal_binary(&data[pos..])?;
155        self.label = vec![0; self.label_bit_length.to_bytes()];
156        if pos + self.label_bit_length.to_bytes() > data.len() {
157            return Err(TreeError::MalformedNode.into());
158        }
159        self.label
160            .clone_from_slice(&data[pos..pos + self.label_bit_length.to_bytes()]);
161        pos += self.label_bit_length.to_bytes();
162        if pos >= data.len() {
163            return Err(TreeError::MalformedNode.into());
164        }
165
166        if data[pos] == NodeKind::None as u8 {
167            self.leaf_node = NodePointer::null_ptr();
168            pos += 1;
169        } else {
170            let mut leaf_node = LeafNode {
171                ..Default::default()
172            };
173            pos += leaf_node.unmarshal_binary(&data[pos..])?;
174            self.leaf_node = Rc::new(RefCell::new(NodePointer {
175                clean: true,
176                hash: leaf_node.get_hash(),
177                node: Some(Rc::new(RefCell::new(NodeBox::Leaf(leaf_node)))),
178                ..Default::default()
179            }));
180        };
181
182        // Hashes are only present in non-compact serialization.
183        if data.len() >= pos + Hash::len() * 2 {
184            let left_hash = Hash::from(&data[pos..pos + Hash::len()]);
185            pos += Hash::len();
186            let right_hash = Hash::from(&data[pos..pos + Hash::len()]);
187            pos += Hash::len();
188
189            if left_hash.is_empty() {
190                self.left = NodePointer::null_ptr();
191            } else {
192                self.left = Rc::new(RefCell::new(NodePointer {
193                    clean: true,
194                    hash: left_hash,
195                    node: None,
196                    ..Default::default()
197                }));
198            }
199            if right_hash.is_empty() {
200                self.right = NodePointer::null_ptr();
201            } else {
202                self.right = Rc::new(RefCell::new(NodePointer {
203                    clean: true,
204                    hash: right_hash,
205                    node: None,
206                    ..Default::default()
207                }));
208            }
209
210            self.update_hash();
211        }
212
213        self.clean = true;
214
215        Ok(pos)
216    }
217}
218
219impl LeafNode {
220    pub fn compact_marshal_binary(&self) -> Result<Vec<u8>> {
221        let mut result: Vec<u8> = Vec::with_capacity(1 + VALUE_LENGTH_SIZE);
222        result.push(NodeKind::Leaf as u8);
223        result.append(&mut self.key.marshal_binary()?);
224        result.append(&mut (self.value.len() as u32).marshal_binary()?);
225        result.extend_from_slice(&self.value);
226
227        Ok(result)
228    }
229}
230
231impl Marshal for LeafNode {
232    fn marshal_binary(&self) -> Result<Vec<u8>> {
233        self.compact_marshal_binary()
234    }
235
236    fn unmarshal_binary(&mut self, data: &[u8]) -> Result<usize> {
237        if data.len() < 1 + size_of::<Depth>() + VALUE_LENGTH_SIZE
238            || data[0] != NodeKind::Leaf as u8
239        {
240            return Err(TreeError::MalformedNode.into());
241        }
242
243        self.clean = true;
244
245        let mut pos = 1;
246        self.key = Key::new();
247        let key_len = self.key.unmarshal_binary(&data[pos..])?;
248        pos += key_len;
249        if pos + VALUE_LENGTH_SIZE > data.len() {
250            return Err(TreeError::MalformedNode.into());
251        }
252
253        self.value = Value::new();
254        let mut value_len = 0u32;
255        value_len.unmarshal_binary(&data[pos..(pos + VALUE_LENGTH_SIZE)])?;
256        pos += VALUE_LENGTH_SIZE;
257        if pos + (value_len as usize) > data.len() {
258            return Err(TreeError::MalformedNode.into());
259        }
260
261        self.value
262            .extend_from_slice(&data[pos..(pos + value_len as usize)]);
263        pos += value_len as usize;
264
265        self.update_hash();
266
267        Ok(pos)
268    }
269}
270
271impl Marshal for Key {
272    fn marshal_binary(&self) -> Result<Vec<u8>> {
273        let mut result: Vec<u8> = Vec::new();
274        result.append(&mut (self.len() as Depth).marshal_binary()?);
275        result.extend_from_slice(self);
276        Ok(result)
277    }
278
279    fn unmarshal_binary(&mut self, data: &[u8]) -> Result<usize> {
280        if data.len() < size_of::<Depth>() {
281            return Err(TreeError::MalformedKey.into());
282        }
283        let mut key_len: Depth = 0;
284        key_len.unmarshal_binary(data)?;
285
286        if data.len() < size_of::<Depth>() + key_len as usize {
287            return Err(TreeError::MalformedKey.into());
288        }
289
290        self.extend_from_slice(&data[size_of::<Depth>()..(size_of::<Depth>() + key_len as usize)]);
291        Ok(size_of::<Depth>() + key_len as usize)
292    }
293}