add(scan): Test the RegisterKeys scan service call

Cargo.lock

@@ -4240,6 +4240,25 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "strum"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "723b93e8addf9aa965ebe2d11da6d7540fa2283fcea14b3371ff055f7ba13f5f"
+
+[[package]]
+name = "strum_macros"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a3417fc93d76740d974a01654a09777cb500428cc874ca9f45edfe0c4d4cd18"
+dependencies = [
+ "heck 0.4.1",
+ "proc-macro2",
+ "quote",
+ "rustversion",
+ "syn 2.0.48",
+]
+
 [[package]]
 name = "subtle"
 version = "2.4.1"
@@ -5748,6 +5767,8 @@ dependencies = [
  "sha2",
  "spandoc",
  "static_assertions",
+ "strum",
+ "strum_macros",
  "thiserror",
  "tinyvec",
  "tokio",
@@ -5933,6 +5954,7 @@ dependencies = [
  "rand 0.8.5",
  "semver 1.0.21",
  "serde",
+ "strum",
  "tokio",
  "tower",
  "tracing",

zebra-chain/Cargo.toml

@@ -147,6 +147,10 @@ rand_chacha = { version = "0.3.1", optional = true }
 
 zebra-test = { path = "../zebra-test/", version = "1.0.0-beta.34", optional = true }
 
+# Helpful macros for working with enums and strings
+strum_macros = "0.26"
+strum = "0.26"
+
 [dev-dependencies]
 # Benchmarks
 criterion = { version = "0.5.1", features = ["html_reports"] }

zebra-chain/src/parameters/network.rs

@@ -12,6 +12,8 @@ use crate::{
 #[cfg(any(test, feature = "proptest-impl"))]
 use proptest_derive::Arbitrary;
 
+use strum_macros::EnumIter;
+
 #[cfg(test)]
 mod tests;
 
@@ -52,7 +54,7 @@ mod tests;
 const ZIP_212_GRACE_PERIOD_DURATION: HeightDiff = 32_256;
 
 /// An enum describing the possible network choices.
-#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash, Serialize, Deserialize)]
+#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash, Serialize, Deserialize, EnumIter)]
 #[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 pub enum Network {
     /// The production mainnet.

zebra-chain/src/primitives/viewing_key/sapling.rs

@@ -74,7 +74,7 @@ impl SaplingViewingKey {
 impl Network {
     /// Returns the human-readable prefix for an Zcash Sapling extended full viewing key
     /// for this network.
-    fn sapling_efvk_hrp(&self) -> &'static str {
+    pub fn sapling_efvk_hrp(&self) -> &'static str {
         if self.is_a_test_network() {
             // Assume custom testnets have the same HRP
             //

zebra-scan/Cargo.toml

@@ -54,7 +54,7 @@ futures = "0.3.30"
 zcash_client_backend = "0.10.0-rc.1"
 zcash_primitives = "0.13.0-rc.1"
 
-zebra-chain = { path = "../zebra-chain", version = "1.0.0-beta.34" }
+zebra-chain = { path = "../zebra-chain", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-state = { path = "../zebra-state", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-node-services = { path = "../zebra-node-services", version = "1.0.0-beta.34", features = ["shielded-scan"] }
 zebra-grpc = { path = "../zebra-grpc", version = "0.1.0-alpha.1" }
@@ -75,6 +75,7 @@ zcash_note_encryption = { version = "0.4.0", optional = true }
 zebra-test = { path = "../zebra-test", version = "1.0.0-beta.34", optional = true }
 
 [dev-dependencies]
+strum = "0.26"
 
 insta = { version = "1.33.0", features = ["ron", "redactions"] }
 tokio = { version = "1.36.0", features = ["test-util"] }

zebra-scan/src/service/scan_task/scan.rs

@@ -21,7 +21,6 @@ use zcash_client_backend::{
     scanning::{ScanError, ScanningKey},
 };
 use zcash_primitives::{
-    constants::*,
     sapling::SaplingIvk,
     zip32::{AccountId, DiversifiableFullViewingKey, Scope},
 };
@@ -61,8 +60,9 @@ const INITIAL_WAIT: Duration = Duration::from_secs(15);
 
 /// The amount of time between checking for new blocks and starting new scans.
 ///
-/// This is just under half the target block interval.
-pub const CHECK_INTERVAL: Duration = Duration::from_secs(30);
+/// TODO: The current value is set to 10 so that tests don't sleep for too long and finish faster.
+/// Set it to 30 after #8250 gets addressed or remove this const completely in the refactor.
+pub const CHECK_INTERVAL: Duration = Duration::from_secs(10);
 
 /// We log an info log with progress after this many blocks.
 const INFO_LOG_INTERVAL: u32 = 10_000;
@@ -81,6 +81,7 @@ pub async fn start(
     info!(?network, "starting scan task");
 
     // Do not scan and notify if we are below sapling activation height.
+    #[cfg(not(test))]
     wait_for_height(
         sapling_activation_height,
         "Sapling activation",
@@ -405,19 +406,11 @@ pub fn scan_block<K: ScanningKey>(
 //       performance: stop returning both the dfvk and ivk for the same key
 // TODO: use `ViewingKey::parse` from zebra-chain instead
 pub fn sapling_key_to_scan_block_keys(
-    sapling_key: &SaplingScanningKey,
+    key: &SaplingScanningKey,
     network: Network,
 ) -> Result<(Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>), Report> {
-    let hrp = if network.is_a_test_network() {
-        // Assume custom testnets have the same HRP
-        //
-        // TODO: add the regtest HRP here
-        testnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY
-    } else {
-        mainnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY
-    };
-
-    let efvk = decode_extended_full_viewing_key(hrp, sapling_key).map_err(|e| eyre!(e))?;
+    let efvk =
+        decode_extended_full_viewing_key(network.sapling_efvk_hrp(), key).map_err(|e| eyre!(e))?;
 
     // Just return all the keys for now, so we can be sure our code supports them.
     let dfvk = efvk.to_diversifiable_full_viewing_key();

zebra-scan/src/service/scan_task/tests/vectors.rs

@@ -18,7 +18,8 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Send some keys to be registered
     let num_keys = 10;
-    let sapling_keys = mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"));
+    let sapling_keys =
+        mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"), network);
     let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> =
         sapling_keys.into_iter().zip((0..).map(Some)).collect();
     mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
@@ -60,7 +61,7 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that keys can't be overridden.
 
-    let sapling_keys = mock_sapling_scanning_keys(20);
+    let sapling_keys = mock_sapling_scanning_keys(20, network);
     let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> = sapling_keys
         .clone()
         .into_iter()
@@ -87,7 +88,7 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that it removes keys correctly
 
-    let sapling_keys = mock_sapling_scanning_keys(30);
+    let sapling_keys = mock_sapling_scanning_keys(30, network);
     let done_rx = mock_scan_task.remove_keys(&sapling_keys)?;
 
     let (new_keys, _new_results_senders) =

zebra-scan/src/service/tests.rs

@@ -1,7 +1,8 @@
 //! Tests for ScanService.
 
+use strum::IntoEnumIterator;
 use tokio::sync::mpsc::error::TryRecvError;
-use tower::{Service, ServiceExt};
+use tower::{Service, ServiceBuilder, ServiceExt};
 
 use color_eyre::{eyre::eyre, Result};
 
@@ -12,7 +13,8 @@ use zebra_state::TransactionIndex;
 use crate::{
     service::{scan_task::ScanTaskCommand, ScanService},
     storage::db::tests::{fake_sapling_results, new_test_storage},
-    tests::ZECPAGES_SAPLING_VIEWING_KEY,
+    tests::{mock_sapling_scanning_keys, ZECPAGES_SAPLING_VIEWING_KEY},
+    Config,
 };
 
 /// Tests that keys are deleted correctly
@@ -254,3 +256,70 @@ pub async fn scan_service_get_results_for_key_correctly() -> Result<()> {
 
     Ok(())
 }
+
+/// Tests that the scan service registers keys correctly.
+#[tokio::test]
+pub async fn scan_service_registers_keys_correctly() -> Result<()> {
+    for network in Network::iter() {
+        scan_service_registers_keys_correctly_for(network).await?;
+    }
+
+    Ok(())
+}
+
+async fn scan_service_registers_keys_correctly_for(network: Network) -> Result<()> {
+    // Mock the state.
+    let (state, _, _, chain_tip_change) = zebra_state::populated_state(vec![], network).await;
+
+    // Instantiate the scan service.
+    let mut scan_service = ServiceBuilder::new()
+        .buffer(2)
+        .service(ScanService::new(&Config::ephemeral(), network, state, chain_tip_change).await);
+
+    // Mock three Sapling keys.
+    let mocked_keys = mock_sapling_scanning_keys(3, network);
+
+    // Add birth heights to the mocked keys.
+    let keys_to_register: Vec<_> = mocked_keys
+        .clone()
+        .into_iter()
+        .zip((0u32..).map(Some))
+        .collect();
+
+    // Register the first key.
+    match scan_service
+        .ready()
+        .await
+        .map_err(|err| eyre!(err))?
+        .call(Request::RegisterKeys(keys_to_register[..1].to_vec()))
+        .await
+        .map_err(|err| eyre!(err))?
+    {
+        Response::RegisteredKeys(registered_keys) => {
+            // The key should be registered.
+            assert_eq!(registered_keys, mocked_keys[..1]);
+        }
+
+        _ => panic!("scan service should have responded with the `RegisteredKeys` response"),
+    }
+
+    // Try registering all three keys.
+    match scan_service
+        .ready()
+        .await
+        .map_err(|err| eyre!(err))?
+        .call(Request::RegisterKeys(keys_to_register))
+        .await
+        .map_err(|err| eyre!(err))?
+    {
+        Response::RegisteredKeys(registered_keys) => {
+            // Only the last two keys should be registered in this service call since the first one
+            // was registered in the previous call.
+            assert_eq!(registered_keys, mocked_keys[1..3]);
+        }
+
+        _ => panic!("scan service should have responded with the `RegisteredKeys` response"),
+    }
+
+    Ok(())
+}

zebra-scan/src/tests.rs

@@ -37,6 +37,7 @@ use zebra_chain::{
     amount::{Amount, NegativeAllowed},
     block::{self, merkle, Block, Header, Height},
     fmt::HexDebug,
+    parameters::Network,
     primitives::{redjubjub, Groth16Proof},
     sapling::{self, PerSpendAnchor, Spend, TransferData},
     serialization::AtLeastOne,
@@ -55,16 +56,16 @@ pub const ZECPAGES_SAPLING_VIEWING_KEY: &str = "zxviews1q0duytgcqqqqpqre26wkl45g
 /// A fake viewing key in an incorrect format.
 pub const FAKE_SAPLING_VIEWING_KEY: &str = "zxviewsfake";
 
-/// Generates `num_keys` of [`SaplingScanningKey`]s for tests.
+/// Generates `num_keys` of [`SaplingScanningKey`]s for tests for the given [`Network`].
 ///
 /// The keys are seeded only from their index in the returned `Vec`, so repeated calls return same
 /// keys at a particular index.
-pub fn mock_sapling_scanning_keys(num_keys: u8) -> Vec<SaplingScanningKey> {
+pub fn mock_sapling_scanning_keys(num_keys: u8, network: Network) -> Vec<SaplingScanningKey> {
     let mut keys: Vec<SaplingScanningKey> = vec![];
 
     for seed in 0..num_keys {
         keys.push(encode_extended_full_viewing_key(
-            "zxviews",
+            network.sapling_efvk_hrp(),
             &mock_sapling_efvk(&[seed]),
         ));
     }