Fix prefix_proof misbehavior (#2947)

Fix #2946

The behavior of `prefix_proof` is currently buggy. When a node isn't
found in the proof that the remote sent us, we simply discard it instead
of adding it to the list of keys to query at the next iteration.
The consequence is that `state_getKeysPaged` was wrong most the time.

In order to fix this, I had to reorganize the entire function.

There's still an issue because we're asking for too much data, which
leads the remote to shutdown the connection, so queries still
occasionally fail. But this other issue is not the fault of the
`prefix_proof` module, and this PR is only about fixing `prefix_proof`.

bin/light-base/src/sync_service.rs

@@ -479,6 +479,7 @@ impl<TPlat: Platform> SyncService<TPlat> {
 
             // TODO: better peers selection ; don't just take the first
             // TODO: handle max_parallel
+            // TODO: is the number of keys is large, split into multiple requests
             for target in self
                 .peers_assumed_know_blocks(block_number, block_hash)
                 .await

bin/wasm-node/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Fix the `state_getKeysPaged` JSON-RPC function returning incorrect results in some situations. ([#2947](https://github.com/paritytech/smoldot/pull/2947))
+
 ## 0.7.4 - 2022-10-27
 
 ### Changed

src/trie/prefix_proof.rs

@@ -30,7 +30,7 @@
 use super::{nibble, proof_verify};
 
 use alloc::{vec, vec::Vec};
-use core::{fmt, iter};
+use core::{fmt, iter, mem};
 
 /// Configuration to pass to [`prefix_scan`].
 pub struct Config<'a> {
@@ -76,31 +76,45 @@ impl PrefixScan {
         mut self,
         proof: impl Iterator<Item = &'a [u8]> + Clone + 'a,
     ) -> Result<ResumeOutcome, (Self, proof_verify::Error)> {
+        let mut non_terminal_queries = mem::take(&mut self.next_queries);
+
         // The entire body is executed as long as verifying at least one proof succeeds.
         for is_first_iteration in iter::once(true).chain(iter::repeat(false)) {
             // Filled with the queries to perform at the next iteration.
             // Capacity assumes a maximum of 2 children per node on average. This value was chosen
             // completely arbitrarily.
-            let mut next = Vec::with_capacity(self.next_queries.len() * 2);
-
-            // True if any proof verification has succeeded during this iteration.
-            // Controls whether we continue iterating.
-            let mut any_successful_proof = false;
+            let mut next = Vec::with_capacity(non_terminal_queries.len() * 2);
+
+            debug_assert!(!non_terminal_queries.is_empty());
+            // TODO: iterating like that is very inefficient, as we verify the proof dozens of times
+            loop {
+                let query = match non_terminal_queries.pop() {
+                    Some(q) => q,
+                    None => break,
+                };
 
-            debug_assert!(!self.next_queries.is_empty());
-            for query in &self.next_queries {
                 let info = match proof_verify::trie_node_info(proof_verify::TrieNodeInfoConfig {
                     requested_key: query.iter().copied(),
                     trie_root_hash: &self.trie_root_hash,
                     proof: proof.clone(),
                 }) {
                     Ok(info) => info,
-                    Err(err) if is_first_iteration => return Err((self, err)),
-                    Err(_) => continue,
+                    Err(proof_verify::Error::MissingProofEntry { .. }) if !is_first_iteration => {
+                        // Node not in the proof. There's no point in adding this node to `next`
+                        // as we will fail again if we try to verify the proof again.
+                        // If `is_first_iteration`, it means that the proof is incorrect.
+                        self.next_queries.push(query);
+                        continue;
+                    }
+                    Err(err) => {
+                        // Push all the non-processed queries back to `next_queries` before
+                        // returning the error, so that we can try again.
+                        self.next_queries.push(query);
+                        self.next_queries.extend(non_terminal_queries);
+                        return Err((self, err));
+                    }
                 };
 
-                any_successful_proof = true;
-
                 if matches!(
                     info.storage_value,
                     proof_verify::StorageValue::Known(_)
@@ -119,31 +133,28 @@ impl PrefixScan {
                     self.final_result.push(key);
                 }
 
-                for child_nibble in info.children.next_nibbles() {
-                    let mut next_query = Vec::with_capacity(query.len() + 1);
-                    next_query.extend_from_slice(query);
-                    next_query.push(child_nibble);
-                    next.push(next_query);
-                }
-            }
-
-            // If we have failed to make any progress during this iteration, return `InProgress`.
-            if !any_successful_proof {
-                debug_assert!(next.is_empty());
-                // Errors are immediately returned if `is_first_iteration`.
-                debug_assert!(!is_first_iteration);
-                break;
+                // For each child of the node, put into `next` the key that goes towards this
+                // child.
+                next.extend(info.children.unfold_append_to_key(query));
             }
 
             // Finished when nothing more to request.
-            if next.is_empty() {
+            if next.is_empty() && self.next_queries.is_empty() {
                 return Ok(ResumeOutcome::Success {
                     keys: self.final_result,
                 });
             }
 
-            // Update `next_queries` for the next iteration.
-            self.next_queries = next;
+            // If we have failed to make any progress during this iteration, return `InProgress`.
+            if next.is_empty() {
+                debug_assert!(!self.next_queries.is_empty());
+                // Errors are immediately returned if `is_first_iteration`.
+                debug_assert!(!is_first_iteration);
+                break;
+            }
+
+            // Update `non_terminal_queries` for the next iteration.
+            non_terminal_queries = next;
         }
 
         Ok(ResumeOutcome::InProgress(self))
@@ -167,3 +178,5 @@ pub enum ResumeOutcome {
         keys: Vec<Vec<u8>>,
     },
 }
+
+// TODO: needs tests

src/trie/proof_verify.rs

@@ -306,6 +306,30 @@ impl Children {
             ),
         }
     }
+
+    /// Iterators over all the children of the node. Returns an iterator producing one element per
+    /// child, where the element is `key` plus the nibble of this child.
+    pub fn unfold_append_to_key(
+        &self,
+        mut key: Vec<nibble::Nibble>,
+    ) -> impl Iterator<Item = Vec<nibble::Nibble>> {
+        match *self {
+            Children::None => either::Left(None.into_iter()),
+            Children::One(nibble) => {
+                key.push(nibble);
+                either::Left(Some(key).into_iter())
+            }
+            Children::Multiple { children_bitmap } => either::Right(
+                nibble::all_nibbles()
+                    .filter(move |n| (children_bitmap & (1 << u8::from(*n)) != 0))
+                    .map(move |nibble| {
+                        let mut k = key.clone();
+                        k.push(nibble);
+                        k
+                    }),
+            ),
+        }
+    }
 }
 
 /// Possible error returned by [`verify_proof`]