Implement shard tree sync process

components/brave_wallet/browser/BUILD.gn

@@ -314,12 +314,18 @@ static_library("browser") {
 
   if (enable_orchard) {
     sources += [
+      "zcash/zcash_blocks_batch_scan_task.cc",
+      "zcash/zcash_blocks_batch_scan_task.h",
       "zcash/zcash_create_shield_transaction_task.cc",
       "zcash/zcash_create_shield_transaction_task.h",
       "zcash/zcash_get_zcash_chain_tip_status_task.cc",
       "zcash/zcash_get_zcash_chain_tip_status_task.h",
+      "zcash/zcash_scan_blocks_task.cc",
+      "zcash/zcash_scan_blocks_task.h",
       "zcash/zcash_shield_sync_service.cc",
       "zcash/zcash_shield_sync_service.h",
+      "zcash/zcash_verify_chain_state_task.cc",
+      "zcash/zcash_verify_chain_state_task.h",
     ]
 
     deps += [

components/brave_wallet/browser/test/BUILD.gn

@@ -163,7 +163,10 @@ source_set("brave_wallet_unit_tests") {
       "//brave/components/brave_wallet/browser/internal/orchard_bundle_manager_unittest.cc",
       "//brave/components/brave_wallet/browser/internal/orchard_storage/orchard_storage_unittest.cc",
       "//brave/components/brave_wallet/browser/internal/orchard_sync_state_unittest.cc",
+      "//brave/components/brave_wallet/browser/zcash/zcash_blocks_batch_scan_task_unittest.cc",
+      "//brave/components/brave_wallet/browser/zcash/zcash_scan_blocks_task_unittest.cc",
       "//brave/components/brave_wallet/browser/zcash/zcash_shield_sync_service_unittest.cc",
+      "//brave/components/brave_wallet/browser/zcash/zcash_verify_chain_state_task_unittest.cc",
     ]
     deps += [
       "//brave/components/brave_wallet/browser/internal:orchard_bundle",

components/brave_wallet/browser/zcash/zcash_blocks_batch_scan_task.cc

@@ -0,0 +1,257 @@
+// Copyright (c) 2024 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#include "brave/components/brave_wallet/browser/zcash/zcash_blocks_batch_scan_task.h"
+
+#include <algorithm>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "base/containers/extend.h"
+#include "brave/components/brave_wallet/browser/zcash/zcash_wallet_service.h"
+#include "brave/components/brave_wallet/common/hex_utils.h"
+
+namespace brave_wallet {
+
+namespace {
+constexpr uint32_t kBlockDownloadBatchSize = 10u;
+}
+
+ZCashBlocksBatchScanTask::ZCashBlocksBatchScanTask(
+    ZCashActionContext& context,
+    ZCashShieldSyncService::OrchardBlockScannerProxy& scanner,
+    uint32_t from,
+    uint32_t to,
+    ZCashBlocksBatchScanTaskCallback callback)
+    : context_(context),
+      scanner_(scanner),
+      from_(from),
+      to_(to),
+      callback_(std::move(callback)) {
+  CHECK_GT(from, kNu5BlockUpdate);
+  CHECK_GE(to, from);
+  frontier_block_height_ = from - 1;
+}
+
+ZCashBlocksBatchScanTask::~ZCashBlocksBatchScanTask() = default;
+
+void ZCashBlocksBatchScanTask::Start() {
+  CHECK(!started_);
+  started_ = true;
+  ScheduleWorkOnTask();
+}
+
+void ZCashBlocksBatchScanTask::ScheduleWorkOnTask() {
+  base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
+      FROM_HERE, base::BindOnce(&ZCashBlocksBatchScanTask::WorkOnTask,
+                                weak_ptr_factory_.GetWeakPtr()));
+}
+
+void ZCashBlocksBatchScanTask::WorkOnTask() {
+  if (error_) {
+    std::move(callback_).Run(base::unexpected(*error_));
+    return;
+  }
+
+  if (!frontier_tree_state_) {
+    GetFrontierTreeState();
+    return;
+  }
+
+  if (!frontier_block_) {
+    GetFrontierBlock();
+    return;
+  }
+
+  if (!scan_result_ && (!downloaded_blocks_ ||
+                        downloaded_blocks_->size() != (to_ - from_ + 1))) {
+    DownloadBlocks();
+    return;
+  }
+
+  if (!scan_result_) {
+    ScanBlocks();
+    return;
+  }
+
+  if (!database_updated_) {
+    UpdateDatabase();
+    return;
+  }
+
+  std::move(callback_).Run(true);
+}
+
+void ZCashBlocksBatchScanTask::GetFrontierTreeState() {
+  auto block_id = zcash::mojom::BlockID::New(frontier_block_height_,
+                                             std::vector<uint8_t>({}));
+  context_->zcash_rpc->GetTreeState(
+      context_->chain_id, std::move(block_id),
+      base::BindOnce(&ZCashBlocksBatchScanTask::OnGetFrontierTreeState,
+                     weak_ptr_factory_.GetWeakPtr()));
+}
+
+void ZCashBlocksBatchScanTask::OnGetFrontierTreeState(
+    base::expected<zcash::mojom::TreeStatePtr, std::string> result) {
+  if (!result.has_value() || !result.value()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kFailedToReceiveTreeState,
+        base::StrCat({"Frontier tree state failed, ", result.error()})};
+    ScheduleWorkOnTask();
+    return;
+  }
+  frontier_tree_state_ = result.value().Clone();
+  ScheduleWorkOnTask();
+}
+
+void ZCashBlocksBatchScanTask::GetFrontierBlock() {
+  context_->zcash_rpc->GetCompactBlocks(
+      context_->chain_id, frontier_block_height_, frontier_block_height_,
+      base::BindOnce(&ZCashBlocksBatchScanTask::OnGetFrontierBlock,
+                     weak_ptr_factory_.GetWeakPtr()));
+}
+
+void ZCashBlocksBatchScanTask::OnGetFrontierBlock(
+    base::expected<std::vector<zcash::mojom::CompactBlockPtr>, std::string>
+        result) {
+  if (!result.has_value() || result.value().size() != 1) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kFailedToDownloadBlocks,
+        result.error()};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  frontier_block_ = std::move(result.value()[0]);
+  ScheduleWorkOnTask();
+}
+
+void ZCashBlocksBatchScanTask::DownloadBlocks() {
+  uint32_t start_index =
+      downloaded_blocks_ ? from_ + downloaded_blocks_->size() : from_;
+  uint32_t end_index = std::min(start_index + kBlockDownloadBatchSize - 1, to_);
+  auto expected_size = end_index - start_index + 1;
+
+  context_->zcash_rpc->GetCompactBlocks(
+      context_->chain_id, start_index, end_index,
+      base::BindOnce(&ZCashBlocksBatchScanTask::OnBlocksDownloaded,
+                     weak_ptr_factory_.GetWeakPtr(), expected_size));
+}
+
+void ZCashBlocksBatchScanTask::OnBlocksDownloaded(
+    size_t expected_size,
+    base::expected<std::vector<zcash::mojom::CompactBlockPtr>, std::string>
+        result) {
+  CHECK(frontier_block_);
+  CHECK(frontier_tree_state_);
+  if (!result.has_value()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kFailedToDownloadBlocks,
+        result.error()};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  if (expected_size != result.value().size()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kFailedToDownloadBlocks,
+        "Expected block count doesn't match actual"};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  if (!downloaded_blocks_) {
+    downloaded_blocks_ = std::vector<zcash::mojom::CompactBlockPtr>();
+  }
+  base::Extend(downloaded_blocks_.value(), std::move(result.value()));
+  ScheduleWorkOnTask();
+}
+
+void ZCashBlocksBatchScanTask::ScanBlocks() {
+  if (!downloaded_blocks_ || downloaded_blocks_->empty()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kScannerError, "No blocks to scan"};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  if (!frontier_block_.value() || !frontier_tree_state_.value() ||
+      !frontier_block_.value()->chain_metadata) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kScannerError, "Frontier error"};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  OrchardTreeState tree_state;
+  {
+    auto frontier_bytes = PrefixedHexStringToBytes(
+        base::StrCat({"0x", frontier_tree_state_.value()->orchardTree}));
+
+    if (!frontier_bytes) {
+      error_ = ZCashShieldSyncService::Error{
+          ZCashShieldSyncService::ErrorCode::kScannerError,
+          "Failed to parse tree state"};
+      ScheduleWorkOnTask();
+      return;
+    }
+
+    tree_state.block_height = frontier_block_.value()->height;
+    tree_state.tree_size =
+        frontier_block_.value()->chain_metadata->orchard_commitment_tree_size;
+    tree_state.frontier = std::move(*frontier_bytes);
+  }
+
+  latest_scanned_block_ = downloaded_blocks_->back().Clone();
+
+  scanner_->ScanBlocks(
+      std::move(tree_state), std::move(downloaded_blocks_.value()),
+      base::BindOnce(&ZCashBlocksBatchScanTask::OnBlocksScanned,
+                     weak_ptr_factory_.GetWeakPtr()));
+}
+
+void ZCashBlocksBatchScanTask::OnBlocksScanned(
+    base::expected<OrchardBlockScanner::Result, OrchardBlockScanner::ErrorCode>
+        result) {
+  if (!result.has_value()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kScannerError,
+        "Failed to scan blocks"};
+    ScheduleWorkOnTask();
+    return;
+  }
+
+  scan_result_ = std::move(result.value());
+  ScheduleWorkOnTask();
+}
+
+void ZCashBlocksBatchScanTask::UpdateDatabase() {
+  CHECK(scan_result_.has_value());
+  CHECK(latest_scanned_block_.has_value());
+  auto latest_scanned_block_hash = ToHex((*latest_scanned_block_)->hash);
+  auto latest_scanned_block_height = (*latest_scanned_block_)->height;
+
+  context_->sync_state->AsyncCall(&OrchardSyncState::ApplyScanResults)
+      .WithArgs(context_->account_id.Clone(), std::move(scan_result_.value()),
+                latest_scanned_block_height, latest_scanned_block_hash)
+      .Then(base::BindOnce(&ZCashBlocksBatchScanTask::OnDatabaseUpdated,
+                           weak_ptr_factory_.GetWeakPtr()));
+}
+
+void ZCashBlocksBatchScanTask::OnDatabaseUpdated(
+    base::expected<OrchardStorage::Result, OrchardStorage::Error> result) {
+  if (!result.has_value()) {
+    error_ = ZCashShieldSyncService::Error{
+        ZCashShieldSyncService::ErrorCode::kFailedToUpdateDatabase,
+        result.error().message};
+    ScheduleWorkOnTask();
+    return;
+  }
+  database_updated_ = true;
+  ScheduleWorkOnTask();
+}
+
+}  // namespace brave_wallet

components/brave_wallet/browser/zcash/zcash_blocks_batch_scan_task.h

@@ -0,0 +1,85 @@
+// Copyright (c) 2024 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#ifndef BRAVE_COMPONENTS_BRAVE_WALLET_BROWSER_ZCASH_ZCASH_BLOCKS_BATCH_SCAN_TASK_H_
+#define BRAVE_COMPONENTS_BRAVE_WALLET_BROWSER_ZCASH_ZCASH_BLOCKS_BATCH_SCAN_TASK_H_
+
+#include <string>
+#include <vector>
+
+#include "brave/components/brave_wallet/browser/zcash/zcash_shield_sync_service.h"
+#include "brave/components/brave_wallet/common/zcash_utils.h"
+
+namespace brave_wallet {
+
+// Scans single scan range.
+class ZCashBlocksBatchScanTask {
+ public:
+  using ZCashBlocksBatchScanTaskCallback = base::OnceCallback<void(
+      base::expected<bool, ZCashShieldSyncService::Error>)>;
+  ZCashBlocksBatchScanTask(
+      ZCashActionContext& context,
+      ZCashShieldSyncService::OrchardBlockScannerProxy& scanner,
+      uint32_t from,
+      uint32_t to,
+      ZCashBlocksBatchScanTaskCallback callback);
+  ~ZCashBlocksBatchScanTask();
+
+  uint32_t from() const { return from_; }
+
+  uint32_t to() const { return to_; }
+
+  void Start();
+
+ private:
+  void WorkOnTask();
+  void ScheduleWorkOnTask();
+
+  void GetFrontierTreeState();
+  void OnGetFrontierTreeState(
+      base::expected<zcash::mojom::TreeStatePtr, std::string>);
+  void GetFrontierBlock();
+  void OnGetFrontierBlock(
+      base::expected<std::vector<zcash::mojom::CompactBlockPtr>, std::string>
+          result);
+
+  void DownloadBlocks();
+  void OnBlocksDownloaded(
+      size_t expected_size,
+      base::expected<std::vector<zcash::mojom::CompactBlockPtr>, std::string>
+          result);
+
+  void ScanBlocks();
+  void OnBlocksScanned(base::expected<OrchardBlockScanner::Result,
+                                      OrchardBlockScanner::ErrorCode> result);
+
+  void UpdateDatabase();
+  void OnDatabaseUpdated(
+      base::expected<OrchardStorage::Result, OrchardStorage::Error> result);
+
+  raw_ref<ZCashActionContext> context_;
+  raw_ref<ZCashShieldSyncService::OrchardBlockScannerProxy> scanner_;
+  uint32_t from_ = 0;
+  uint32_t to_ = 0;
+  ZCashBlocksBatchScanTaskCallback callback_;
+
+  uint32_t frontier_block_height_ = 0;
+
+  std::optional<ZCashShieldSyncService::Error> error_;
+  std::optional<zcash::mojom::TreeStatePtr> frontier_tree_state_;
+  std::optional<zcash::mojom::CompactBlockPtr> frontier_block_;
+  std::optional<std::vector<zcash::mojom::CompactBlockPtr>> downloaded_blocks_;
+  std::optional<OrchardBlockScanner::Result> scan_result_;
+  std::optional<zcash::mojom::CompactBlockPtr> latest_scanned_block_;
+
+  bool started_ = false;
+  bool database_updated_ = false;
+
+  base::WeakPtrFactory<ZCashBlocksBatchScanTask> weak_ptr_factory_{this};
+};
+
+}  // namespace brave_wallet
+
+#endif  // BRAVE_COMPONENTS_BRAVE_WALLET_BROWSER_ZCASH_ZCASH_BLOCKS_BATCH_SCAN_TASK_H_