|
| 1 | +// Copyright 2019 Kodebox, Inc. |
| 2 | +// This file is part of CodeChain. |
| 3 | +// |
| 4 | +// This program is free software: you can redistribute it and/or modify |
| 5 | +// it under the terms of the GNU Affero General Public License as |
| 6 | +// published by the Free Software Foundation, either version 3 of the |
| 7 | +// License, or (at your option) any later version. |
| 8 | +// |
| 9 | +// This program is distributed in the hope that it will be useful, |
| 10 | +// but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 11 | +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 12 | +// GNU Affero General Public License for more details. |
| 13 | +// |
| 14 | +// You should have received a copy of the GNU Affero General Public License |
| 15 | +// along with this program. If not, see <https://www.gnu.org/licenses/>. |
| 16 | + |
| 17 | +use std::collections::VecDeque; |
| 18 | +use std::convert::From; |
| 19 | + |
| 20 | +use ccrypto::BLAKE_NULL_RLP; |
| 21 | +use hashdb::{DBValue, HashDB}; |
| 22 | +use primitives::H256; |
| 23 | + |
| 24 | +use super::error::{ChunkError, Error}; |
| 25 | +use super::{DecodedPathSlice, PathSlice, CHUNK_HEIGHT}; |
| 26 | +use crate::nibbleslice::NibbleSlice; |
| 27 | +use crate::{Node, TrieDBMut}; |
| 28 | + |
| 29 | +#[derive(RlpEncodable, RlpDecodable, Eq, PartialEq)] |
| 30 | +pub struct TerminalNode { |
| 31 | + // Relative path from the chunk root. |
| 32 | + pub path_slice: PathSlice, |
| 33 | + pub node_rlp: Vec<u8>, |
| 34 | +} |
| 35 | + |
| 36 | +impl std::fmt::Debug for TerminalNode { |
| 37 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { |
| 38 | + let path_slice = NibbleSlice::from_encoded(&self.path_slice); |
| 39 | + f.debug_struct("TerminalNode") |
| 40 | + .field("path_slice", &path_slice) |
| 41 | + .field("node_rlp", &NodeDebugAdaptor { |
| 42 | + rlp: &self.node_rlp, |
| 43 | + }) |
| 44 | + .finish() |
| 45 | + } |
| 46 | +} |
| 47 | + |
| 48 | +struct NodeDebugAdaptor<'a> { |
| 49 | + rlp: &'a [u8], |
| 50 | +} |
| 51 | + |
| 52 | +impl<'a> std::fmt::Debug for NodeDebugAdaptor<'a> { |
| 53 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { |
| 54 | + match Node::decoded(&self.rlp) { |
| 55 | + Some(node) => write!(f, "{:?}", &node), |
| 56 | + None => write!(f, "{:?}", self.rlp), |
| 57 | + } |
| 58 | + } |
| 59 | +} |
| 60 | + |
| 61 | +/// An unverified chunk from the network |
| 62 | +#[derive(Debug)] |
| 63 | +pub struct RawChunk { |
| 64 | + pub nodes: Vec<TerminalNode>, |
| 65 | +} |
| 66 | + |
| 67 | +/// Fully recovered, and re-hydrated chunk. |
| 68 | +pub struct RecoveredChunk { |
| 69 | + pub(crate) root: H256, |
| 70 | + /// contains all nodes including non-terminal nodes and terminal nodes. |
| 71 | + /// You can blindly pour all items in `nodes` into `HashDB`. |
| 72 | + pub(crate) nodes: Vec<(H256, DBValue)>, |
| 73 | + /// Their path slices are relative to this chunk root. |
| 74 | + pub(crate) unresolved_chunks: Vec<UnresolvedChunk>, |
| 75 | +} |
| 76 | + |
| 77 | +impl RawChunk { |
| 78 | + /// Verify and recover the chunk |
| 79 | + pub fn recover(&self, expected_chunk_root: H256) -> Result<RecoveredChunk, Error> { |
| 80 | + let mut memorydb = memorydb::MemoryDB::new(); |
| 81 | + let mut chunk_root = H256::new(); |
| 82 | + |
| 83 | + { |
| 84 | + let mut trie = TrieDBMut::new(&mut memorydb, &mut chunk_root); |
| 85 | + for node in self.nodes.iter() { |
| 86 | + let old_val = match Node::decoded(&node.node_rlp) { |
| 87 | + Some(Node::Branch(slice, child)) => { |
| 88 | + let encoded = DecodedPathSlice::from_encoded(&node.path_slice).with_slice(slice).encode(); |
| 89 | + trie.insert_raw(Node::Branch(NibbleSlice::from_encoded(&encoded), child))? |
| 90 | + } |
| 91 | + Some(Node::Leaf(slice, data)) => { |
| 92 | + let encoded = DecodedPathSlice::from_encoded(&node.path_slice).with_slice(slice).encode(); |
| 93 | + trie.insert_raw(Node::Leaf(NibbleSlice::from_encoded(&encoded), data))? |
| 94 | + } |
| 95 | + None => return Err(ChunkError::InvalidContent.into()), |
| 96 | + }; |
| 97 | + |
| 98 | + if let Some(old_val) = old_val { |
| 99 | + if old_val.as_ref() != node.node_rlp.as_slice() { |
| 100 | + return Err(ChunkError::InvalidContent.into()) |
| 101 | + } |
| 102 | + } |
| 103 | + } |
| 104 | + } |
| 105 | + |
| 106 | + // Some nodes in the chunk is different from the expected. |
| 107 | + if chunk_root != expected_chunk_root { |
| 108 | + return Err(ChunkError::ChunkRootMismatch { |
| 109 | + expected: expected_chunk_root, |
| 110 | + actual: chunk_root, |
| 111 | + } |
| 112 | + .into()) |
| 113 | + } |
| 114 | + |
| 115 | + let mut nodes = Vec::new(); |
| 116 | + let mut unresolved_chunks = Vec::new(); |
| 117 | + let mut queue: VecDeque<NodePath> = VecDeque::from(vec![NodePath::new(chunk_root)]); |
| 118 | + while let Some(path) = queue.pop_front() { |
| 119 | + let node = match memorydb.get(&path.key) { |
| 120 | + Some(x) => x, |
| 121 | + None => { |
| 122 | + // all unresolved should depth == CHUNK_HEIGHT + 1 |
| 123 | + if path.depth != CHUNK_HEIGHT + 1 { |
| 124 | + return Err(ChunkError::InvalidHeight.into()) |
| 125 | + } |
| 126 | + |
| 127 | + unresolved_chunks.push(UnresolvedChunk::from(path)); |
| 128 | + continue |
| 129 | + } |
| 130 | + }; |
| 131 | + |
| 132 | + if path.depth > CHUNK_HEIGHT { |
| 133 | + return Err(ChunkError::InvalidHeight.into()) |
| 134 | + } |
| 135 | + nodes.push((path.key, node.clone())); |
| 136 | + |
| 137 | + let node = Node::decoded(&node).expect("Chunk root was verified; Node can't be wrong"); |
| 138 | + if let Node::Branch(slice, children) = node { |
| 139 | + for (index, child) in children.iter().enumerate() { |
| 140 | + if let Some(child) = child { |
| 141 | + queue.push_back(path.with_slice_and_index(slice, index, *child)); |
| 142 | + } |
| 143 | + } |
| 144 | + } |
| 145 | + } |
| 146 | + |
| 147 | + Ok(RecoveredChunk { |
| 148 | + root: expected_chunk_root, |
| 149 | + nodes, |
| 150 | + unresolved_chunks, |
| 151 | + }) |
| 152 | + } |
| 153 | +} |
| 154 | + |
| 155 | +impl std::fmt::Debug for RecoveredChunk { |
| 156 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { |
| 157 | + struct Adapter<'a>(&'a [(H256, DBValue)]); |
| 158 | + impl<'a> std::fmt::Debug for Adapter<'a> { |
| 159 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { |
| 160 | + f.debug_list() |
| 161 | + .entries(self.0.iter().map(|(hash, rlp)| { |
| 162 | + (hash, NodeDebugAdaptor { |
| 163 | + rlp, |
| 164 | + }) |
| 165 | + })) |
| 166 | + .finish() |
| 167 | + } |
| 168 | + } |
| 169 | + |
| 170 | + f.debug_struct("RecoveredChunk") |
| 171 | + .field("root", &self.root) |
| 172 | + .field("nodes", &Adapter(&self.nodes)) |
| 173 | + .field("unresolved_chunks", &self.unresolved_chunks) |
| 174 | + .finish() |
| 175 | + } |
| 176 | +} |
| 177 | + |
| 178 | +/// Chunk obtained from the state db. |
| 179 | +#[derive(Debug)] |
| 180 | +pub struct Chunk { |
| 181 | + pub root: H256, |
| 182 | + pub terminal_nodes: Vec<TerminalNode>, |
| 183 | +} |
| 184 | + |
| 185 | +impl Chunk { |
| 186 | + pub(crate) fn from_chunk_root(db: &dyn HashDB, chunk_root: H256) -> Chunk { |
| 187 | + let mut unresolved: VecDeque<NodePath> = VecDeque::from(vec![NodePath::new(chunk_root)]); |
| 188 | + let mut terminal_nodes: Vec<TerminalNode> = Vec::new(); |
| 189 | + while let Some(path) = unresolved.pop_front() { |
| 190 | + assert!(path.key != BLAKE_NULL_RLP, "Empty DB"); |
| 191 | + assert!(path.depth <= CHUNK_HEIGHT); |
| 192 | + let node = db.get(&path.key).expect("Can't find the node in a db. DB is inconsistent"); |
| 193 | + let node_decoded = Node::decoded(&node).expect("Node cannot be decoded. DB is inconsistent"); |
| 194 | + |
| 195 | + match node_decoded { |
| 196 | + // Continue to BFS |
| 197 | + Node::Branch(slice, ref children) if path.depth < CHUNK_HEIGHT => { |
| 198 | + for (i, hash) in children.iter().enumerate() { |
| 199 | + if let Some(hash) = hash { |
| 200 | + unresolved.push_back(path.with_slice_and_index(slice, i, *hash)); |
| 201 | + } |
| 202 | + } |
| 203 | + } |
| 204 | + // Reached the terminal node. Branch at path.depth == CHUNK_HEIGHT || Leaf |
| 205 | + _ => terminal_nodes.push(TerminalNode { |
| 206 | + path_slice: path.path_slice.encode(), |
| 207 | + node_rlp: node.to_vec(), |
| 208 | + }), |
| 209 | + }; |
| 210 | + } |
| 211 | + Chunk { |
| 212 | + root: chunk_root, |
| 213 | + terminal_nodes, |
| 214 | + } |
| 215 | + } |
| 216 | + |
| 217 | + // Returns path slices to unresolved chunk roots relative to this chunk root |
| 218 | + pub(crate) fn unresolved_chunks(&self) -> Vec<UnresolvedChunk> { |
| 219 | + let mut result = Vec::new(); |
| 220 | + for node in self.terminal_nodes.iter() { |
| 221 | + let decoded = Node::decoded(&node.node_rlp).expect("All terminal nodes should be valid"); |
| 222 | + if let Node::Branch(slice, children) = decoded { |
| 223 | + for (i, child) in children.iter().enumerate() { |
| 224 | + if let Some(child) = child { |
| 225 | + result.push(UnresolvedChunk { |
| 226 | + path_slice: DecodedPathSlice::from_encoded(&node.path_slice).with_slice_and_index(slice, i), |
| 227 | + chunk_root: *child, |
| 228 | + }) |
| 229 | + } |
| 230 | + } |
| 231 | + } |
| 232 | + } |
| 233 | + result |
| 234 | + } |
| 235 | + |
| 236 | + #[cfg(test)] |
| 237 | + pub(crate) fn into_raw_chunk(self) -> RawChunk { |
| 238 | + RawChunk { |
| 239 | + nodes: self.terminal_nodes, |
| 240 | + } |
| 241 | + } |
| 242 | +} |
| 243 | + |
| 244 | +/// path slice to `chunk_root` is relative to the root of originating chunk. |
| 245 | +#[derive(Debug)] |
| 246 | +pub(crate) struct UnresolvedChunk { |
| 247 | + pub path_slice: DecodedPathSlice, |
| 248 | + pub chunk_root: H256, |
| 249 | +} |
| 250 | + |
| 251 | +impl From<NodePath> for UnresolvedChunk { |
| 252 | + fn from(path: NodePath) -> Self { |
| 253 | + Self { |
| 254 | + path_slice: path.path_slice, |
| 255 | + chunk_root: path.key, |
| 256 | + } |
| 257 | + } |
| 258 | +} |
| 259 | + |
| 260 | +#[derive(Debug)] |
| 261 | +struct NodePath { |
| 262 | + // path slice to the node relative to chunk_root |
| 263 | + path_slice: DecodedPathSlice, |
| 264 | + depth: usize, |
| 265 | + key: H256, |
| 266 | +} |
| 267 | + |
| 268 | +impl NodePath { |
| 269 | + fn new(key: H256) -> NodePath { |
| 270 | + NodePath { |
| 271 | + path_slice: DecodedPathSlice::new(), |
| 272 | + depth: 1, |
| 273 | + key, |
| 274 | + } |
| 275 | + } |
| 276 | + |
| 277 | + fn with_slice_and_index(&self, slice: NibbleSlice, index: usize, key: H256) -> NodePath { |
| 278 | + NodePath { |
| 279 | + path_slice: self.path_slice.with_slice_and_index(slice, index), |
| 280 | + depth: self.depth + 1, |
| 281 | + key, |
| 282 | + } |
| 283 | + } |
| 284 | +} |
0 commit comments