Merge pull request #2650 from subspace/storage_proof_unused_nodes

ensure storage proof contains only the nodes that are read during the proof check

crates/sp-domains-fraud-proof/src/execution_prover.rs

@@ -9,8 +9,7 @@ use hash_db::{HashDB, Hasher, Prefix};
 use sc_client_api::backend::Backend;
 use sp_api::StorageProof;
 use sp_core::traits::CodeExecutor;
-use sp_core::H256;
-use sp_runtime::traits::{BlakeTwo256, Block as BlockT, HashingFor};
+use sp_runtime::traits::{Block as BlockT, HashingFor};
 use sp_state_machine::backend::AsTrieBackend;
 use sp_state_machine::{TrieBackend, TrieBackendBuilder, TrieBackendStorage};
 use sp_trie::DBValue;
@@ -86,40 +85,6 @@ where
             .map_err(Into::into)
         }
     }
-
-    /// Runs the execution using the partial state constructed from the given storage proof and
-    /// returns the execution result.
-    ///
-    /// The execution result contains the information of state root after applying the execution
-    /// so that it can be used to compare with the one specified in the fraud proof.
-    pub fn check_execution_proof(
-        &self,
-        at: Block::Hash,
-        execution_phase: &ExecutionPhase,
-        call_data: &[u8],
-        pre_execution_root: H256,
-        proof: StorageProof,
-    ) -> sp_blockchain::Result<Vec<u8>> {
-        let state = self.backend.state_at(at)?;
-
-        let trie_backend = state.as_trie_backend();
-
-        let state_runtime_code = sp_state_machine::backend::BackendRuntimeCode::new(trie_backend);
-        let runtime_code = state_runtime_code
-            .runtime_code()
-            .map_err(sp_blockchain::Error::RuntimeCode)?;
-
-        sp_state_machine::execution_proof_check::<BlakeTwo256, _>(
-            pre_execution_root,
-            proof,
-            &mut Default::default(),
-            &*self.executor,
-            execution_phase.execution_method(),
-            call_data,
-            &runtime_code,
-        )
-        .map_err(Into::into)
-    }
 }
 
 /// Create a new trie backend with memory DB delta changes.

crates/sp-domains-fraud-proof/src/host_functions.rs

@@ -7,13 +7,13 @@ use crate::{
 };
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
-use codec::{Decode, Encode};
+use codec::{Codec, Decode, Encode};
 use domain_block_preprocessor::inherents::extract_domain_runtime_upgrade_code;
 use domain_block_preprocessor::stateless_runtime::StatelessRuntime;
 use domain_runtime_primitives::{
     BlockNumber, CheckExtrinsicsValidityError, CHECK_EXTRINSICS_AND_DO_PRE_DISPATCH_METHOD_NAME,
 };
-use hash_db::Hasher;
+use hash_db::{HashDB, Hasher};
 use sc_client_api::execution_extensions::ExtensionsFactory;
 use sc_client_api::BlockBackend;
 use sc_executor::RuntimeVersionOf;
@@ -27,8 +27,8 @@ use sp_externalities::Extensions;
 use sp_messenger::MessengerApi;
 use sp_runtime::traits::{Block as BlockT, Hash as HashT, Header as HeaderT, NumberFor};
 use sp_runtime::OpaqueExtrinsic;
-use sp_state_machine::{create_proof_check_backend, Error, OverlayedChanges, StateMachine};
-use sp_trie::StorageProof;
+use sp_state_machine::{Error, OverlayedChanges, StateMachine, TrieBackend, TrieBackendBuilder};
+use sp_trie::{MemoryDB, StorageProof};
 use std::borrow::Cow;
 use std::marker::PhantomData;
 use std::sync::Arc;
@@ -669,6 +669,47 @@ where
     }
 }
 
+type CreateProofCheckBackedResult<H> =
+    Result<(TrieBackend<MemoryDB<H>, H>, sp_trie::recorder::Recorder<H>), Box<dyn Error>>;
+
+/// Creates a memory db backed with a recorder.
+/// Fork of `sp_state_machine::create_proof_check_backend` with recorder enabled.
+fn create_proof_check_backend_with_recorder<H>(
+    root: H::Out,
+    proof: StorageProof,
+) -> CreateProofCheckBackedResult<H>
+where
+    H: Hasher,
+    H::Out: Codec,
+{
+    let db = proof.into_memory_db();
+    let recorder = sp_trie::recorder::Recorder::<H>::default();
+
+    if db.contains(&root, hash_db::EMPTY_PREFIX) {
+        let backend = TrieBackendBuilder::new(db, root)
+            .with_recorder(recorder.clone())
+            .build();
+        Ok((backend, recorder))
+    } else {
+        Err(Box::new(sp_state_machine::ExecutionError::InvalidProof))
+    }
+}
+
+/// Execution Proof check error.
+#[derive(Debug)]
+pub enum ExecutionProofCheckError {
+    /// Holds the actual execution proof error
+    ExecutionError(Box<dyn sp_state_machine::Error>),
+    /// Error when storage proof contains unused node keys.
+    UnusedNodes,
+}
+
+impl From<Box<dyn sp_state_machine::Error>> for ExecutionProofCheckError {
+    fn from(value: Box<dyn Error>) -> Self {
+        Self::ExecutionError(value)
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 /// Executes the given proof using the runtime
 /// The only difference between sp_state_machine::execution_proof_check is Extensions
@@ -681,13 +722,14 @@ pub(crate) fn execution_proof_check<H, Exec>(
     call_data: &[u8],
     runtime_code: &RuntimeCode,
     extensions: &mut Extensions,
-) -> Result<Vec<u8>, Box<dyn Error>>
+) -> Result<Vec<u8>, ExecutionProofCheckError>
 where
     H: Hasher,
     H::Out: Ord + 'static + codec::Codec,
     Exec: CodeExecutor + Clone + 'static,
 {
-    let trie_backend = create_proof_check_backend::<H>(root, proof)?;
+    let expected_nodes_to_be_read = proof.iter_nodes().count();
+    let (trie_backend, recorder) = create_proof_check_backend_with_recorder::<H>(root, proof)?;
     let result = StateMachine::<_, H, Exec>::new(
         &trie_backend,
         overlay,
@@ -698,7 +740,12 @@ where
         runtime_code,
         CallContext::Offchain,
     )
-    .execute();
+    .execute()?;
+
+    let recorded_proof = recorder.drain_storage_proof();
+    if recorded_proof.iter_nodes().count() != expected_nodes_to_be_read {
+        return Err(ExecutionProofCheckError::UnusedNodes);
+    }
 
-    result
+    Ok(result)
 }

crates/sp-domains/src/proof_provider_and_verifier.rs

@@ -1,6 +1,8 @@
 #[cfg(not(feature = "std"))]
 extern crate alloc;
 
+#[cfg(not(feature = "std"))]
+use alloc::collections::BTreeSet;
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use frame_support::PalletError;
@@ -18,6 +20,8 @@ use sp_std::fmt::Debug;
 use sp_std::marker::PhantomData;
 use sp_trie::{read_trie_value, LayoutV1, StorageProof};
 #[cfg(feature = "std")]
+use std::collections::BTreeSet;
+#[cfg(feature = "std")]
 use trie_db::{DBValue, TrieDBMutBuilder, TrieLayout, TrieMut};
 
 /// Verification error.
@@ -29,6 +33,8 @@ pub enum VerificationError {
     MissingValue,
     /// Failed to decode value.
     FailedToDecode,
+    /// Storage proof contains unused nodes after reading the necessary keys.
+    UnusedNodesInTheProof,
 }
 
 /// Type that provides utilities to verify the storage proof.
@@ -48,15 +54,35 @@ impl<H: Hasher> StorageProofVerifier<H> {
     }
 
     /// Returns the value against a given key.
+    /// Note: Storage proof should contain nodes that are expected else this function errors out.
     pub fn get_bare_value(
         state_root: &H::Out,
         proof: StorageProof,
         key: StorageKey,
     ) -> Result<Vec<u8>, VerificationError> {
+        let expected_nodes_to_be_read = proof.iter_nodes().count();
+        let mut recorder = sp_trie::Recorder::<LayoutV1<H>>::new();
         let db = proof.into_memory_db::<H>();
-        let val = read_trie_value::<LayoutV1<H>, _>(&db, state_root, key.as_ref(), None, None)
-            .map_err(|_| VerificationError::InvalidProof)?
-            .ok_or(VerificationError::MissingValue)?;
+        let val = read_trie_value::<LayoutV1<H>, _>(
+            &db,
+            state_root,
+            key.as_ref(),
+            Some(&mut recorder),
+            None,
+        )
+        .map_err(|_| VerificationError::InvalidProof)?
+        .ok_or(VerificationError::MissingValue)?;
+
+        // check if the storage proof has any extra nodes that are not read.
+        let visited_nodes = recorder
+            .drain()
+            .into_iter()
+            .map(|record| record.data)
+            .collect::<BTreeSet<_>>();
+
+        if expected_nodes_to_be_read != visited_nodes.len() {
+            return Err(VerificationError::UnusedNodesInTheProof);
+        }
 
         Ok(val)
     }

crates/sp-domains/src/valued_trie.rs

@@ -251,15 +251,18 @@ where
 
 #[cfg(test)]
 mod test {
-    use crate::proof_provider_and_verifier::{StorageProofProvider, StorageProofVerifier};
+    use crate::proof_provider_and_verifier::{
+        StorageProofProvider, StorageProofVerifier, VerificationError,
+    };
     use crate::valued_trie::valued_ordered_trie_root;
+    use frame_support::assert_err;
     use parity_scale_codec::{Compact, Encode};
     use rand::rngs::StdRng;
     use rand::{Rng, SeedableRng};
     use sp_core::storage::StorageKey;
     use sp_core::H256;
     use sp_runtime::traits::{BlakeTwo256, Hash};
-    use sp_trie::LayoutV1;
+    use sp_trie::{LayoutV1, StorageProof};
     use trie_db::node::Value;
 
     #[test]
@@ -333,7 +336,7 @@ mod test {
             );
 
             // there is a possibility that wrong key ends up being a different leaf in the merkle tree
-            // but the data that key holds is neither valid extrinsic nor the the one we expect.
+            // but the data that key holds is neither valid extrinsic nor the one we expect.
             if let Ok(data) = result {
                 assert_ne!(data, ext.clone())
             }
@@ -347,4 +350,49 @@ mod test {
             .is_none()
         );
     }
+
+    fn craft_valid_storage_proof_with_multiple_keys() -> (sp_core::H256, StorageProof) {
+        use sp_state_machine::backend::Backend;
+        use sp_state_machine::{prove_read, InMemoryBackend};
+
+        let state_version = sp_runtime::StateVersion::V1;
+
+        // construct storage proof
+        let backend = <InMemoryBackend<sp_core::Blake2Hasher>>::from((
+            vec![
+                (None, vec![(b"key1".to_vec(), Some(b"value1".to_vec()))]),
+                (None, vec![(b"key2".to_vec(), Some(b"value2".to_vec()))]),
+                (None, vec![(b"key3".to_vec(), Some(b"value3".to_vec()))]),
+                (
+                    None,
+                    vec![(b"key4".to_vec(), Some((42u64, 42u32, 42u16, 42u8).encode()))],
+                ),
+                // Value is too big to fit in a branch node
+                (None, vec![(b"key11".to_vec(), Some(vec![0u8; 32]))]),
+            ],
+            state_version,
+        ));
+        let root = backend.storage_root(std::iter::empty(), state_version).0;
+        let proof = prove_read(
+            backend,
+            &[&b"key1"[..], &b"key2"[..], &b"key4"[..], &b"key22"[..]],
+        )
+        .unwrap();
+
+        (root, proof)
+    }
+
+    #[test]
+    fn test_storage_proof_with_unused_nodes() {
+        let (root, storage_proof) = craft_valid_storage_proof_with_multiple_keys();
+        // Verifying the proof with unused nodes should fail
+        assert_err!(
+            StorageProofVerifier::<BlakeTwo256>::get_bare_value(
+                &root,
+                storage_proof,
+                StorageKey(b"key2".to_vec()),
+            ),
+            VerificationError::UnusedNodesInTheProof
+        );
+    }
 }