Handle rlp decoding Result in patricia trie (#8166)

* Decode patricia node with UntrustedRlp

(cherry picked from commit efb993b)

* Replace Rlp with UntrustedRlp in triedbmut

* Handle node decode results in trie

util/patricia_trie/src/lookup.rs

@@ -18,7 +18,6 @@
 
 use hashdb::HashDB;
 use nibbleslice::NibbleSlice;
-use rlp::Rlp;
 use ethereum_types::H256;
 
 use super::{TrieError, Query};
@@ -56,7 +55,7 @@ impl<'a, Q: Query> Lookup<'a, Q> {
 			// without incrementing the depth.
 			let mut node_data = &node_data[..];
 			loop {
-				match Node::decoded(node_data) {
+				match Node::decoded(node_data).expect("rlp read from db; qed") {
 					Node::Leaf(slice, value) => {
 						return Ok(match slice == key {
 							true => Some(self.query.decode(value)),
@@ -82,9 +81,8 @@ impl<'a, Q: Query> Lookup<'a, Q> {
 				}
 
 				// check if new node data is inline or hash.
-				let r = Rlp::new(node_data);
-				if r.is_data() && r.size() == 32 {
-					hash = r.as_val();
+				if let Some(h) = Node::try_decode_hash(&node_data) {
+					hash = h;
 					break
 				}
 			}

util/patricia_trie/src/node.rs

@@ -14,11 +14,12 @@
 // You should have received a copy of the GNU General Public License
 // along with Parity.  If not, see <http://www.gnu.org/licenses/>.
 
+use ethereum_types::H256;
 use elastic_array::ElasticArray36;
 use nibbleslice::NibbleSlice;
 use nibblevec::NibbleVec;
 use bytes::*;
-use rlp::*;
+use rlp::{UntrustedRlp, RlpStream, Prototype, DecoderError};
 use hashdb::DBValue;
 
 /// Partial node key type.
@@ -39,30 +40,30 @@ pub enum Node<'a> {
 
 impl<'a> Node<'a> {
 	/// Decode the `node_rlp` and return the Node.
-	pub fn decoded(node_rlp: &'a [u8]) -> Self {
-		let r = Rlp::new(node_rlp);
-		match r.prototype() {
+	pub fn decoded(node_rlp: &'a [u8]) -> Result<Self, DecoderError> {
+		let r = UntrustedRlp::new(node_rlp);
+		match r.prototype()? {
 			// either leaf or extension - decode first item with NibbleSlice::???
 			// and use is_leaf return to figure out which.
 			// if leaf, second item is a value (is_data())
 			// if extension, second item is a node (either SHA3 to be looked up and
 			// fed back into this function or inline RLP which can be fed back into this function).
-			Prototype::List(2) => match NibbleSlice::from_encoded(r.at(0).data()) {
-				(slice, true) => Node::Leaf(slice, r.at(1).data()),
-				(slice, false) => Node::Extension(slice, r.at(1).as_raw()),
+			Prototype::List(2) => match NibbleSlice::from_encoded(r.at(0)?.data()?) {
+				(slice, true) => Ok(Node::Leaf(slice, r.at(1)?.data()?)),
+				(slice, false) => Ok(Node::Extension(slice, r.at(1)?.as_raw())),
 			},
 			// branch - first 16 are nodes, 17th is a value (or empty).
 			Prototype::List(17) => {
 				let mut nodes = [&[] as &[u8]; 16];
 				for i in 0..16 {
-					nodes[i] = r.at(i).as_raw();
+					nodes[i] = r.at(i)?.as_raw();
 				}
-				Node::Branch(nodes, if r.at(16).is_empty() { None } else { Some(r.at(16).data()) })
+				Ok(Node::Branch(nodes, if r.at(16)?.is_empty() { None } else { Some(r.at(16)?.data()?) }))
 			},
 			// an empty branch index.
-			Prototype::Data(0) => Node::Empty,
+			Prototype::Data(0) => Ok(Node::Empty),
 			// something went wrong.
-			_ => panic!("Rlp is not valid.")
+			_ => Err(DecoderError::Custom("Rlp is not valid."))
 		}
 	}
 
@@ -102,6 +103,15 @@ impl<'a> Node<'a> {
 			}
 		}
 	}
+
+	pub fn try_decode_hash(node_data: &[u8]) -> Option<H256> {
+		let r = UntrustedRlp::new(node_data);
+		if r.is_data() && r.size() == 32 {
+			Some(r.as_val().expect("Hash is the correct size of 32 bytes; qed"))
+		} else {
+			None
+		}
+	}
 }
 
 /// An owning node type. Useful for trie iterators.

util/patricia_trie/src/triedb.rs

@@ -17,7 +17,6 @@
 use std::fmt;
 use hashdb::*;
 use nibbleslice::NibbleSlice;
-use rlp::*;
 use super::node::{Node, OwnedNode};
 use super::lookup::Lookup;
 use super::{Trie, TrieItem, TrieError, TrieIterator, Query};
@@ -97,7 +96,10 @@ impl<'db> TrieDB<'db> {
 			Node::Extension(ref slice, ref item) => {
 				write!(f, "'{:?} ", slice)?;
 				if let Ok(node) = self.get_raw_or_lookup(&*item) {
-					self.fmt_all(Node::decoded(&node), f, deepness)?;
+					match Node::decoded(&node) {
+						Ok(n) => self.fmt_all(n, f, deepness)?,
+						Err(err) => writeln!(f, "ERROR decoding node extension Rlp: {}", err)?,
+					}
 				}
 			},
 			Node::Branch(ref nodes, ref value) => {
@@ -108,12 +110,19 @@ impl<'db> TrieDB<'db> {
 				}
 				for i in 0..16 {
 					let node = self.get_raw_or_lookup(&*nodes[i]);
-					match node.as_ref().map(|n| Node::decoded(&*n)) {
-						Ok(Node::Empty) => {},
+					match node.as_ref() { 
 						Ok(n) => {
-							self.fmt_indent(f, deepness + 1)?;
-							write!(f, "'{:x} ", i)?;
-							self.fmt_all(n, f, deepness + 1)?;
+							match Node::decoded(&*n) {
+								Ok(Node::Empty) => {},
+								Ok(n) => {
+									self.fmt_indent(f, deepness + 1)?;
+									write!(f, "'{:x} ", i)?;
+									self.fmt_all(n, f, deepness + 1)?;
+								}
+								Err(e) => {
+									write!(f, "ERROR decoding node branch Rlp: {}", e)?
+								}
+							}
 						}
 						Err(e) => {
 							write!(f, "ERROR: {}", e)?;
@@ -133,16 +142,18 @@ impl<'db> TrieDB<'db> {
 	/// This could be a simple identity operation in the case that the node is sufficiently small, but
 	/// may require a database lookup.
 	fn get_raw_or_lookup(&'db self, node: &'db [u8]) -> super::Result<DBValue> {
-		// check if its keccak + len
-		let r = Rlp::new(node);
-		match r.is_data() && r.size() == 32 {
-			true => {
-				let key = r.as_val::<H256>();
+		match Node::try_decode_hash(node) {
+			Some(key) => {
 				self.db.get(&key).ok_or_else(|| Box::new(TrieError::IncompleteDatabase(key)))
 			}
-			false => Ok(DBValue::from_slice(node))
+			None => Ok(DBValue::from_slice(node))
 		}
 	}
+
+	/// Create a node from raw rlp bytes, assumes valid rlp because encoded locally
+	fn decode_node(node: &'db [u8]) -> Node {
+		Node::decoded(node).expect("rlp read from db; qed")
+	}
 }
 
 impl<'db> Trie for TrieDB<'db> {
@@ -167,7 +178,10 @@ impl<'db> fmt::Debug for TrieDB<'db> {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		writeln!(f, "c={:?} [", self.hash_count)?;
 		let root_rlp = self.db.get(self.root).expect("Trie root not found!");
-		self.fmt_all(Node::decoded(&root_rlp), f, 0)?;
+		match Node::decoded(&root_rlp) {
+			Ok(node) => self.fmt_all(node, f, 0)?,
+			Err(e) => writeln!(f, "ERROR decoding node rlp: {}", e)?,
+		}
 		writeln!(f, "]")
 	}
 }
@@ -222,7 +236,7 @@ impl<'a> TrieDBIterator<'a> {
 	fn seek<'key>(&mut self, mut node_data: DBValue, mut key: NibbleSlice<'key>) -> super::Result<()> {
 		loop {
 			let (data, mid) = {
-				let node = Node::decoded(&node_data);
+				let node = TrieDB::decode_node(&node_data);
 				match node {
 					Node::Leaf(slice, _) => {
 						if slice == key {
@@ -284,7 +298,7 @@ impl<'a> TrieDBIterator<'a> {
 
 	/// Descend into a payload.
 	fn descend(&mut self, d: &[u8]) -> super::Result<()> {
-		let node = Node::decoded(&self.db.get_raw_or_lookup(d)?).into();
+		let node = TrieDB::decode_node(&self.db.get_raw_or_lookup(d)?).into();
 		Ok(self.descend_into_node(node))
 	}
 
@@ -382,7 +396,7 @@ impl<'a> Iterator for TrieDBIterator<'a> {
 					self.trail.pop();
 				},
 				IterStep::Descend(Ok(d)) => {
-					self.descend_into_node(Node::decoded(&d).into())
+					self.descend_into_node(TrieDB::decode_node(&d).into())
 				},
 				IterStep::Descend(Err(e)) => {
 					return Some(Err(e))

util/patricia_trie/src/triedbmut.rs

@@ -24,7 +24,7 @@ use super::node::NodeKey;
 use hashdb::HashDB;
 use bytes::ToPretty;
 use nibbleslice::NibbleSlice;
-use rlp::{Rlp, RlpStream};
+use rlp::{UntrustedRlp, RlpStream};
 use hashdb::DBValue;
 
 use std::collections::{HashSet, VecDeque};
@@ -88,18 +88,17 @@ enum Node {
 impl Node {
 	// load an inline node into memory or get the hash to do the lookup later.
 	fn inline_or_hash(node: &[u8], db: &HashDB, storage: &mut NodeStorage) -> NodeHandle {
-		let r = Rlp::new(node);
-		if r.is_data() && r.size() == 32 {
-			NodeHandle::Hash(r.as_val::<H256>())
-		} else {
-			let child = Node::from_rlp(node, db, storage);
-			NodeHandle::InMemory(storage.alloc(Stored::New(child)))
-		}
+		RlpNode::try_decode_hash(&node)
+			.map(NodeHandle::Hash)
+			.unwrap_or_else(|| {
+				let child = Node::from_rlp(node, db, storage);
+				NodeHandle::InMemory(storage.alloc(Stored::New(child)))
+			})
 	}
 
 	// decode a node from rlp without getting its children.
 	fn from_rlp(rlp: &[u8], db: &HashDB, storage: &mut NodeStorage) -> Self {
-		match RlpNode::decoded(rlp) {
+		match RlpNode::decoded(rlp).expect("rlp read from db; qed") {
 			RlpNode::Empty => Node::Empty,
 			RlpNode::Leaf(k, v) => Node::Leaf(k.encoded(true), DBValue::from_slice(&v)),
 			RlpNode::Extension(key, cb) => {
@@ -108,7 +107,7 @@ impl Node {
 			RlpNode::Branch(ref children_rlp, val) => {
 				let mut child = |i| {
 					let raw = children_rlp[i];
-					let child_rlp = Rlp::new(raw);
+					let child_rlp = UntrustedRlp::new(raw);
 					if !child_rlp.is_empty() {
 						Some(Self::inline_or_hash(raw, db, storage))
 					} else {