Prevent account storage leakage (#270)

* WIP

* Iteration over all keys with the specified prefix

* Add clear_prefix in runtime-io

* Introduce a custom storage impl: Double Map

* Remove prefix

* Impl for_keys_with_prefix for light client

* Fix wasm_executor

* Test storage removal leads to removal of stroage

* Check for ok result in storage tests.

* Add docs.

* Remove commented code under decl_storage!

* Add clear_prefix test in runtime-io

* Add test for wasm_executor

* Prefix walking test.

* Rebuild binaries.

substrate/client/src/light/backend.rs

@@ -145,6 +145,10 @@ impl<Block, F> StateBackend for OnDemandState<Block, F> where Block: BlockT, F:
 		Err(ClientErrorKind::NotAvailableOnLightClient.into()) // TODO: fetch from remote node
 	}
 
+	fn for_keys_with_prefix<A: FnMut(&[u8])>(&self, _prefix: &[u8], _action: A) {
+		// whole state is not available on light node
+	}
+
 	fn storage_root<I>(&self, _delta: I) -> ([u8; 32], Self::Transaction)
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)> {
 		([0; 32], ())

substrate/executor/src/wasm_executor.rs

@@ -225,6 +225,11 @@ impl_function_executor!(this: FunctionExecutor<'e, E>,
 		this.ext.clear_storage(&key);
 		Ok(())
 	},
+	ext_clear_prefix(prefix_data: *const u8, prefix_len: u32) => {
+		let prefix = this.memory.get(prefix_data, prefix_len as usize).map_err(|_| DummyUserError)?;
+		this.ext.clear_prefix(&prefix);
+		Ok(())
+	},
 	// return 0 and place u32::max_value() into written_out if no value exists for the key.
 	ext_get_allocated_storage(key_data: *const u8, key_len: u32, written_out: *mut u32) -> *mut u8 => {
 		let key = this.memory.get(key_data, key_len as usize).map_err(|_| DummyUserError)?;
@@ -577,6 +582,29 @@ mod tests {
 		assert_eq!(expected, ext);
 	}
 
+	#[test]
+	fn clear_prefix_should_work() {
+		let mut ext = TestExternalities::default();
+		ext.set_storage(b"aaa".to_vec(), b"1".to_vec());
+		ext.set_storage(b"aab".to_vec(), b"2".to_vec());
+		ext.set_storage(b"aba".to_vec(), b"3".to_vec());
+		ext.set_storage(b"abb".to_vec(), b"4".to_vec());
+		ext.set_storage(b"bbb".to_vec(), b"5".to_vec());
+		let test_code = include_bytes!("../wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm");
+
+		// This will clear all entries which prefix is "ab".
+		let output = WasmExecutor.call(&mut ext, &test_code[..], "test_clear_prefix", b"ab").unwrap();
+
+		assert_eq!(output, b"all ok!".to_vec());
+
+		let expected: HashMap<_, _> = map![
+			b"aaa".to_vec() => b"1".to_vec(),
+			b"aab".to_vec() => b"2".to_vec(),
+			b"bbb".to_vec() => b"5".to_vec()
+		];
+		assert_eq!(expected, ext);
+	}
+
 	#[test]
 	fn blake2_256_should_work() {
 		let mut ext = TestExternalities::default();

substrate/executor/wasm/src/lib.rs

@@ -11,7 +11,7 @@ extern crate substrate_runtime_io as runtime_io;
 extern crate substrate_runtime_sandbox as sandbox;
 
 use runtime_io::{
-	set_storage, storage, print, blake2_256,
+	set_storage, storage, clear_prefix, print, blake2_256,
 	twox_128, twox_256, ed25519_verify, enumerated_trie_root
 };
 
@@ -29,6 +29,10 @@ impl_stubs!(
 		print("finished!");
 		b"all ok!".to_vec()
 	},
+	test_clear_prefix NO_DECODE => |input| {
+		clear_prefix(input);
+		b"all ok!".to_vec()
+	},
 	test_empty_return NO_DECODE => |_| Vec::new(),
 	test_panic NO_DECODE => |_| panic!("test panic"),
 	test_conditional_panic NO_DECODE => |input: &[u8]| {

substrate/runtime-io/with_std.rs

@@ -68,6 +68,13 @@ pub fn clear_storage(key: &[u8]) {
 	);
 }
 
+/// Clear the storage entries key of which starts with the given prefix.
+pub fn clear_prefix(prefix: &[u8]) {
+	ext::with(|ext|
+		ext.clear_prefix(prefix)
+	);
+}
+
 /// The current relay chain identifier.
 pub fn chain_id() -> u64 {
 	ext::with(|ext|
@@ -211,4 +218,23 @@ mod std_tests {
 			assert_eq!(&w, b"Hello world");
 		});
 	}
+
+	#[test]
+	fn clear_prefix_works() {
+		let mut t: TestExternalities = map![
+			b":a".to_vec() => b"\x0b\0\0\0Hello world".to_vec(),
+			b":abcd".to_vec() => b"\x0b\0\0\0Hello world".to_vec(),
+			b":abc".to_vec() => b"\x0b\0\0\0Hello world".to_vec(),
+			b":abdd".to_vec() => b"\x0b\0\0\0Hello world".to_vec()
+		];
+
+		with_externalities(&mut t, || {
+			clear_prefix(b":abc");
+
+			assert!(storage(b":a").is_some());
+			assert!(storage(b":abdd").is_some());
+			assert!(storage(b":abcd").is_none());
+			assert!(storage(b":abc").is_none());
+		});
+	}
 }

substrate/runtime-io/without_std.rs

@@ -56,6 +56,7 @@ extern "C" {
 	fn ext_print_num(value: u64);
 	fn ext_set_storage(key_data: *const u8, key_len: u32, value_data: *const u8, value_len: u32);
 	fn ext_clear_storage(key_data: *const u8, key_len: u32);
+	fn ext_clear_prefix(prefix_data: *const u8, prefix_len: u32);
 	fn ext_get_allocated_storage(key_data: *const u8, key_len: u32, written_out: *mut u32) -> *mut u8;
 	fn ext_get_storage_into(key_data: *const u8, key_len: u32, value_data: *mut u8, value_len: u32, value_offset: u32) -> u32;
 	fn ext_storage_root(result: *mut u8);
@@ -80,7 +81,7 @@ pub fn storage(key: &[u8]) -> Option<Vec<u8>> {
 	}
 }
 
-/// Set the storage to some particular key.
+/// Set the storage of some particular key to Some value.
 pub fn set_storage(key: &[u8], value: &[u8]) {
 	unsafe {
 		ext_set_storage(
@@ -90,7 +91,7 @@ pub fn set_storage(key: &[u8], value: &[u8]) {
 	}
 }
 
-/// Set the storage to some particular key.
+/// Clear the storage of some particular key.
 pub fn clear_storage(key: &[u8]) {
 	unsafe {
 		ext_clear_storage(
@@ -99,6 +100,16 @@ pub fn clear_storage(key: &[u8]) {
 	}
 }
 
+/// Clear the storage entries key of which starts with the given prefix.
+pub fn clear_prefix(prefix: &[u8]) {
+	unsafe {
+		ext_clear_prefix(
+			prefix.as_ptr(),
+			prefix.len() as u32
+		);
+	}
+}
+
 /// Get `key` from storage, placing the value into `value_out` (as much as possible) and return
 /// the number of bytes that the key in storage was beyond the offset.
 pub fn read_storage(key: &[u8], value_out: &mut [u8], value_offset: usize) -> Option<usize> {

substrate/runtime-support/src/storage/mod.rs

@@ -459,6 +459,11 @@ pub mod unhashed {
 		runtime_io::clear_storage(key);
 	}
 
+	/// Ensure keys with the given `prefix` have no entries in storage.
+	pub fn kill_prefix(prefix: &[u8]) {
+		runtime_io::clear_prefix(prefix);
+	}
+
 	/// Get a Vec of bytes from storage.
 	pub fn get_raw(key: &[u8]) -> Option<Vec<u8>> {
 		runtime_io::storage(key)

substrate/runtime/staking/src/account_db.rs

@@ -20,6 +20,7 @@ use rstd::prelude::*;
 use rstd::cell::RefCell;
 use rstd::collections::btree_map::{BTreeMap, Entry};
 use runtime_support::StorageMap;
+use double_map::StorageDoubleMap;
 use super::*;
 
 pub struct ChangeEntry<T: Trait> {
@@ -61,7 +62,7 @@ pub trait AccountDb<T: Trait> {
 pub struct DirectAccountDb;
 impl<T: Trait> AccountDb<T> for DirectAccountDb {
 	fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option<Vec<u8>> {
-		<StorageOf<T>>::get(&(account.clone(), location.to_vec()))
+		<StorageOf<T>>::get(account.clone(), location.to_vec())
 	}
 	fn get_code(&self, account: &T::AccountId) -> Vec<u8> {
 		<CodeOf<T>>::get(account)
@@ -105,9 +106,9 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 			}
 			for (k, v) in changed.storage.into_iter() {
 				if let Some(value) = v {
-					<StorageOf<T>>::insert((address.clone(), k), &value);
+					<StorageOf<T>>::insert(address.clone(), k, value);
 				} else {
-					<StorageOf<T>>::remove((address.clone(), k));
+					<StorageOf<T>>::remove(address.clone(), k);
 				}
 			}
 		}

substrate/runtime/staking/src/double_map.rs

@@ -0,0 +1,90 @@
+// Copyright 2017 Parity Technologies (UK) Ltd.
+// This file is part of Substrate Demo.
+
+// Substrate Demo is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate Demo is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate Demo.  If not, see <http://www.gnu.org/licenses/>.
+
+//! An implementation of double map backed by storage.
+//!
+//! This implementation is somewhat specialized to the tracking of the storage of accounts.
+
+use rstd::prelude::*;
+use codec::Slicable;
+use runtime_support::storage::unhashed;
+use runtime_io::{blake2_256, twox_128};
+
+/// Returns only a first part of the storage key.
+///
+/// Hashed by XX.
+fn first_part_of_key<M: StorageDoubleMap + ?Sized>(k1: M::Key1) -> [u8; 16] {
+	let mut raw_prefix = Vec::new();
+	raw_prefix.extend(M::PREFIX);
+	raw_prefix.extend(Slicable::encode(&k1));
+	twox_128(&raw_prefix)
+}
+
+/// Returns a compound key that consist of the two parts: (prefix, `k1`) and `k2`.
+///
+/// The first part is hased by XX and then concatenated with a blake2 hash of `k2`.
+fn full_key<M: StorageDoubleMap + ?Sized>(k1: M::Key1, k2: M::Key2) -> Vec<u8> {
+	let first_part = first_part_of_key::<M>(k1);
+	let second_part = blake2_256(&Slicable::encode(&k2));
+
+	let mut k = Vec::new();
+	k.extend(&first_part);
+	k.extend(&second_part);
+	k
+}
+
+/// An implementation of a map with a two keys.
+///
+/// It provides an important ability to efficiently remove all entries
+/// that have a common first key.
+///
+/// # Mapping of keys to a storage path
+///
+/// The storage key (i.e. the key under which the `Value` will be stored) is created from two parts.
+/// The first part is a XX hash of a concatenation of the `PREFIX` and `Key1`. And the second part
+/// is a blake2 hash of a `Key2`.
+///
+/// Blake2 is used for `Key2` is because it will be used as a for a key for contract's storage and
+/// thus will be susceptible for a untrusted input.
+pub trait StorageDoubleMap {
+	type Key1: Slicable;
+	type Key2: Slicable;
+	type Value: Slicable + Default;
+
+	const PREFIX: &'static [u8];
+
+	/// Insert an entry into this map.
+	fn insert(k1: Self::Key1, k2: Self::Key2, val: Self::Value) {
+		unhashed::put(&full_key::<Self>(k1, k2)[..], &val);
+	}
+
+	/// Remove an entry from this map.
+	fn remove(k1: Self::Key1, k2: Self::Key2) {
+		unhashed::kill(&full_key::<Self>(k1, k2)[..]);
+	}
+
+	/// Get an entry from this map.
+	///
+	/// If there is entry stored under the given keys, returns `None`.
+	fn get(k1: Self::Key1, k2: Self::Key2) -> Option<Self::Value> {
+		unhashed::get(&full_key::<Self>(k1, k2)[..])
+	}
+
+	/// Removes all entries that shares the `k1` as the first key.
+	fn remove_prefix(k1: Self::Key1) {
+		unhashed::kill_prefix(&first_part_of_key::<Self>(k1))
+	}
+}

substrate/runtime/staking/src/lib.rs

@@ -57,12 +57,14 @@ use session::OnSessionChange;
 use primitives::traits::{Zero, One, Bounded, RefInto, SimpleArithmetic, Executable, MakePayment,
 	As, AuxLookup, Hashing as HashingT, Member};
 use address::Address as RawAddress;
+use double_map::StorageDoubleMap;
 
 pub mod address;
 mod mock;
 mod tests;
 mod genesis_config;
 mod account_db;
+mod double_map;
 
 #[cfg(feature = "std")]
 pub use genesis_config::GenesisConfig;
@@ -232,16 +234,17 @@ decl_storage! {
 
 	// The code associated with an account.
 	pub CodeOf: b"sta:cod:" => default map [ T::AccountId => Vec<u8> ];	// TODO Vec<u8> values should be optimised to not do a length prefix.
+}
 
-	// The storage items associated with an account/key.
-	// TODO: keys should also be able to take AsRef<KeyType> to ensure Vec<u8>s can be passed as &[u8]
-	// TODO: This will need to be stored as a double-map, with `T::AccountId` using the usual XX hash
-	// function, and then the output of this concatenated onto a separate blake2 hash of the `Vec<u8>`
-	// key. We will then need a `remove_prefix` in addition to `set_storage` which removes all
-	// storage items with a particular prefix otherwise we'll suffer leakage with the removal
-	// of smart contracts.
-//	pub StorageOf: b"sta:sto:" => map [ T::AccountId => map(blake2) Vec<u8> => Vec<u8> ];
-	pub StorageOf: b"sta:sto:" => map [ (T::AccountId, Vec<u8>) => Vec<u8> ];
+/// The storage items associated with an account/key.
+///
+/// TODO: keys should also be able to take AsRef<KeyType> to ensure Vec<u8>s can be passed as &[u8]
+pub(crate) struct StorageOf<T>(::rstd::marker::PhantomData<T>);
+impl<T: Trait> double_map::StorageDoubleMap for StorageOf<T> {
+	type Key1 = T::AccountId;
+	type Key2 = Vec<u8>;
+	type Value = Vec<u8>;
+	const PREFIX: &'static [u8] = b"sta:sto:";
 }
 
 enum NewAccountOutcome {
@@ -623,7 +626,7 @@ impl<T: Trait> Module<T> {
 			.map(|v| (Self::voting_balance(&v) + Self::nomination_balance(&v), v))
 			.collect::<Vec<_>>();
 		intentions.sort_unstable_by(|&(ref b1, _), &(ref b2, _)| b2.cmp(&b1));
-		
+
 		<StakeThreshold<T>>::put(
 			if intentions.len() > 0 {
 				let i = (<ValidatorCount<T>>::get() as usize).min(intentions.len() - 1);
@@ -736,7 +739,7 @@ impl<T: Trait> Module<T> {
 		<FreeBalance<T>>::remove(who);
 		<Bondage<T>>::remove(who);
 		<CodeOf<T>>::remove(who);
-		// TODO: <StorageOf<T>>::remove_prefix(address.clone());
+		<StorageOf<T>>::remove_prefix(who.clone());
 
 		if Self::reserved_balance(who).is_zero() {
 			<system::AccountNonce<T>>::remove(who);

substrate/runtime/staking/src/tests.rs

@@ -585,3 +585,35 @@ fn transferring_incomplete_reserved_balance_should_work() {
 		assert_eq!(Staking::free_balance(&2), 42);
 	});
 }
+
+#[test]
+fn account_removal_removes_storage() {
+	with_externalities(&mut new_test_ext(100, 1, 3, 1, false, 0), || {
+		// Setup two accounts with free balance above than exsistential threshold.
+		{
+			<FreeBalance<Test>>::insert(1, 110);
+			<StorageOf<Test>>::insert(1, b"foo".to_vec(), b"1".to_vec());
+			<StorageOf<Test>>::insert(1, b"bar".to_vec(), b"2".to_vec());
+
+			<FreeBalance<Test>>::insert(2, 110);
+			<StorageOf<Test>>::insert(2, b"hello".to_vec(), b"3".to_vec());
+			<StorageOf<Test>>::insert(2, b"world".to_vec(), b"4".to_vec());
+		}
+
+		// Transfer funds from account 1 of such amount that after this transfer
+		// the balance of account 1 is will be below than exsistential threshold.
+		//
+		// This should lead to the removal of all storage associated with this account.
+		assert_ok!(Staking::transfer(&1, 2.into(), 20));
+
+		// Verify that all entries from account 1 is removed, while
+		// entries from account 2 is in place.
+		{
+			assert_eq!(<StorageOf<Test>>::get(1, b"foo".to_vec()), None);
+			assert_eq!(<StorageOf<Test>>::get(1, b"bar".to_vec()), None);
+
+			assert_eq!(<StorageOf<Test>>::get(2, b"hello".to_vec()), Some(b"3".to_vec()));
+			assert_eq!(<StorageOf<Test>>::get(2, b"world".to_vec()), Some(b"4".to_vec()));
+		}
+	});
+}