From 0a654afecc1c7e1ee2ed87a918cd07d4ebcbaf4e Mon Sep 17 00:00:00 2001
From: cheme <emericchevalier.pro@gmail.com>
Date: Mon, 19 Aug 2019 23:31:57 +0200
Subject: [PATCH] Fix rlp decode for inline trie nodes. (#10980)

* Fix rlp decode test.

* Proper fix, with associated tests.

* Put tests in their own module.
---
 util/patricia-trie-ethereum/src/lib.rs        | 43 +++++++++++++++++++
 .../src/rlp_node_codec.rs                     | 15 +++++--
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/util/patricia-trie-ethereum/src/lib.rs b/util/patricia-trie-ethereum/src/lib.rs
index b2fb6ff8f..d7b6a261c 100644
--- a/util/patricia-trie-ethereum/src/lib.rs
+++ b/util/patricia-trie-ethereum/src/lib.rs
@@ -145,3 +145,46 @@ pub type TrieFactory = trie::TrieFactory<Layout>;
 pub type TrieError = trie::TrieError<H256, DecoderError>;
 /// Convenience type alias for Keccak/Rlp flavoured trie results
 pub type Result<T> = trie::Result<T, H256, DecoderError>;
+
+#[cfg(test)]
+mod tests {
+
+	use ethereum_types::H256;
+	use crate::{TrieDB, TrieDBMut, trie::TrieMut};
+	use trie::Trie;
+
+	#[test]
+	fn test_inline_encoding_branch() {
+		let mut memdb = journaldb::new_memory_db();
+		let mut root = H256::zero();
+		{
+			let mut triedbmut = TrieDBMut::new(&mut memdb, &mut root);
+			triedbmut.insert(b"foo", b"bar").unwrap();
+			triedbmut.insert(b"fog", b"b").unwrap();
+			triedbmut.insert(b"fot", &vec![0u8;33][..]).unwrap();
+		}
+		let t = TrieDB::new(&memdb, &root).unwrap();
+		assert!(t.contains(b"foo").unwrap());
+		assert!(t.contains(b"fog").unwrap());
+		assert_eq!(t.get(b"foo").unwrap().unwrap(), b"bar".to_vec());
+		assert_eq!(t.get(b"fog").unwrap().unwrap(), b"b".to_vec());
+		assert_eq!(t.get(b"fot").unwrap().unwrap(), vec![0u8;33]);
+	}
+
+	#[test]
+	fn test_inline_encoding_extension() {
+		let mut memdb = journaldb::new_memory_db();
+		let mut root = H256::zero();
+		{
+			let mut triedbmut = TrieDBMut::new(&mut memdb, &mut root);
+			triedbmut.insert(b"foo", b"b").unwrap();
+			triedbmut.insert(b"fog", b"a").unwrap();
+		}
+		let t = TrieDB::new(&memdb, &root).unwrap();
+		assert!(t.contains(b"foo").unwrap());
+		assert!(t.contains(b"fog").unwrap());
+		assert_eq!(t.get(b"foo").unwrap().unwrap(), b"b".to_vec());
+		assert_eq!(t.get(b"fog").unwrap().unwrap(), b"a".to_vec());
+	}
+
+}
diff --git a/util/patricia-trie-ethereum/src/rlp_node_codec.rs b/util/patricia-trie-ethereum/src/rlp_node_codec.rs
index bc8f975bd..0421541ae 100644
--- a/util/patricia-trie-ethereum/src/rlp_node_codec.rs
+++ b/util/patricia-trie-ethereum/src/rlp_node_codec.rs
@@ -98,7 +98,14 @@ impl NodeCodec<KeccakHasher> for RlpNodeCodec<KeccakHasher> {
 				};
 				match from_encoded {
 					(slice, true) => Ok(Node::Leaf(slice, r.at(1)?.data()?)),
-					(slice, false) => Ok(Node::Extension(slice, r.at(1)?.data()?)),
+					(slice, false) => Ok(Node::Extension(slice, {
+						let value = r.at(1)?;
+						if value.is_data() && value.size() == KeccakHasher::LENGTH {
+							value.data()?
+						} else {
+							value.as_raw()
+						}
+					})),
 				}
 			},
 			// branch - first 16 are nodes, 17th is a value (or empty).
@@ -112,7 +119,7 @@ impl NodeCodec<KeccakHasher> for RlpNodeCodec<KeccakHasher> {
 						if value.is_data() && value.size() == KeccakHasher::LENGTH {
 							nodes[i] = Some(value.data()?);
 						} else {
-							return Err(DecoderError::Custom("Rlp is not valid."));
+							nodes[i] = Some(value.as_raw());
 						}
 					}
 				}
@@ -176,7 +183,9 @@ impl NodeCodec<KeccakHasher> for RlpNodeCodec<KeccakHasher> {
 		for child_ref in children {
 			match child_ref.borrow() {
 				Some(c) => match c {
-					ChildReference::Hash(h) => stream.append(h),
+					ChildReference::Hash(h) => {
+						stream.append(h)
+					},
 					ChildReference::Inline(inline_data, length) => {
 						let bytes = &AsRef::<[u8]>::as_ref(inline_data)[..*length];
 						stream.append_raw(bytes, 1)