use std::sync::Arc;
use parking_lot::RwLock;
use codec::{Decode, Encode, Codec};
use log::debug;
use hash_db::{Hasher, HashDB, EMPTY_PREFIX, Prefix};
use sp_trie::{
MemoryDB, default_child_trie_root, read_trie_value_with, read_child_trie_value_with,
record_all_keys
};
pub use sp_trie::Recorder;
pub use sp_trie::trie_types::{Layout, TrieError};
use crate::trie_backend::TrieBackend;
use crate::trie_backend_essence::{Ephemeral, TrieBackendEssence, TrieBackendStorage};
use crate::{Error, ExecutionError, Backend};
use std::collections::{HashMap, HashSet};
use crate::DBValue;
use sp_core::storage::ChildInfo;
pub struct ProvingBackendRecorder<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
pub(crate) backend: &'a TrieBackendEssence<S, H>,
pub(crate) proof_recorder: &'a mut Recorder<H::Out>,
}
#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
pub struct StorageProof {
trie_nodes: Vec<Vec<u8>>,
}
impl StorageProof {
pub fn new(trie_nodes: Vec<Vec<u8>>) -> Self {
StorageProof { trie_nodes }
}
pub fn empty() -> Self {
StorageProof {
trie_nodes: Vec::new(),
}
}
pub fn is_empty(&self) -> bool {
self.trie_nodes.is_empty()
}
pub fn iter_nodes(self) -> StorageProofNodeIterator {
StorageProofNodeIterator::new(self)
}
}
pub struct StorageProofNodeIterator {
inner: <Vec<Vec<u8>> as IntoIterator>::IntoIter,
}
impl StorageProofNodeIterator {
fn new(proof: StorageProof) -> Self {
StorageProofNodeIterator {
inner: proof.trie_nodes.into_iter(),
}
}
}
impl Iterator for StorageProofNodeIterator {
type Item = Vec<u8>;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
}
pub fn merge_storage_proofs<I>(proofs: I) -> StorageProof
where I: IntoIterator<Item=StorageProof>
{
let trie_nodes = proofs.into_iter()
.flat_map(|proof| proof.iter_nodes())
.collect::<HashSet<_>>()
.into_iter()
.collect();
StorageProof { trie_nodes }
}
impl<'a, S, H> ProvingBackendRecorder<'a, S, H>
where
S: TrieBackendStorage<H>,
H: Hasher,
H::Out: Codec,
{
pub fn storage(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
let mut read_overlay = S::Overlay::default();
let eph = Ephemeral::new(
self.backend.backend_storage(),
&mut read_overlay,
);
let map_e = |e| format!("Trie lookup error: {}", e);
read_trie_value_with::<Layout<H>, _, Ephemeral<S, H>>(
&eph,
self.backend.root(),
key,
&mut *self.proof_recorder,
).map_err(map_e)
}
pub fn child_storage(
&mut self,
storage_key: &[u8],
child_info: ChildInfo,
key: &[u8]
) -> Result<Option<Vec<u8>>, String> {
let root = self.storage(storage_key)?
.and_then(|r| Decode::decode(&mut &r[..]).ok())
.unwrap_or(default_child_trie_root::<Layout<H>>(storage_key));
let mut read_overlay = S::Overlay::default();
let eph = Ephemeral::new(
self.backend.backend_storage(),
&mut read_overlay,
);
let map_e = |e| format!("Trie lookup error: {}", e);
read_child_trie_value_with::<Layout<H>, _, _>(
storage_key,
child_info.keyspace(),
&eph,
&root.as_ref(),
key,
&mut *self.proof_recorder
).map_err(map_e)
}
pub fn record_all_keys(&mut self) {
let mut read_overlay = S::Overlay::default();
let eph = Ephemeral::new(
self.backend.backend_storage(),
&mut read_overlay,
);
let mut iter = move || -> Result<(), Box<TrieError<H::Out>>> {
let root = self.backend.root();
record_all_keys::<Layout<H>, _>(&eph, root, &mut *self.proof_recorder)
};
if let Err(e) = iter() {
debug!(target: "trie", "Error while recording all keys: {}", e);
}
}
}
pub type ProofRecorder<H> = Arc<RwLock<HashMap<<H as Hasher>::Out, Option<DBValue>>>>;
pub struct ProvingBackend<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> (
TrieBackend<ProofRecorderBackend<'a, S, H>, H>,
);
pub struct ProofRecorderBackend<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
backend: &'a S,
proof_recorder: ProofRecorder<H>,
}
impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> ProvingBackend<'a, S, H>
where H::Out: Codec
{
pub fn new(backend: &'a TrieBackend<S, H>) -> Self {
let proof_recorder = Default::default();
Self::new_with_recorder(backend, proof_recorder)
}
pub fn new_with_recorder(
backend: &'a TrieBackend<S, H>,
proof_recorder: ProofRecorder<H>,
) -> Self {
let essence = backend.essence();
let root = essence.root().clone();
let recorder = ProofRecorderBackend {
backend: essence.backend_storage(),
proof_recorder: proof_recorder,
};
ProvingBackend(TrieBackend::new(recorder, root))
}
pub fn extract_proof(&self) -> StorageProof {
let trie_nodes = self.0.essence().backend_storage().proof_recorder
.read()
.iter()
.filter_map(|(_k, v)| v.as_ref().map(|v| v.to_vec()))
.collect();
StorageProof::new(trie_nodes)
}
}
impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> TrieBackendStorage<H>
for ProofRecorderBackend<'a, S, H>
{
type Overlay = S::Overlay;
fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>, String> {
if let Some(v) = self.proof_recorder.read().get(key) {
return Ok(v.clone());
}
let backend_value = self.backend.get(key, prefix)?;
self.proof_recorder.write().insert(key.clone(), backend_value.clone());
Ok(backend_value)
}
}
impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> std::fmt::Debug
for ProvingBackend<'a, S, H>
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "ProvingBackend")
}
}
impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
where
S: 'a + TrieBackendStorage<H>,
H: 'a + Hasher,
H::Out: Ord + Codec,
{
type Error = String;
type Transaction = S::Overlay;
type TrieBackendStorage = S;
fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
self.0.storage(key)
}
fn child_storage(
&self,
storage_key: &[u8],
child_info: ChildInfo,
key: &[u8],
) -> Result<Option<Vec<u8>>, Self::Error> {
self.0.child_storage(storage_key, child_info, key)
}
fn for_keys_in_child_storage<F: FnMut(&[u8])>(
&self,
storage_key: &[u8],
child_info: ChildInfo,
f: F,
) {
self.0.for_keys_in_child_storage(storage_key, child_info, f)
}
fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
self.0.next_storage_key(key)
}
fn next_child_storage_key(
&self,
storage_key: &[u8],
child_info: ChildInfo,
key: &[u8],
) -> Result<Option<Vec<u8>>, Self::Error> {
self.0.next_child_storage_key(storage_key, child_info, key)
}
fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
self.0.for_keys_with_prefix(prefix, f)
}
fn for_key_values_with_prefix<F: FnMut(&[u8], &[u8])>(&self, prefix: &[u8], f: F) {
self.0.for_key_values_with_prefix(prefix, f)
}
fn for_child_keys_with_prefix<F: FnMut(&[u8])>(
&self,
storage_key: &[u8],
child_info: ChildInfo,
prefix: &[u8],
f: F,
) {
self.0.for_child_keys_with_prefix(storage_key, child_info, prefix, f)
}
fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
self.0.pairs()
}
fn keys(&self, prefix: &[u8]) -> Vec<Vec<u8>> {
self.0.keys(prefix)
}
fn child_keys(
&self,
storage_key: &[u8],
child_info: ChildInfo,
prefix: &[u8],
) -> Vec<Vec<u8>> {
self.0.child_keys(storage_key, child_info, prefix)
}
fn storage_root<I>(&self, delta: I) -> (H::Out, Self::Transaction)
where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
{
self.0.storage_root(delta)
}
fn child_storage_root<I>(
&self,
storage_key: &[u8],
child_info: ChildInfo,
delta: I,
) -> (H::Out, bool, Self::Transaction)
where
I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
H::Out: Ord
{
self.0.child_storage_root(storage_key, child_info, delta)
}
}
pub fn create_proof_check_backend<H>(
root: H::Out,
proof: StorageProof,
) -> Result<TrieBackend<MemoryDB<H>, H>, Box<dyn Error>>
where
H: Hasher,
H::Out: Codec,
{
let db = create_proof_check_backend_storage(proof);
if db.contains(&root, EMPTY_PREFIX) {
Ok(TrieBackend::new(db, root))
} else {
Err(Box::new(ExecutionError::InvalidProof))
}
}
pub fn create_proof_check_backend_storage<H>(
proof: StorageProof,
) -> MemoryDB<H>
where
H: Hasher,
{
let mut db = MemoryDB::default();
for item in proof.iter_nodes() {
db.insert(EMPTY_PREFIX, &item);
}
db
}
#[cfg(test)]
mod tests {
use crate::InMemoryBackend;
use crate::trie_backend::tests::test_trie;
use super::*;
use sp_core::{Blake2Hasher, storage::ChildStorageKey};
use crate::proving_backend::create_proof_check_backend;
use sp_trie::PrefixedMemoryDB;
const CHILD_INFO_1: ChildInfo<'static> = ChildInfo::new_default(b"unique_id_1");
const CHILD_INFO_2: ChildInfo<'static> = ChildInfo::new_default(b"unique_id_2");
fn test_proving<'a>(
trie_backend: &'a TrieBackend<PrefixedMemoryDB<Blake2Hasher>,Blake2Hasher>,
) -> ProvingBackend<'a, PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher> {
ProvingBackend::new(trie_backend)
}
#[test]
fn proof_is_empty_until_value_is_read() {
let trie_backend = test_trie();
assert!(test_proving(&trie_backend).extract_proof().is_empty());
}
#[test]
fn proof_is_non_empty_after_value_is_read() {
let trie_backend = test_trie();
let backend = test_proving(&trie_backend);
assert_eq!(backend.storage(b"key").unwrap(), Some(b"value".to_vec()));
assert!(!backend.extract_proof().is_empty());
}
#[test]
fn proof_is_invalid_when_does_not_contains_root() {
use sp_core::H256;
let result = create_proof_check_backend::<Blake2Hasher>(
H256::from_low_u64_be(1),
StorageProof::empty()
);
assert!(result.is_err());
}
#[test]
fn passes_through_backend_calls() {
let trie_backend = test_trie();
let proving_backend = test_proving(&trie_backend);
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);
assert_eq!(trie_mdb.drain(), proving_mdb.drain());
}
#[test]
fn proof_recorded_and_checked() {
let contents = (0..64).map(|i| (vec![i], Some(vec![i]))).collect::<Vec<_>>();
let in_memory = InMemoryBackend::<Blake2Hasher>::default();
let mut in_memory = in_memory.update(vec![(None, contents)]);
let in_memory_root = in_memory.storage_root(::std::iter::empty()).0;
(0..64).for_each(|i| assert_eq!(in_memory.storage(&[i]).unwrap().unwrap(), vec![i]));
let trie = in_memory.as_trie_backend().unwrap();
let trie_root = trie.storage_root(::std::iter::empty()).0;
assert_eq!(in_memory_root, trie_root);
(0..64).for_each(|i| assert_eq!(trie.storage(&[i]).unwrap().unwrap(), vec![i]));
let proving = ProvingBackend::new(trie);
assert_eq!(proving.storage(&[42]).unwrap().unwrap(), vec![42]);
let proof = proving.extract_proof();
let proof_check = create_proof_check_backend::<Blake2Hasher>(in_memory_root.into(), proof).unwrap();
assert_eq!(proof_check.storage(&[42]).unwrap().unwrap(), vec![42]);
}
#[test]
fn proof_recorded_and_checked_with_child() {
let subtrie1 = ChildStorageKey::from_slice(b":child_storage:default:sub1").unwrap();
let subtrie2 = ChildStorageKey::from_slice(b":child_storage:default:sub2").unwrap();
let own1 = subtrie1.into_owned();
let own2 = subtrie2.into_owned();
let contents = vec![
(None, (0..64).map(|i| (vec![i], Some(vec![i]))).collect()),
(Some((own1.clone(), CHILD_INFO_1.to_owned())),
(28..65).map(|i| (vec![i], Some(vec![i]))).collect()),
(Some((own2.clone(), CHILD_INFO_2.to_owned())),
(10..15).map(|i| (vec![i], Some(vec![i]))).collect()),
];
let in_memory = InMemoryBackend::<Blake2Hasher>::default();
let mut in_memory = in_memory.update(contents);
let in_memory_root = in_memory.full_storage_root::<_, Vec<_>, _>(
::std::iter::empty(),
in_memory.child_storage_keys().map(|k|(k.0.to_vec(), Vec::new(), k.1.to_owned()))
).0;
(0..64).for_each(|i| assert_eq!(
in_memory.storage(&[i]).unwrap().unwrap(),
vec![i]
));
(28..65).for_each(|i| assert_eq!(
in_memory.child_storage(&own1[..], CHILD_INFO_1, &[i]).unwrap().unwrap(),
vec![i]
));
(10..15).for_each(|i| assert_eq!(
in_memory.child_storage(&own2[..], CHILD_INFO_2, &[i]).unwrap().unwrap(),
vec![i]
));
let trie = in_memory.as_trie_backend().unwrap();
let trie_root = trie.storage_root(::std::iter::empty()).0;
assert_eq!(in_memory_root, trie_root);
(0..64).for_each(|i| assert_eq!(
trie.storage(&[i]).unwrap().unwrap(),
vec![i]
));
let proving = ProvingBackend::new(trie);
assert_eq!(proving.storage(&[42]).unwrap().unwrap(), vec![42]);
let proof = proving.extract_proof();
let proof_check = create_proof_check_backend::<Blake2Hasher>(
in_memory_root.into(),
proof
).unwrap();
assert!(proof_check.storage(&[0]).is_err());
assert_eq!(proof_check.storage(&[42]).unwrap().unwrap(), vec![42]);
assert_eq!(proof_check.storage(&[41]).unwrap().unwrap(), vec![41]);
assert_eq!(proof_check.storage(&[64]).unwrap(), None);
let proving = ProvingBackend::new(trie);
assert_eq!(proving.child_storage(&own1[..], CHILD_INFO_1, &[64]), Ok(Some(vec![64])));
let proof = proving.extract_proof();
let proof_check = create_proof_check_backend::<Blake2Hasher>(
in_memory_root.into(),
proof
).unwrap();
assert_eq!(
proof_check.child_storage(&own1[..], CHILD_INFO_1, &[64]).unwrap().unwrap(),
vec![64]
);
}
}