mirror of
https://github.com/pezkuwichain/pezkuwi-subxt.git
synced 2026-06-11 08:21:05 +00:00
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.
This commit is contained in:
committed by
Gav Wood
parent
aa747e3fae
commit
2510774f3b
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -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], ())
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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]| {
|
||||
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -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());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
@@ -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()));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -35,6 +35,10 @@ pub trait Backend: TryIntoTrieBackend {
|
||||
/// Get keyed storage associated with specific address, or None if there is nothing associated.
|
||||
fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error>;
|
||||
|
||||
/// Retrieve all entries keys of which start with the given prefix and
|
||||
/// call `f` for each of those keys.
|
||||
fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F);
|
||||
|
||||
/// Calculate the storage root, with given delta over what is already stored in
|
||||
/// the backend, and produce a "transaction" that can be used to commit.
|
||||
fn storage_root<I>(&self, delta: I) -> ([u8; 32], Self::Transaction)
|
||||
@@ -105,6 +109,10 @@ impl Backend for InMemory {
|
||||
Ok(self.inner.get(key).map(Clone::clone))
|
||||
}
|
||||
|
||||
fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
|
||||
self.inner.keys().filter(|key| key.starts_with(prefix)).map(|k| &**k).for_each(f);
|
||||
}
|
||||
|
||||
fn storage_root<I>(&self, delta: I) -> ([u8; 32], Self::Transaction)
|
||||
where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
|
||||
{
|
||||
|
||||
@@ -74,6 +74,13 @@ impl<'a, B: 'a + Backend> Ext<'a, B> {
|
||||
let _ = self.storage_root();
|
||||
self.transaction.expect("transaction always set after calling storage root; qed").0
|
||||
}
|
||||
|
||||
/// Invalidates the currently cached storage root and the db transaction.
|
||||
///
|
||||
/// Called when there are changes that likely will invalidate the storage root.
|
||||
fn mark_dirty(&mut self) {
|
||||
self.transaction = None;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -101,10 +108,17 @@ impl<'a, B: 'a> Externalities for Ext<'a, B>
|
||||
}
|
||||
|
||||
fn place_storage(&mut self, key: Vec<u8>, value: Option<Vec<u8>>) {
|
||||
self.transaction = None; // wipe out the transaction since root will no longer be the same.
|
||||
self.mark_dirty();
|
||||
self.overlay.set_storage(key, value);
|
||||
}
|
||||
|
||||
fn clear_prefix(&mut self, prefix: &[u8]) {
|
||||
self.mark_dirty();
|
||||
self.backend.for_keys_with_prefix(prefix, |key| {
|
||||
self.overlay.set_storage(key.to_vec(), None);
|
||||
});
|
||||
}
|
||||
|
||||
fn chain_id(&self) -> u64 {
|
||||
42
|
||||
}
|
||||
|
||||
@@ -132,6 +132,9 @@ pub trait Externalities {
|
||||
self.place_storage(key.to_vec(), None);
|
||||
}
|
||||
|
||||
/// Clear storage entries which keys are start with the given prefix.
|
||||
fn clear_prefix(&mut self, prefix: &[u8]);
|
||||
|
||||
/// Set or clear a storage entry (`key`) of current contract being called (effective immediately).
|
||||
fn place_storage(&mut self, key: Vec<u8>, value: Option<Vec<u8>>);
|
||||
|
||||
|
||||
@@ -69,6 +69,10 @@ impl Backend for ProvingBackend {
|
||||
.get_with(key, &mut *proof_recorder).map(|x| x.map(|val| val.to_vec())).map_err(map_e)
|
||||
}
|
||||
|
||||
fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
|
||||
self.backend.for_keys_with_prefix(prefix, f)
|
||||
}
|
||||
|
||||
fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
|
||||
self.backend.pairs()
|
||||
}
|
||||
@@ -134,7 +138,7 @@ mod tests {
|
||||
let proving_backend = test_proving();
|
||||
assert_eq!(trie_backend.storage(b"key").unwrap(), proving_backend.storage(b"key").unwrap());
|
||||
assert_eq!(trie_backend.pairs(), proving_backend.pairs());
|
||||
|
||||
|
||||
let (trie_root, mut trie_mdb) = trie_backend.storage_root(::std::iter::empty());
|
||||
let (proving_root, mut proving_mdb) = proving_backend.storage_root(::std::iter::empty());
|
||||
assert_eq!(trie_root, proving_root);
|
||||
|
||||
@@ -35,6 +35,12 @@ impl Externalities for TestExternalities {
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_prefix(&mut self, prefix: &[u8]) {
|
||||
self.retain(|key, _|
|
||||
!key.starts_with(prefix)
|
||||
)
|
||||
}
|
||||
|
||||
fn chain_id(&self) -> u64 { 42 }
|
||||
|
||||
fn storage_root(&mut self) -> [u8; 32] {
|
||||
|
||||
@@ -99,6 +99,37 @@ impl Backend for TrieBackend {
|
||||
.get(key).map(|x| x.map(|val| val.to_vec())).map_err(map_e)
|
||||
}
|
||||
|
||||
fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], mut f: F) {
|
||||
let mut read_overlay = MemoryDB::default();
|
||||
let eph = Ephemeral {
|
||||
storage: &self.storage,
|
||||
overlay: &mut read_overlay,
|
||||
};
|
||||
|
||||
let mut iter = move || -> Result<(), Box<TrieError>> {
|
||||
let trie = TrieDB::new(&eph, &self.root)?;
|
||||
let mut iter = trie.iter()?;
|
||||
|
||||
iter.seek(prefix)?;
|
||||
|
||||
for x in iter {
|
||||
let (key, _) = x?;
|
||||
|
||||
if !key.starts_with(prefix) {
|
||||
break;
|
||||
}
|
||||
|
||||
f(&key);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
};
|
||||
|
||||
if let Err(e) = iter() {
|
||||
debug!(target: "trie", "Error while iterating by prefix: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
|
||||
let mut read_overlay = MemoryDB::default();
|
||||
let eph = Ephemeral {
|
||||
@@ -238,6 +269,7 @@ impl TrieBackendStorage {
|
||||
#[cfg(test)]
|
||||
pub mod tests {
|
||||
use super::*;
|
||||
use std::collections::HashSet;
|
||||
|
||||
fn test_db() -> (MemoryDB, TrieH256) {
|
||||
let mut root = TrieH256::default();
|
||||
@@ -293,4 +325,20 @@ pub mod tests {
|
||||
assert!(!tx.drain().is_empty());
|
||||
assert!(new_root != test_trie().storage_root(::std::iter::empty()).0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_walking_works() {
|
||||
let trie = test_trie();
|
||||
|
||||
let mut seen = HashSet::new();
|
||||
trie.for_keys_with_prefix(b"value", |key| {
|
||||
let for_first_time = seen.insert(key.to_vec());
|
||||
assert!(for_first_time, "Seen key '{:?}' more than once", key);
|
||||
});
|
||||
|
||||
let mut expected = HashSet::new();
|
||||
expected.insert(b"value1".to_vec());
|
||||
expected.insert(b"value2".to_vec());
|
||||
assert_eq!(seen, expected);
|
||||
}
|
||||
}
|
||||
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
Reference in New Issue
Block a user