Check storage_key for validity (#2316)

* Intro `ChildStorageKey` for checked child keys

* Get rid of Into in Externalities trait

* Use Cow in ChildStorageKey

* Fix tests for state-machine.

* Clean

* child_storage_root always return a value

* Don't return Option from Ext::child_storage_root

* Return 42 in child_storage_root

* Return CHILD_STORAGE_KEY_PREFIX from trie id gen

* Bump spec and impl version.

* Require `:default:` in `is_child_trie_key_valid`

* Add `default:` prefix.

* Introduce `into_owned` for `ChildStorageKey`.

* Add documentation.

* Fix state-machine tests

* Remove outdated TODO

I check out with Emeric and he is ok with that

* child_storage_root is infailable

* Nit

* Move assert after check.

* Apply suggestions from @DemiMarie-parity

Co-Authored-By: pepyakin <s.pepyakin@gmail.com>

* Formatting nit in core/executor/src/wasm_executor.rs

Co-Authored-By: DemiMarie-parity <48690212+DemiMarie-parity@users.noreply.github.com>

* Formatting nits from @thiolliere

substrate/core/state-machine/src/basic.rs

@@ -23,7 +23,7 @@ use heapsize::HeapSizeOf;
 use trie::trie_root;
 use primitives::storage::well_known_keys::{CHANGES_TRIE_CONFIG, CODE, HEAP_PAGES};
 use parity_codec::Encode;
-use super::{Externalities, OverlayedChanges};
+use super::{ChildStorageKey, Externalities, OverlayedChanges};
 use log::warn;
 
 /// Simple HashMap-based Externalities impl.
@@ -115,7 +115,7 @@ impl<H: Hasher> Externalities<H> for BasicExternalities where H::Out: Ord + Heap
 		Externalities::<H>::storage(self, key)
 	}
 
-	fn child_storage(&self, _storage_key: &[u8], _key: &[u8]) -> Option<Vec<u8>> {
+	fn child_storage(&self, _storage_key: ChildStorageKey<H>, _key: &[u8]) -> Option<Vec<u8>> {
 		None
 	}
 
@@ -132,11 +132,10 @@ impl<H: Hasher> Externalities<H> for BasicExternalities where H::Out: Ord + Heap
 		}
 	}
 
-	fn place_child_storage(&mut self, _storage_key: Vec<u8>, _key: Vec<u8>, _value: Option<Vec<u8>>) -> bool {
-		false
+	fn place_child_storage(&mut self, _storage_key: ChildStorageKey<H>, _key: Vec<u8>, _value: Option<Vec<u8>>) {
 	}
 
-	fn kill_child_storage(&mut self, _storage_key: &[u8]) { }
+	fn kill_child_storage(&mut self, _storage_key: ChildStorageKey<H>) { }
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {
 		self.changes.clear_prefix(prefix);
@@ -149,8 +148,8 @@ impl<H: Hasher> Externalities<H> for BasicExternalities where H::Out: Ord + Heap
 		trie_root::<H, _, _, _>(self.inner.clone())
 	}
 
-	fn child_storage_root(&mut self, _storage_key: &[u8]) -> Option<Vec<u8>> {
-		None
+	fn child_storage_root(&mut self, _storage_key: ChildStorageKey<H>) -> Vec<u8> {
+		vec![42]
 	}
 
 	fn storage_changes_root(&mut self, _parent: H::Out, _parent_num: u64) -> Option<H::Out> {

substrate/core/state-machine/src/ext.rs

@@ -20,10 +20,10 @@ use std::{error, fmt, cmp::Ord};
 use log::warn;
 use crate::backend::{Backend, Consolidate};
 use crate::changes_trie::{AnchorBlockId, Storage as ChangesTrieStorage, compute_changes_trie_root};
-use crate::{Externalities, OverlayedChanges, OffchainExt};
+use crate::{Externalities, OverlayedChanges, OffchainExt, ChildStorageKey};
 use hash_db::Hasher;
 use primitives::storage::well_known_keys::is_child_storage_key;
-use trie::{MemoryDB, TrieDBMut, TrieMut, default_child_trie_root, is_child_trie_key_valid};
+use trie::{MemoryDB, TrieDBMut, TrieMut, default_child_trie_root};
 use heapsize::HeapSizeOf;
 
 const EXT_NOT_ALLOWED_TO_FAIL: &str = "Externalities not allowed to fail within runtime";
@@ -214,10 +214,10 @@ where
 		self.backend.storage_hash(key).expect(EXT_NOT_ALLOWED_TO_FAIL)
 	}
 
-	fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
+	fn child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> Option<Vec<u8>> {
 		let _guard = panic_handler::AbortGuard::new(true);
-		self.overlay.child_storage(storage_key, key).map(|x| x.map(|x| x.to_vec())).unwrap_or_else(||
-			self.backend.child_storage(storage_key, key).expect(EXT_NOT_ALLOWED_TO_FAIL))
+		self.overlay.child_storage(storage_key.as_ref(), key).map(|x| x.map(|x| x.to_vec())).unwrap_or_else(||
+			self.backend.child_storage(storage_key.as_ref(), key).expect(EXT_NOT_ALLOWED_TO_FAIL))
 	}
 
 	fn exists_storage(&self, key: &[u8]) -> bool {
@@ -228,11 +228,12 @@ where
 		}
 	}
 
-	fn exists_child_storage(&self, storage_key: &[u8], key: &[u8]) -> bool {
+	fn exists_child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> bool {
 		let _guard = panic_handler::AbortGuard::new(true);
-		match self.overlay.child_storage(storage_key, key) {
+
+		match self.overlay.child_storage(storage_key.as_ref(), key) {
 			Some(x) => x.is_some(),
-			_ => self.backend.exists_child_storage(storage_key, key).expect(EXT_NOT_ALLOWED_TO_FAIL),
+			_ => self.backend.exists_child_storage(storage_key.as_ref(), key).expect(EXT_NOT_ALLOWED_TO_FAIL),
 		}
 	}
 
@@ -247,28 +248,20 @@ where
 		self.overlay.set_storage(key, value);
 	}
 
-	fn place_child_storage(&mut self, storage_key: Vec<u8>, key: Vec<u8>, value: Option<Vec<u8>>) -> bool {
+	fn place_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: Vec<u8>, value: Option<Vec<u8>>) {
 		let _guard = panic_handler::AbortGuard::new(true);
-		if !is_child_storage_key(&storage_key) || !is_child_trie_key_valid::<H>(&storage_key) {
-			return false;
-		}
 
 		self.mark_dirty();
-		self.overlay.set_child_storage(storage_key, key, value);
-
-		true
+		self.overlay.set_child_storage(storage_key.into_owned(), key, value);
 	}
 
-	fn kill_child_storage(&mut self, storage_key: &[u8]) {
+	fn kill_child_storage(&mut self, storage_key: ChildStorageKey<H>) {
 		let _guard = panic_handler::AbortGuard::new(true);
-		if !is_child_storage_key(storage_key) || !is_child_trie_key_valid::<H>(storage_key) {
-			return;
-		}
 
 		self.mark_dirty();
-		self.overlay.clear_child_storage(storage_key);
-		self.backend.for_keys_in_child_storage(storage_key, |key| {
-			self.overlay.set_child_storage(storage_key.to_vec(), key.to_vec(), None);
+		self.overlay.clear_child_storage(storage_key.as_ref());
+		self.backend.for_keys_in_child_storage(storage_key.as_ref(), |key| {
+			self.overlay.set_child_storage(storage_key.as_ref().to_vec(), key.to_vec(), None);
 		});
 	}
 
@@ -315,17 +308,17 @@ where
 		root
 	}
 
-	fn child_storage_root(&mut self, storage_key: &[u8]) -> Option<Vec<u8>> {
+	fn child_storage_root(&mut self, storage_key: ChildStorageKey<H>) -> Vec<u8> {
 		let _guard = panic_handler::AbortGuard::new(true);
-		if !is_child_storage_key(storage_key) || !is_child_trie_key_valid::<H>(storage_key) {
-			return None;
-		}
-
 		if self.storage_transaction.is_some() {
-			return Some(self.storage(storage_key).unwrap_or(default_child_trie_root::<H>(storage_key)));
+			self
+				.storage(storage_key.as_ref())
+				.unwrap_or(
+					default_child_trie_root::<H>(storage_key.as_ref())
+				)
+		} else {
+			self.child_storage_root_transaction(storage_key.as_ref()).0
 		}
-
-		Some(self.child_storage_root_transaction(storage_key).0)
 	}
 
 	fn storage_changes_root(&mut self, parent: H::Out, parent_num: u64) -> Option<H::Out> {

substrate/core/state-machine/src/lib.rs

@@ -19,6 +19,7 @@
 #![warn(missing_docs)]
 
 use std::{fmt, panic::UnwindSafe, result, marker::PhantomData};
+use std::borrow::Cow;
 use log::warn;
 use hash_db::Hasher;
 use heapsize::HeapSizeOf;
@@ -55,6 +56,58 @@ pub use proving_backend::{create_proof_check_backend, create_proof_check_backend
 pub use trie_backend_essence::{TrieBackendStorage, Storage};
 pub use trie_backend::TrieBackend;
 
+/// A wrapper around a child storage key.
+///
+/// This wrapper ensures that the child storage key is correct and properly used.  It is
+/// impossible to create an instance of this struct without providing a correct `storage_key`.
+pub struct ChildStorageKey<'a, H: Hasher> {
+	storage_key: Cow<'a, [u8]>,
+	_hasher: PhantomData<H>,
+}
+
+impl<'a, H: Hasher> ChildStorageKey<'a, H> {
+	fn new(storage_key: Cow<'a, [u8]>) -> Option<Self> {
+		if !trie::is_child_trie_key_valid::<H>(&storage_key) {
+			return None;
+		}
+
+		Some(ChildStorageKey {
+			storage_key,
+			_hasher: PhantomData,
+		})
+	}
+
+	/// Create a new `ChildStorageKey` from a vector.
+	///
+	/// `storage_key` has should start with `:child_storage:default:`
+	/// See `is_child_trie_key_valid` for more details.
+	pub fn from_vec(key: Vec<u8>) -> Option<Self> {
+		Self::new(Cow::Owned(key))
+	}
+
+	/// Create a new `ChildStorageKey` from a slice.
+	///
+	/// `storage_key` has should start with `:child_storage:default:`
+	/// See `is_child_trie_key_valid` for more details.
+	pub fn from_slice(key: &'a [u8]) -> Option<Self> {
+		Self::new(Cow::Borrowed(key))
+	}
+
+	/// Get access to the byte representation of the storage key.
+	///
+	/// This key is guaranteed to be correct.
+	pub fn as_ref(&self) -> &[u8] {
+		&*self.storage_key
+	}
+
+	/// Destruct this instance into an owned vector that represents the storage key.
+	///
+	/// This key is guaranteed to be correct.
+	pub fn into_owned(self) -> Vec<u8> {
+		self.storage_key.into_owned()
+	}
+}
+
 /// State Machine Error bound.
 ///
 /// This should reflect WASM error type bound for future compatibility.
@@ -105,7 +158,7 @@ pub trait Externalities<H: Hasher> {
 	}
 
 	/// Read child runtime storage.
-	fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>>;
+	fn child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> Option<Vec<u8>>;
 
 	/// Set storage entry `key` of current contract being called (effective immediately).
 	fn set_storage(&mut self, key: Vec<u8>, value: Vec<u8>) {
@@ -113,7 +166,7 @@ pub trait Externalities<H: Hasher> {
 	}
 
 	/// Set child storage entry `key` of current contract being called (effective immediately).
-	fn set_child_storage(&mut self, storage_key: Vec<u8>, key: Vec<u8>, value: Vec<u8>) -> bool {
+	fn set_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: Vec<u8>, value: Vec<u8>) {
 		self.place_child_storage(storage_key, key, Some(value))
 	}
 
@@ -123,8 +176,8 @@ pub trait Externalities<H: Hasher> {
 	}
 
 	/// Clear a child storage entry (`key`) of current contract being called (effective immediately).
-	fn clear_child_storage(&mut self, storage_key: &[u8], key: &[u8]) -> bool {
-		self.place_child_storage(storage_key.to_vec(), key.to_vec(), None)
+	fn clear_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: &[u8]) {
+		self.place_child_storage(storage_key, key.to_vec(), None)
 	}
 
 	/// Whether a storage entry exists.
@@ -133,12 +186,12 @@ pub trait Externalities<H: Hasher> {
 	}
 
 	/// Whether a child storage entry exists.
-	fn exists_child_storage(&self, storage_key: &[u8], key: &[u8]) -> bool {
+	fn exists_child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> bool {
 		self.child_storage(storage_key, key).is_some()
 	}
 
 	/// Clear an entire child storage.
-	fn kill_child_storage(&mut self, storage_key: &[u8]);
+	fn kill_child_storage(&mut self, storage_key: ChildStorageKey<H>);
 
 	/// Clear storage entries which keys are start with the given prefix.
 	fn clear_prefix(&mut self, prefix: &[u8]);
@@ -147,7 +200,7 @@ pub trait Externalities<H: Hasher> {
 	fn place_storage(&mut self, key: Vec<u8>, value: Option<Vec<u8>>);
 
 	/// Set or clear a child storage entry. Return whether the operation succeeds.
-	fn place_child_storage(&mut self, storage_key: Vec<u8>, key: Vec<u8>, value: Option<Vec<u8>>) -> bool;
+	fn place_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: Vec<u8>, value: Option<Vec<u8>>);
 
 	/// Get the identity of the chain.
 	fn chain_id(&self) -> u64;
@@ -155,10 +208,11 @@ pub trait Externalities<H: Hasher> {
 	/// Get the trie root of the current storage map. This will also update all child storage keys in the top-level storage map.
 	fn storage_root(&mut self) -> H::Out where H::Out: Ord;
 
-	/// Get the trie root of a child storage map. This will also update the value of the child storage keys in the top-level storage map. If the storage root equals default hash as defined by trie, the key in top-level storage map will be removed.
-	///
-	/// Returns None if key provided is not a storage key. This can due to not being started with CHILD_STORAGE_KEY_PREFIX, or the trie implementation regards the key as invalid.
-	fn child_storage_root(&mut self, storage_key: &[u8]) -> Option<Vec<u8>>;
+	/// Get the trie root of a child storage map. This will also update the value of the child
+	/// storage keys in the top-level storage map.
+	/// If the storage root equals the default hash as defined by the trie, the key in the top-level
+	/// storage map will be removed.
+	fn child_storage_root(&mut self, storage_key: ChildStorageKey<H>) -> Vec<u8>;
 
 	/// Get the change trie root of the current storage overlay at a block with given parent.
 	fn storage_changes_root(&mut self, parent: H::Out, parent_num: u64) -> Option<H::Out> where H::Out: Ord;
@@ -905,12 +959,35 @@ mod tests {
 		let backend = InMemory::<Blake2Hasher>::default().try_into_trie_backend().unwrap();
 		let changes_trie_storage = InMemoryChangesTrieStorage::new();
 		let mut overlay = OverlayedChanges::default();
-		let mut ext = Ext::new(&mut overlay, &backend, Some(&changes_trie_storage), NeverOffchainExt::new());
+		let mut ext = Ext::new(
+			&mut overlay,
+			&backend,
+			Some(&changes_trie_storage),
+			NeverOffchainExt::new()
+		);
 
-		assert!(ext.set_child_storage(b":child_storage:testchild".to_vec(), b"abc".to_vec(), b"def".to_vec()));
-		assert_eq!(ext.child_storage(b":child_storage:testchild", b"abc"), Some(b"def".to_vec()));
-		ext.kill_child_storage(b":child_storage:testchild");
-		assert_eq!(ext.child_storage(b":child_storage:testchild", b"abc"), None);
+		ext.set_child_storage(
+			ChildStorageKey::from_slice(b":child_storage:default:testchild").unwrap(),
+			b"abc".to_vec(),
+			b"def".to_vec()
+		);
+		assert_eq!(
+			ext.child_storage(
+				ChildStorageKey::from_slice(b":child_storage:default:testchild").unwrap(),
+				b"abc"
+			),
+			Some(b"def".to_vec())
+		);
+		ext.kill_child_storage(
+			ChildStorageKey::from_slice(b":child_storage:default:testchild").unwrap()
+		);
+		assert_eq!(
+			ext.child_storage(
+				ChildStorageKey::from_slice(b":child_storage:default:testchild").unwrap(),
+				b"abc"
+			),
+			None
+		);
 	}
 
 	#[test]

substrate/core/state-machine/src/testing.rs

@@ -25,7 +25,7 @@ use crate::backend::InMemory;
 use crate::changes_trie::{compute_changes_trie_root, InMemoryStorage as ChangesTrieInMemoryStorage, AnchorBlockId};
 use primitives::storage::well_known_keys::{CHANGES_TRIE_CONFIG, CODE, HEAP_PAGES};
 use parity_codec::Encode;
-use super::{Externalities, OverlayedChanges};
+use super::{ChildStorageKey, Externalities, OverlayedChanges};
 
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H: Hasher> where H::Out: HeapSizeOf {
@@ -122,8 +122,8 @@ impl<H: Hasher> Externalities<H> for TestExternalities<H> where H::Out: Ord + He
 		self.storage(key)
 	}
 
-	fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
-		self.changes.child_storage(storage_key, key)?.map(Vec::from)
+	fn child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> Option<Vec<u8>> {
+		self.changes.child_storage(storage_key.as_ref(), key)?.map(Vec::from)
 	}
 
 	fn place_storage(&mut self, key: Vec<u8>, maybe_value: Option<Vec<u8>>) {
@@ -139,14 +139,12 @@ impl<H: Hasher> Externalities<H> for TestExternalities<H> where H::Out: Ord + He
 		}
 	}
 
-	fn place_child_storage(&mut self, storage_key: Vec<u8>, key: Vec<u8>, value: Option<Vec<u8>>) -> bool {
-		self.changes.set_child_storage(storage_key, key, value);
-		// TODO place_child_storage and set_child_storage should always be valid (create child on set)?
-		true
+	fn place_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: Vec<u8>, value: Option<Vec<u8>>) {
+		self.changes.set_child_storage(storage_key.into_owned(), key, value);
 	}
 
-	fn kill_child_storage(&mut self, storage_key: &[u8]) {
-		self.changes.clear_child_storage(storage_key);
+	fn kill_child_storage(&mut self, storage_key: ChildStorageKey<H>) {
+		self.changes.clear_child_storage(storage_key.as_ref());
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {
@@ -160,8 +158,8 @@ impl<H: Hasher> Externalities<H> for TestExternalities<H> where H::Out: Ord + He
 		trie_root::<H, _, _, _>(self.inner.clone())
 	}
 
-	fn child_storage_root(&mut self, _storage_key: &[u8]) -> Option<Vec<u8>> {
-		None
+	fn child_storage_root(&mut self, _storage_key: ChildStorageKey<H>) -> Vec<u8> {
+		unimplemented!()
 	}
 
 	fn storage_changes_root(&mut self, parent: H::Out, parent_num: u64) -> Option<H::Out> {