state sync proof

core/api/service/state/impl/state_api_impl.cpp

@@ -15,6 +15,7 @@
 #include "common/hexutil.hpp"
 #include "common/monadic_utils.hpp"
 #include "runtime/executor.hpp"
+#include "storage/trie/on_read.hpp"
 
 OUTCOME_CPP_DEFINE_CATEGORY(kagome::api, StateApiImpl::Error, e) {
   using E = kagome::api::StateApiImpl::Error;
@@ -193,15 +194,14 @@ namespace kagome::api {
       std::optional<primitives::BlockHash> opt_at) const {
     auto at =
         opt_at.has_value() ? opt_at.value() : block_tree_->bestBlock().hash;
-    std::unordered_set<common::Buffer> proof;
-    auto prove = [&](common::BufferView raw) { proof.emplace(raw); };
+    storage::trie::OnRead db;
     OUTCOME_TRY(header, header_repo_->getBlockHeader(at));
-    OUTCOME_TRY(trie,
-                storage_->getProofReaderBatchAt(header.state_root, prove));
+    OUTCOME_TRY(
+        trie, storage_->getProofReaderBatchAt(header.state_root, db.onRead()));
     for (auto &key : keys) {
       OUTCOME_TRY(trie->tryGet(key));
     }
-    return ReadProof{at, {proof.begin(), proof.end()}};
+    return ReadProof{at, db.vec()};
   }
 
   outcome::result<primitives::Version> StateApiImpl::getRuntimeVersion(

core/blockchain/impl/cached_tree.cpp

@@ -153,7 +153,7 @@ namespace kagome::blockchain {
     handle(subtree_root_node);
   }
 
-  inline TreeMeta::Weight TreeMeta::getWeight(
+  TreeMeta::Weight TreeMeta::getWeight(
       std::shared_ptr<TreeNode> node) const {
     auto finalized = last_finalized.lock();
     BOOST_ASSERT(finalized);

core/common/buffer_view.hpp

@@ -20,13 +20,15 @@ namespace kagome::common {
 namespace kagome::common {
 
   class BufferView : public gsl::span<const uint8_t> {
-    using Span = gsl::span<const uint8_t>;
-
    public:
-    using Span::Span;
-    using Span::operator=;
+    using span::span;
+
+    BufferView(const span &other) noexcept : span(other) {}
 
-    BufferView(const Span &other) noexcept : Span(other) {}
+    template <typename T>
+    decltype(auto) operator=(T &&t) {
+      return span::operator=(std::forward<T>(t));
+    }
 
     std::string toHex() const {
       return hex_lower(*this);

core/network/impl/protocols/light.cpp

@@ -10,6 +10,7 @@
 #include "runtime/common/trie_storage_provider_impl.hpp"
 #include "runtime/executor.hpp"
 #include "runtime/module_repository.hpp"
+#include "storage/trie/on_read.hpp"
 
 namespace kagome::network {
   LightProtocol::LightProtocol(
@@ -35,11 +36,11 @@ namespace kagome::network {
 
   std::optional<outcome::result<LightProtocol::ResponseType>>
   LightProtocol::onRxRequest(RequestType req, std::shared_ptr<Stream>) {
-    std::unordered_set<common::Buffer> proof;
-    auto prove = [&](common::BufferView raw) { proof.emplace(raw); };
+    storage::trie::OnRead proof;
     OUTCOME_TRY(header, repository_->getBlockHeader(req.block));
-    OUTCOME_TRY(batch,
-                storage_->getProofReaderBatchAt(header.state_root, prove));
+    OUTCOME_TRY(
+        batch,
+        storage_->getProofReaderBatchAt(header.state_root, proof.onRead()));
     auto call = boost::get<LightProtocolRequest::Call>(&req.op);
     if (call) {
       OUTCOME_TRY(instance,
@@ -65,7 +66,6 @@ namespace kagome::network {
         OUTCOME_TRY(trie.get().tryGet(key));
       }
     }
-    return {
-        LightProtocolResponse{{proof.begin(), proof.end()}, call != nullptr}};
+    return LightProtocolResponse{proof.vec(), call != nullptr};
   }
 }  // namespace kagome::network

core/network/impl/state_protocol_observer_impl.cpp

@@ -7,14 +7,35 @@
 #include <boost/algorithm/string/predicate.hpp>
 #include <boost/bind/storage.hpp>
 #include <libp2p/outcome/outcome.hpp>
+#include <unordered_set>
 
 #include "blockchain/block_header_repository.hpp"
 #include "common/buffer.hpp"
 #include "network/types/state_response.hpp"
 #include "storage/predefined_keys.hpp"
+#include "storage/trie/compact_encode.hpp"
 #include "storage/trie/trie_storage.hpp"
 #include "storage/trie/types.hpp"
 
+/*
+Example:
+  /sync/2 proof=true
+  trie = {
+    "": value1,
+    "key": value2,
+    "key_child": {
+      "": value3,
+      "key": value4,
+    },
+  }
+  request([]) = [value1, value2, value3, value4]
+  request([""]) = [value1, value2, value3, value4]
+  request(["key"]) = [value2, value3, value4]
+  request(["key_child"]) = []
+  request(["key_child", ""]) = [value3, value4]
+  request(["key_child", "key"]) = [value4]
+*/
+
 constexpr unsigned MAX_RESPONSE_BYTES = 2 * 1024 * 1024;
 
 OUTCOME_CPP_DEFINE_CATEGORY(kagome::network,
@@ -26,6 +47,8 @@ OUTCOME_CPP_DEFINE_CATEGORY(kagome::network,
       return "Expected child root hash prefix.";
     case E::NOTFOUND_CHILD_ROOTHASH:
       return "Child storage root hash not found.";
+    case E::VALUE_NOT_FOUND:
+      return "Value not found";
   }
   return "unknown error";
 }
@@ -79,7 +102,19 @@ namespace kagome::network {
 
   outcome::result<network::StateResponse>
   StateProtocolObserverImpl::onStateRequest(const StateRequest &request) const {
+    if (request.start.size() > 2) {
+      return Error::INVALID_CHILD_ROOTHASH;
+    }
+    if (request.start.size() == 2
+        and not boost::starts_with(request.start[0],
+                                   storage::kChildStoragePrefix)) {
+      return Error::INVALID_CHILD_ROOTHASH;
+    }
     OUTCOME_TRY(header, blocks_headers_->getBlockHeader(request.hash));
+    if (not request.no_proof) {
+      OUTCOME_TRY(proof, prove(header.state_root, request.start));
+      return StateResponse{{}, proof};
+    }
     OUTCOME_TRY(batch, storage_->getEphemeralBatchAt(header.state_root));
 
     auto cursor = batch->trieCursor();
@@ -97,10 +132,6 @@ namespace kagome::network {
     // Second key is child state storage key
     if (request.start.size() == 2) {
       const auto &parent_key = request.start[0];
-      const auto &child_prefix = storage::kChildStorageDefaultPrefix;
-      if (!boost::starts_with(parent_key, child_prefix)) {
-        return Error::INVALID_CHILD_ROOTHASH;
-      }
       if (auto value_res = batch->tryGet(parent_key);
           value_res.has_value() && value_res.value().has_value()) {
         OUTCOME_TRY(hash,
@@ -150,4 +181,81 @@ namespace kagome::network {
 
     return response;
   }
+
+  outcome::result<common::Buffer> StateProtocolObserverImpl::prove(
+      const common::Hash256 &root,
+      const std::vector<common::Buffer> &keys) const {
+    storage::trie::OnRead db;
+    auto get_batch = [&](const common::Hash256 &root) {
+      return storage_->getProofReaderBatchAt(root, db.onRead());
+    };
+
+    std::vector<std::unique_ptr<storage::BufferStorageCursor>> stack;
+    OUTCOME_TRY(batch, get_batch(root));
+    stack.emplace_back(batch->cursor());
+
+    std::unordered_set<common::Hash256> child_roots;
+    common::BufferView seek;
+    if (keys.size() == 2) {
+      auto &cursor = stack.back();
+      OUTCOME_TRY(cursor->seek(keys[0]));
+      if (cursor->key() != keys[0]) {
+        return Error::NOTFOUND_CHILD_ROOTHASH;
+      }
+      auto value = cursor->value();
+      if (not value) {
+        return Error::VALUE_NOT_FOUND;
+      }
+      OUTCOME_TRY(root, common::Hash256::fromSpan(*value));
+      OUTCOME_TRY(batch, get_batch(root));
+      child_roots.emplace(root);
+      stack.emplace_back(batch->cursor());
+      seek = keys[1];
+    } else if (not keys.empty()) {
+      seek = keys[0];
+    }
+    auto &cursor = stack.back();
+    OUTCOME_TRY(cursor->seek(seek));
+    if (cursor->key() == seek) {
+      if (not cursor->value()) {
+        return Error::VALUE_NOT_FOUND;
+      }
+      OUTCOME_TRY(cursor->next());
+    }
+
+    while (not stack.empty()) {
+      auto pop = true;
+      auto &cursor = stack.back();
+      while (cursor->isValid()) {
+        auto value = cursor->value();
+        if (not value) {
+          return Error::VALUE_NOT_FOUND;
+        }
+        auto key = cursor->key().value();
+        if (stack.size() == 1
+            and boost::starts_with(key, storage::kChildStoragePrefix)) {
+          OUTCOME_TRY(root, common::Hash256::fromSpan(*value));
+          if (child_roots.emplace(root).second) {
+            OUTCOME_TRY(batch, get_batch(root));
+            stack.emplace_back(batch->cursor());
+            OUTCOME_TRY(stack.back()->seekFirst());
+            pop = false;
+            break;
+          }
+        } else if (db.size >= MAX_RESPONSE_BYTES) {
+          stack.clear();
+          pop = false;
+          break;
+        }
+        OUTCOME_TRY(cursor->next());
+      }
+      if (pop) {
+        stack.pop_back();
+        if (not stack.empty()) {
+          OUTCOME_TRY(stack.back()->next());
+        }
+      }
+    }
+    return storage::trie::compactEncode(db, root);
+  }
 }  // namespace kagome::network

core/network/impl/state_protocol_observer_impl.hpp

@@ -28,7 +28,11 @@ namespace kagome::network {
       : public StateProtocolObserver,
         public std::enable_shared_from_this<StateProtocolObserverImpl> {
    public:
-    enum class Error { INVALID_CHILD_ROOTHASH = 1, NOTFOUND_CHILD_ROOTHASH };
+    enum class Error {
+      INVALID_CHILD_ROOTHASH = 1,
+      NOTFOUND_CHILD_ROOTHASH,
+      VALUE_NOT_FOUND,
+    };
 
     StateProtocolObserverImpl(
         std::shared_ptr<blockchain::BlockHeaderRepository> blocks_headers,
@@ -45,6 +49,10 @@ namespace kagome::network {
         const common::Buffer &key,
         size_t limit) const;
 
+    outcome::result<common::Buffer> prove(
+        const common::Hash256 &root,
+        const std::vector<common::Buffer> &keys) const;
+
     std::shared_ptr<blockchain::BlockHeaderRepository> blocks_headers_;
     std::shared_ptr<storage::trie::TrieStorage> storage_;
     log::Logger log_;