Decode Merkle proofs entirely at once

bin/light-base/src/json_rpc_service/chain_head.rs

@@ -300,6 +300,15 @@ impl<TPlat: Platform> Background<TPlat> {
                         }
                         .to_json_call_object_parameters(None)
                     }
+                    Some(Err(runtime_service::RuntimeCallError::MissingProofEntry)) => {
+                        methods::ServerToClient::chainHead_unstable_callEvent {
+                            subscription: (&subscription_id).into(),
+                            result: methods::ChainHeadCallEvent::Error {
+                                error: "incomplete call proof".into(),
+                            },
+                        }
+                        .to_json_call_object_parameters(None)
+                    }
                     Some(Err(runtime_service::RuntimeCallError::CallProof(error))) => {
                         methods::ServerToClient::chainHead_unstable_callEvent {
                             subscription: (&subscription_id).into(),

bin/light-base/src/runtime_service.rs

@@ -54,9 +54,10 @@
 //! large, the subscription is force-killed by the [`RuntimeService`].
 //!
 
-use crate::{network_service, platform::Platform, sync_service};
+use crate::{platform::Platform, sync_service};
 
 use alloc::{
+    borrow::ToOwned as _,
     boxed::Box,
     collections::BTreeMap,
     format,
@@ -82,7 +83,7 @@ use smoldot::{
     executor, header,
     informant::{BytesDisplay, HashDisplay},
     network::protocol,
-    trie::{self, proof_verify},
+    trie::{self, proof_decode},
 };
 
 /// Configuration for a runtime service.
@@ -803,6 +804,14 @@ impl<'a, TPlat: Platform> RuntimeLock<'a, TPlat> {
             .await
             .map_err(RuntimeCallError::CallProof);
 
+        let call_proof = call_proof.and_then(|call_proof| {
+            proof_decode::decode_and_verify_proof(proof_decode::VerifyProofConfig {
+                proof: call_proof.decode().into_iter().map(|v| v.to_owned()), // TODO: to_owned() inefficiency, need some help from the networking to obtain the owned data
+                trie_root_hash: &self.block_state_root_hash,
+            })
+            .map_err(RuntimeCallError::StorageRetrieval)
+        });
+
         let (guarded, virtual_machine) = match self.runtime.runtime.as_ref() {
             Ok(r) => {
                 let mut lock = r.virtual_machine.lock().await;
@@ -829,7 +838,7 @@ impl<'a, TPlat: Platform> RuntimeLock<'a, TPlat> {
 pub struct RuntimeCallLock<'a> {
     guarded: MutexGuard<'a, Option<executor::host::HostVmPrototype>>,
     block_state_root_hash: [u8; 32],
-    call_proof: Result<network_service::EncodedMerkleProof, RuntimeCallError>,
+    call_proof: Result<trie::proof_decode::DecodedTrieProof<Vec<u8>>, RuntimeCallError>,
 }
 
 impl<'a> RuntimeCallLock<'a> {
@@ -845,17 +854,13 @@ impl<'a> RuntimeCallLock<'a> {
     // TODO: if proof is invalid, we should give the option to fetch another call proof
     pub fn storage_entry(&self, requested_key: &[u8]) -> Result<Option<&[u8]>, RuntimeCallError> {
         let call_proof = match &self.call_proof {
-            Ok(p) => p.decode(),
+            Ok(p) => p,
             Err(err) => return Err(err.clone()),
         };
 
-        match proof_verify::verify_proof(proof_verify::VerifyProofConfig {
-            requested_key,
-            trie_root_hash: self.block_storage_root(),
-            proof: call_proof.iter().map(|v| &v[..]),
-        }) {
-            Ok(v) => Ok(v),
-            Err(err) => Err(RuntimeCallError::StorageRetrieval(err)),
+        match call_proof.storage_value(requested_key) {
+            Some(v) => Ok(v),
+            None => Err(RuntimeCallError::MissingProofEntry),
         }
     }
 
@@ -870,52 +875,39 @@ impl<'a> RuntimeCallLock<'a> {
         &'_ self,
         prefix: &[u8],
     ) -> Result<impl Iterator<Item = impl AsRef<[u8]> + '_>, RuntimeCallError> {
-        // TODO: this is sub-optimal as we iterate over the proof multiple times and do a lot of Vec allocations
+        // TODO: this could be a function in the proof_decode module
         let mut to_find = vec![trie::bytes_to_nibbles(prefix.iter().copied()).collect::<Vec<_>>()];
         let mut output = Vec::new();
 
         let call_proof = match &self.call_proof {
-            Ok(p) => p.decode(),
+            Ok(p) => p,
             Err(err) => return Err(err.clone()),
         };
 
         for key in mem::take(&mut to_find) {
-            let node_info = proof_verify::trie_node_info(proof_verify::TrieNodeInfoConfig {
-                requested_key: key.iter().cloned(),
-                trie_root_hash: self.block_storage_root(),
-                proof: call_proof.iter().map(|v| &v[..]),
-            })
-            .map_err(RuntimeCallError::StorageRetrieval)?;
+            let node_info = call_proof
+                .trie_node_info(&key)
+                .ok_or(RuntimeCallError::MissingProofEntry)?;
 
             if matches!(
                 node_info.storage_value,
-                proof_verify::StorageValue::Known(_)
-                    | proof_verify::StorageValue::HashKnownValueMissing(_)
+                proof_decode::StorageValue::Known(_)
+                    | proof_decode::StorageValue::HashKnownValueMissing(_)
             ) {
                 assert_eq!(key.len() % 2, 0);
                 output.push(
                     trie::nibbles_to_bytes_suffix_extend(key.iter().copied()).collect::<Vec<_>>(),
                 );
             }
 
-            match node_info.children {
-                proof_verify::Children::None => {}
-                proof_verify::Children::One(nibble) => {
-                    let mut child = key.clone();
-                    child.push(nibble);
-                    to_find.push(child);
+            for nibble in trie::all_nibbles() {
+                if !node_info.children.has_child(nibble) {
+                    continue;
                 }
-                proof_verify::Children::Multiple { children_bitmap } => {
-                    for nibble in trie::all_nibbles() {
-                        if (children_bitmap & (1 << u8::from(nibble))) == 0 {
-                            continue;
-                        }
 
-                        let mut child = key.clone();
-                        child.push(nibble);
-                        to_find.push(child);
-                    }
-                }
+                let mut child = key.clone();
+                child.push(nibble);
+                to_find.push(child);
             }
         }
 
@@ -952,7 +944,9 @@ pub enum RuntimeCallError {
     /// Error while retrieving the storage item from other nodes.
     // TODO: change error type?
     #[display(fmt = "Error in call proof: {}", _0)]
-    StorageRetrieval(proof_verify::Error),
+    StorageRetrieval(proof_decode::Error),
+    /// One or more entries are missing from the call proof.
+    MissingProofEntry,
     /// Error while retrieving the call proof from the network.
     #[display(fmt = "Error when retrieving the call proof: {}", _0)]
     CallProof(sync_service::CallProofQueryError),
@@ -968,8 +962,9 @@ impl RuntimeCallError {
         match self {
             RuntimeCallError::InvalidRuntime(_) => false,
             // TODO: as a temporary hack, we consider `TrieRootNotFound` as the remote not knowing about the requested block; see https://github.com/paritytech/substrate/pull/8046
-            RuntimeCallError::StorageRetrieval(proof_verify::Error::TrieRootNotFound) => true,
+            RuntimeCallError::StorageRetrieval(proof_decode::Error::TrieRootNotFound) => true,
             RuntimeCallError::StorageRetrieval(_) => false,
+            RuntimeCallError::MissingProofEntry => false,
             RuntimeCallError::CallProof(err) => err.is_network_problem(),
             RuntimeCallError::StorageQuery(err) => err.is_network_problem(),
         }

bin/light-base/src/sync_service.rs

@@ -40,7 +40,7 @@ use smoldot::{
     executor::host,
     libp2p::PeerId,
     network::{protocol, service},
-    trie::{self, prefix_proof, proof_verify},
+    trie::{self, prefix_proof, proof_decode},
 };
 
 mod parachain;
@@ -429,16 +429,20 @@ impl<TPlat: Platform> SyncService<TPlat> {
                 .map_err(StorageQueryErrorDetail::Network)
                 .and_then(|outcome| {
                     let decoded = outcome.decode();
+                    let decoded =
+                        proof_decode::decode_and_verify_proof(proof_decode::VerifyProofConfig {
+                            proof: decoded.iter().copied(),
+                            trie_root_hash: &storage_trie_root,
+                        })
+                        .map_err(StorageQueryErrorDetail::ProofVerification)?;
+
                     let mut result = Vec::with_capacity(requested_keys.clone().count());
                     for key in requested_keys.clone() {
                         result.push(
-                            proof_verify::verify_proof(proof_verify::VerifyProofConfig {
-                                proof: decoded.iter().map(|nv| &nv[..]),
-                                requested_key: key.as_ref(),
-                                trie_root_hash: &storage_trie_root,
-                            })
-                            .map_err(StorageQueryErrorDetail::ProofVerification)?
-                            .map(|v| v.to_owned()),
+                            decoded
+                                .storage_value(key.as_ref())
+                                .ok_or(StorageQueryErrorDetail::MissingProofEntry)?
+                                .map(|v| v.to_owned()),
                         );
                     }
                     debug_assert_eq!(result.len(), result.capacity());
@@ -516,8 +520,14 @@ impl<TPlat: Platform> SyncService<TPlat> {
                             }
                             Err((scan, err)) => {
                                 prefix_scan = scan;
-                                outcome_errors
-                                    .push(StorageQueryErrorDetail::ProofVerification(err));
+                                outcome_errors.push(match err {
+                                    prefix_proof::Error::InvalidProof(err) => {
+                                        StorageQueryErrorDetail::ProofVerification(err)
+                                    }
+                                    prefix_proof::Error::MissingProofEntry => {
+                                        StorageQueryErrorDetail::MissingProofEntry
+                                    }
+                                });
                             }
                         }
                     }
@@ -616,10 +626,11 @@ impl StorageQueryError {
                 ),
             ) => false,
             // TODO: as a temporary hack, we consider `TrieRootNotFound` as the remote not knowing about the requested block; see https://github.com/paritytech/substrate/pull/8046
-            StorageQueryErrorDetail::ProofVerification(proof_verify::Error::TrieRootNotFound) => {
+            StorageQueryErrorDetail::ProofVerification(proof_decode::Error::TrieRootNotFound) => {
                 true
             }
-            StorageQueryErrorDetail::ProofVerification(_) => false,
+            StorageQueryErrorDetail::ProofVerification(_)
+            | StorageQueryErrorDetail::MissingProofEntry => false,
         })
     }
 }
@@ -646,7 +657,9 @@ pub enum StorageQueryErrorDetail {
     Network(network_service::StorageProofRequestError),
     /// Error verifying the proof.
     #[display(fmt = "{}", _0)]
-    ProofVerification(proof_verify::Error),
+    ProofVerification(proof_decode::Error),
+    /// Proof is missing one or more desired storage items.
+    MissingProofEntry,
 }
 
 /// Error that can happen when calling [`SyncService::call_proof_query`].

bin/light-base/src/sync_service/standalone.rs

@@ -33,7 +33,7 @@ use smoldot::{
     libp2p,
     network::{self, protocol},
     sync::all,
-    trie::proof_verify,
+    trie::proof_decode,
 };
 
 /// Starts a sync service background task to synchronize a standalone chain (relay chain or not).
@@ -579,17 +579,23 @@ impl<TPlat: Platform> Task<TPlat> {
                         // TODO: lots of copying around
                         // TODO: log what happens
                         let decoded = outcome.decode();
-                        keys.iter()
-                            .map(|key| {
-                                proof_verify::verify_proof(proof_verify::VerifyProofConfig {
-                                    proof: decoded.iter().map(|nv| &nv[..]),
-                                    requested_key: key.as_ref(),
-                                    trie_root_hash: &state_trie_root,
-                                })
-                                .map_err(|_| ())
-                                .map(|v| v.map(|v| v.to_vec()))
+                        if let Ok(decoded) =
+                            proof_decode::decode_and_verify_proof(proof_decode::VerifyProofConfig {
+                                proof: decoded.iter().copied(),
+                                trie_root_hash: &state_trie_root,
                             })
-                            .collect::<Result<Vec<_>, ()>>()
+                        {
+                            keys.iter()
+                                .map(|key| {
+                                    decoded
+                                        .storage_value(key)
+                                        .ok_or(())
+                                        .map(|v| v.map(|v| v.to_vec()))
+                                })
+                                .collect::<Result<Vec<_>, ()>>()
+                        } else {
+                            Err(())
+                        }
                     } else {
                         Err(())
                     }

bin/wasm-node/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Changed
 
 - In earlier versions of smoldot, `setTimeout(callback, 0)` was frequently used in order split execution of CPU-intensive tasks in multiple smaller ones while still giving back control to the execution environment (such as NodeJS or the browser). Unfortunately, when a web page is in the background, browsers set a minimum delay of one second for `setTimeout`. For this reason, the usage of ̀`setTimeout` has now been reduced to the strict minimum, except when the environment is browser and `document.visibilityState` is equal to `visible`. ([#2999](https://github.com/paritytech/smoldot/pull/2999))
+- Optimize the Merkle proof verification. The complexity has been reduced from `O(n^2)` to `O(n * log n)`.
 
 ## 0.7.7 - 2022-11-11
 

src/sync/all.rs

@@ -1600,8 +1600,7 @@ impl<TRq, TSrc, TBl> AllSync<TRq, TSrc, TBl> {
     /// Inject a response to a previously-emitted call proof request.
     ///
     /// On success, must contain the encoded Merkle proof. See the
-    /// [`trie`](crate::trie::proof_verify) module for a description of the format of Merkle
-    /// proofs.
+    /// [`trie`](crate::trie) module for a description of the format of Merkle proofs.
     ///
     /// # Panic
     ///