add(scan): Implement `RegisterKeys` service request to register new keys (#8251)

* Refactor obtaining of activation heights

* Impl the `RegisterKeys` service request

* Mock viewing keys for tests

* Refactor tests

* Apply suggestions from code review

Co-authored-by: Arya <aryasolhi@gmail.com>

* Remove a redundant comment

I don't think we need to assume the genesis block doesn't contain
shielded data as the comment says.

* Avoid using a single-letter variable

* Refactor mocking Sapling scanning keys

---------

Co-authored-by: Arya <aryasolhi@gmail.com>

zebra-chain/src/parameters/network.rs

@@ -129,6 +129,13 @@ impl Network {
     pub fn is_a_test_network(&self) -> bool {
         *self != Network::Mainnet
     }
+
+    /// Returns the Sapling activation height for this network.
+    pub fn sapling_activation_height(self) -> Height {
+        super::NetworkUpgrade::Sapling
+            .activation_height(self)
+            .expect("Sapling activation height needs to be set")
+    }
 }
 
 impl FromStr for Network {

zebra-node-services/src/scan_service/request.rs

@@ -14,8 +14,8 @@ pub enum Request {
     /// TODO: Accept `KeyHash`es and return key hashes that are registered
     CheckKeyHashes(Vec<()>),
 
-    /// TODO: Accept `ViewingKeyWithHash`es and return Ok(()) if successful or an error
+    /// Submits viewing keys with their optional birth-heights for scanning.
-    RegisterKeys(Vec<()>),
+    RegisterKeys(Vec<(String, Option<u32>)>),
 
     /// Deletes viewing keys and their results from the database.
     DeleteKeys(Vec<String>),

zebra-node-services/src/scan_service/response.rs

@@ -10,23 +10,28 @@ use zebra_chain::{block::Height, transaction::Hash};
 #[derive(Debug)]
 /// Response types for `zebra_scan::service::ScanService`
 pub enum Response {
-    /// Response to the `Info` request
+    /// Response to the [`Info`](super::request::Request::Info) request
     Info {
         /// The minimum sapling birthday height for the shielded scanner
         min_sapling_birthday_height: Height,
     },
 
-    /// Response to Results request
+    /// Response to [`RegisterKeys`](super::request::Request::RegisterKeys) request
+    ///
+    /// Contains the keys that the scanner accepted.
+    RegisteredKeys(Vec<String>),
+
+    /// Response to [`Results`](super::request::Request::Results) request
     ///
     /// We use the nested `BTreeMap` so we don't repeat any piece of response data.
     Results(BTreeMap<String, BTreeMap<Height, Vec<Hash>>>),
 
-    /// Response to DeleteKeys request
+    /// Response to [`DeleteKeys`](super::request::Request::DeleteKeys) request
     DeletedKeys,
 
-    /// Response to ClearResults request
+    /// Response to [`ClearResults`](super::request::Request::ClearResults) request
     ClearedResults,
 
-    /// Response to SubscribeResults request
+    /// Response to `SubscribeResults` request
     SubscribeResults(mpsc::Receiver<Arc<Hash>>),
 }

zebra-scan/src/service.rs

@@ -91,7 +91,7 @@ impl Service<Request> for ScanService {
 
                 return async move {
                     Ok(Response::Info {
-                        min_sapling_birthday_height: db.min_sapling_birthday_height(),
+                        min_sapling_birthday_height: db.network().sapling_activation_height(),
                     })
                 }
                 .boxed();
@@ -101,10 +101,15 @@ impl Service<Request> for ScanService {
                 // TODO: check that these entries exist in db
             }
 
-            Request::RegisterKeys(_viewing_key_with_hashes) => {
+            Request::RegisterKeys(keys) => {
-                // TODO:
+                let mut scan_task = self.scan_task.clone();
-                //  - add these keys as entries in db
+
-                //  - send new keys to scan task
+                return async move {
+                    Ok(Response::RegisteredKeys(
+                        scan_task.register_keys(keys)?.await?,
+                    ))
+                }
+                .boxed();
             }
 
             Request::DeleteKeys(keys) => {

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

@@ -12,9 +12,11 @@ use color_eyre::{eyre::eyre, Report};
 use tokio::sync::oneshot;
 
 use zcash_primitives::{sapling::SaplingIvk, zip32::DiversifiableFullViewingKey};
-use zebra_chain::{block::Height, transaction::Transaction};
+use zebra_chain::{block::Height, parameters::Network, transaction::Transaction};
 use zebra_state::SaplingScanningKey;
 
+use crate::scan::sapling_key_to_scan_block_keys;
+
 use super::ScanTask;
 
 #[derive(Debug)]
@@ -23,10 +25,9 @@ pub enum ScanTaskCommand {
     /// Start scanning for new viewing keys
     RegisterKeys {
         /// New keys to start scanning for
-        keys: HashMap<
+        keys: Vec<(String, Option<u32>)>,
-            SaplingScanningKey,
+        /// Returns the set of keys the scanner accepted.
-            (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>, Height),
+        rsp_tx: oneshot::Sender<Vec<SaplingScanningKey>>,
-        >,
     },
 
     /// Stop scanning for deleted viewing keys
@@ -61,11 +62,13 @@ impl ScanTask {
             SaplingScanningKey,
             (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>),
         >,
+        network: Network,
     ) -> Result<
         HashMap<SaplingScanningKey, (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>, Height)>,
         Report,
     > {
         let mut new_keys = HashMap::new();
+        let sapling_activation_height = network.sapling_activation_height();
 
         loop {
             let cmd = match cmd_receiver.try_recv() {
@@ -79,25 +82,44 @@ impl ScanTask {
             };
 
             match cmd {
-                ScanTaskCommand::RegisterKeys { keys } => {
+                ScanTaskCommand::RegisterKeys { keys, rsp_tx } => {
+                    // Determine what keys we pass to the scanner.
                     let keys: Vec<_> = keys
                         .into_iter()
-                        .filter(|(key, _)| {
+                        .filter_map(|key| {
-                            !parsed_keys.contains_key(key) || new_keys.contains_key(key)
+                            // Don't accept keys that:
+                            // 1. the scanner already has, and
+                            // 2. were already submitted.
+                            if parsed_keys.contains_key(&key.0) && !new_keys.contains_key(&key.0) {
+                                return None;
+                            }
+
+                            let birth_height = if let Some(height) = key.1 {
+                                match Height::try_from(height) {
+                                    Ok(height) => height,
+                                    // Don't accept the key if its birth height is not a valid height.
+                                    Err(_) => return None,
+                                }
+                            } else {
+                                // Use the Sapling activation height if the key has no birth height.
+                                sapling_activation_height
+                            };
+
+                            sapling_key_to_scan_block_keys(&key.0, network)
+                                .ok()
+                                .map(|parsed| (key.0, (parsed.0, parsed.1, birth_height)))
                         })
                         .collect();
 
-                    if !keys.is_empty() {
+                    // Send the accepted keys back.
-                        new_keys.extend(keys.clone());
+                    let _ = rsp_tx.send(keys.iter().map(|key| key.0.clone()).collect());
 
-                        let keys =
+                    new_keys.extend(keys.clone());
-                            keys.into_iter()
-                                .map(|(key, (decoded_dfvks, decoded_ivks, _h))| {
-                                    (key, (decoded_dfvks, decoded_ivks))
-                                });
 
-                        parsed_keys.extend(keys);
+                    parsed_keys.extend(
-                    }
+                        keys.into_iter()
+                            .map(|(key, (dfvks, ivks, _))| (key, (dfvks, ivks))),
+                    );
                 }
 
                 ScanTaskCommand::RemoveKeys { done_tx, keys } => {
@@ -143,11 +165,12 @@ impl ScanTask {
     /// Sends a message to the scan task to start scanning for the provided viewing keys.
     pub fn register_keys(
         &mut self,
-        keys: HashMap<
+        keys: Vec<(String, Option<u32>)>,
-            SaplingScanningKey,
+    ) -> Result<oneshot::Receiver<Vec<String>>, mpsc::SendError<ScanTaskCommand>> {
-            (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>, Height),
+        let (rsp_tx, rsp_rx) = oneshot::channel();
-        >,
+
-    ) -> Result<(), mpsc::SendError<ScanTaskCommand>> {
+        self.send(ScanTaskCommand::RegisterKeys { keys, rsp_tx })?;
-        self.send(ScanTaskCommand::RegisterKeys { keys })
+
+        Ok(rsp_rx)
     }
 }

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

@@ -75,7 +75,7 @@ pub async fn start(
     cmd_receiver: Receiver<ScanTaskCommand>,
 ) -> Result<(), Report> {
     let network = storage.network();
-    let sapling_activation_height = storage.min_sapling_birthday_height();
+    let sapling_activation_height = network.sapling_activation_height();
 
     // Do not scan and notify if we are below sapling activation height.
     wait_for_height(
@@ -125,7 +125,7 @@ pub async fn start(
             }
         }
 
-        let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+        let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
         // TODO: Check if the `start_height` is at or above the current height
         if !new_keys.is_empty() {

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

@@ -71,7 +71,7 @@ pub async fn scan_range(
     state: State,
     storage: Storage,
 ) -> Result<(), BoxError> {
-    let sapling_activation_height = storage.min_sapling_birthday_height();
+    let sapling_activation_height = storage.network().sapling_activation_height();
     // Do not scan and notify if we are below sapling activation height.
     wait_for_height(
         sapling_activation_height,

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

@@ -4,25 +4,23 @@ use std::collections::HashMap;
 
 use color_eyre::Report;
 
-use zebra_chain::block::Height;
+use crate::{service::ScanTask, tests::mock_sapling_scanning_keys};
-
-use crate::service::ScanTask;
 
 /// Test that [`ScanTask::process_messages`] adds and removes keys as expected for `RegisterKeys` and `DeleteKeys` command
 #[tokio::test]
 async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
     let (mut mock_scan_task, cmd_receiver) = ScanTask::mock();
     let mut parsed_keys = HashMap::new();
+    let network = Default::default();
 
     // Send some keys to be registered
     let num_keys = 10;
-    mock_scan_task.register_keys(
+    let sapling_keys = mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"));
-        (0..num_keys)
+    let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> =
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
+        sapling_keys.into_iter().zip((0..).map(Some)).collect();
-            .collect(),
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
-    )?;
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     // Check that it updated parsed_keys correctly and returned the right new keys when starting with an empty state
 
@@ -38,15 +36,11 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
         "should add all received keys to parsed keys"
     );
 
-    mock_scan_task.register_keys(
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
-        (0..num_keys)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
-            .collect(),
-    )?;
 
     // Check that no key should be added if they are all already known and the heights are the same
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     assert_eq!(
         parsed_keys.len(),
@@ -59,21 +53,19 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
         "should not return known keys as new keys"
     );
 
-    // Check that it returns the last seen start height for a key as the new key when receiving 2 register key messages
+    // Check that keys can't be overridden.
 
-    mock_scan_task.register_keys(
+    let sapling_keys = mock_sapling_scanning_keys(20);
-        (10..20)
+    let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> = sapling_keys
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
+        .clone()
-            .collect(),
+        .into_iter()
-    )?;
+        .map(|key| (key, Some(0)))
+        .collect();
 
-    mock_scan_task.register_keys(
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[10..20].to_vec())?;
-        (10..15)
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[10..15].to_vec())?;
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     assert_eq!(
         parsed_keys.len(),
@@ -87,22 +79,13 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
         "should add 10 of received keys to new keys"
     );
 
-    for (new_key, (_, _, start_height)) in new_keys {
-        if (10..15).contains(&new_key.parse::<i32>().expect("should parse into int")) {
-            assert_eq!(
-                start_height,
-                Height::MAX,
-                "these key heights should have been overwritten by the second message"
-            );
-        }
-    }
-
     // Check that it removes keys correctly
 
-    let done_rx =
+    let sapling_keys = mock_sapling_scanning_keys(30);
-        mock_scan_task.remove_keys(&(0..200).map(|i| i.to_string()).collect::<Vec<_>>())?;
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let done_rx = mock_scan_task.remove_keys(&sapling_keys)?;
+
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     // Check that it sends the done notification successfully before returning and dropping `done_tx`
     done_rx.await?;
@@ -116,15 +99,11 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that it doesn't return removed keys as new keys when processing a batch of messages
 
-    mock_scan_task.register_keys(
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
-        (0..200)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
 
-    mock_scan_task.remove_keys(&(0..200).map(|i| i.to_string()).collect::<Vec<_>>())?;
+    mock_scan_task.remove_keys(&sapling_keys)?;
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     assert!(
         new_keys.is_empty(),
@@ -133,21 +112,13 @@ async fn scan_task_processes_messages_correctly() -> Result<(), Report> {
 
     // Check that it does return registered keys if they were removed in a prior message when processing a batch of messages
 
-    mock_scan_task.register_keys(
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
-        (0..200)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
 
-    mock_scan_task.remove_keys(&(0..200).map(|i| i.to_string()).collect::<Vec<_>>())?;
+    mock_scan_task.remove_keys(&sapling_keys)?;
 
-    mock_scan_task.register_keys(
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[..2].to_vec())?;
-        (0..2)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
 
-    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys)?;
+    let new_keys = ScanTask::process_messages(&cmd_receiver, &mut parsed_keys, network)?;
 
     assert_eq!(
         new_keys.len(),

zebra-scan/src/storage.rs

@@ -2,10 +2,7 @@
 
 use std::collections::{BTreeMap, HashMap};
 
-use zebra_chain::{
+use zebra_chain::{block::Height, parameters::Network};
-    block::Height,
-    parameters::{Network, NetworkUpgrade},
-};
 use zebra_state::TransactionIndex;
 
 use crate::config::Config;
@@ -126,19 +123,4 @@ impl Storage {
     ) -> BTreeMap<Height, Vec<SaplingScannedResult>> {
         self.sapling_results_for_key(sapling_key)
     }
-
-    // Parameters
-
-    /// Returns the minimum sapling birthday height for the configured network.
-    pub fn min_sapling_birthday_height(&self) -> Height {
-        // Assume that the genesis block never contains shielded inputs or outputs.
-        //
-        // # Consensus
-        //
-        // For Zcash mainnet and the public testnet, Sapling activates above genesis,
-        // so this is always true.
-        NetworkUpgrade::Sapling
-            .activation_height(self.network())
-            .unwrap_or(Height(0))
-    }
 }

zebra-scan/src/storage/db/sapling.rs

@@ -213,7 +213,7 @@ impl Storage {
         sapling_key: &SaplingScanningKey,
         birthday_height: Option<Height>,
     ) {
-        let min_birthday_height = self.min_sapling_birthday_height();
+        let min_birthday_height = self.network().sapling_activation_height();
 
         // The birthday height must be at least the minimum height for that pool.
         let birthday_height = birthday_height

zebra-scan/src/tests.rs

@@ -12,8 +12,11 @@ use ff::{Field, PrimeField};
 use group::GroupEncoding;
 use rand::{rngs::OsRng, thread_rng, RngCore};
 
-use zcash_client_backend::proto::compact_formats::{
+use zcash_client_backend::{
-    ChainMetadata, CompactBlock, CompactSaplingOutput, CompactSaplingSpend, CompactTx,
+    encoding::encode_extended_full_viewing_key,
+    proto::compact_formats::{
+        ChainMetadata, CompactBlock, CompactSaplingOutput, CompactSaplingSpend, CompactTx,
+    },
 };
 use zcash_note_encryption::Domain;
 use zcash_primitives::{
@@ -27,7 +30,7 @@ use zcash_primitives::{
         value::NoteValue,
         Note, Nullifier,
     },
-    zip32::DiversifiableFullViewingKey,
+    zip32,
 };
 
 use zebra_chain::{
@@ -41,6 +44,7 @@ use zebra_chain::{
     transparent::{CoinbaseData, Input},
     work::{difficulty::CompactDifficulty, equihash::Solution},
 };
+use zebra_state::SaplingScanningKey;
 
 #[cfg(test)]
 mod vectors;
@@ -51,6 +55,31 @@ 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.
+///
+/// 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> {
+    let mut keys: Vec<SaplingScanningKey> = vec![];
+
+    for seed in 0..num_keys {
+        keys.push(encode_extended_full_viewing_key(
+            "zxviews",
+            &mock_sapling_efvk(&[seed]),
+        ));
+    }
+
+    keys
+}
+
+/// Generates an [`zip32::sapling::ExtendedFullViewingKey`] from `seed` for tests.
+#[allow(deprecated)]
+pub fn mock_sapling_efvk(seed: &[u8]) -> zip32::sapling::ExtendedFullViewingKey {
+    // TODO: Use `to_diversifiable_full_viewing_key` since `to_extended_full_viewing_key` is
+    // deprecated.
+    zip32::sapling::ExtendedSpendingKey::master(seed).to_extended_full_viewing_key()
+}
+
 /// Generates a fake block containing a Sapling output decryptable by `dfvk`.
 ///
 /// The fake block has the following transactions in this order:
@@ -61,7 +90,7 @@ pub const FAKE_SAPLING_VIEWING_KEY: &str = "zxviewsfake";
 pub fn fake_block(
     height: BlockHeight,
     nf: Nullifier,
-    dfvk: &DiversifiableFullViewingKey,
+    dfvk: &zip32::sapling::DiversifiableFullViewingKey,
     value: u64,
     tx_after: bool,
     initial_sapling_tree_size: Option<u32>,
@@ -135,7 +164,7 @@ pub fn fake_compact_block(
     height: BlockHeight,
     prev_hash: BlockHash,
     nf: Nullifier,
-    dfvk: &DiversifiableFullViewingKey,
+    dfvk: &zip32::sapling::DiversifiableFullViewingKey,
     value: u64,
     tx_after: bool,
     initial_sapling_tree_size: Option<u32>,

zebra-scan/src/tests/vectors.rs

@@ -5,7 +5,8 @@ use std::sync::Arc;
 use color_eyre::Result;
 
 use zcash_client_backend::{
-    encoding::decode_extended_full_viewing_key, proto::compact_formats::ChainMetadata,
+    encoding::{decode_extended_full_viewing_key, encode_extended_full_viewing_key},
+    proto::compact_formats::ChainMetadata,
 };
 use zcash_primitives::{
     constants::mainnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
@@ -24,7 +25,7 @@ use zebra_state::{SaplingScannedResult, TransactionIndex};
 use crate::{
     scan::{block_to_compact, scan_block},
     storage::db::tests::new_test_storage,
-    tests::{fake_block, ZECPAGES_SAPLING_VIEWING_KEY},
+    tests::{fake_block, mock_sapling_efvk, ZECPAGES_SAPLING_VIEWING_KEY},
 };
 
 /// This test:
@@ -146,31 +147,30 @@ async fn scanning_zecpages_from_populated_zebra_state() -> Result<()> {
 ///
 /// The purpose of this test is to check if our database and our scanning code are compatible.
 #[test]
-#[allow(deprecated)]
 fn scanning_fake_blocks_store_key_and_results() -> Result<()> {
+    let network = Network::Mainnet;
+
     // Generate a key
-    let extsk = ExtendedSpendingKey::master(&[]);
+    let efvk = mock_sapling_efvk(&[]);
-    // TODO: find out how to do it with `to_diversifiable_full_viewing_key` as `to_extended_full_viewing_key` is deprecated.
+    let dfvk = efvk.to_diversifiable_full_viewing_key();
-    let extfvk = extsk.to_extended_full_viewing_key();
+    let key_to_be_stored = encode_extended_full_viewing_key("zxviews", &efvk);
-    let dfvk: DiversifiableFullViewingKey = extsk.to_diversifiable_full_viewing_key();
-    let key_to_be_stored =
-        zcash_client_backend::encoding::encode_extended_full_viewing_key("zxviews", &extfvk);
 
     // Create a database
-    let mut s = new_test_storage(Network::Mainnet);
+    let mut storage = new_test_storage(network);
 
     // Insert the generated key to the database
-    s.add_sapling_key(&key_to_be_stored, None);
+    storage.add_sapling_key(&key_to_be_stored, None);
 
     // Check key was added
-    assert_eq!(s.sapling_keys_last_heights().len(), 1);
+    assert_eq!(storage.sapling_keys_last_heights().len(), 1);
     assert_eq!(
-        s.sapling_keys_last_heights()
+        storage
+            .sapling_keys_last_heights()
             .get(&key_to_be_stored)
             .expect("height is stored")
             .next()
             .expect("height is not maximum"),
-        s.min_sapling_birthday_height()
+        network.sapling_activation_height()
     );
 
     let nf = Nullifier([7; 32]);
@@ -186,7 +186,7 @@ fn scanning_fake_blocks_store_key_and_results() -> Result<()> {
         SaplingScannedResult::from_bytes_in_display_order(*result.transactions()[0].txid.as_ref());
 
     // Add result to database
-    s.add_sapling_results(
+    storage.add_sapling_results(
         &key_to_be_stored,
         Height(1),
         [(TransactionIndex::from_usize(0), result)].into(),
@@ -194,7 +194,7 @@ fn scanning_fake_blocks_store_key_and_results() -> Result<()> {
 
     // Check the result was added
     assert_eq!(
-        s.sapling_results(&key_to_be_stored).get(&Height(1)),
+        storage.sapling_results(&key_to_be_stored).get(&Height(1)),
         Some(&vec![result])
     );
 

zebrad/tests/common/shielded_scan/scans_for_new_key.rs

@@ -105,7 +105,12 @@ pub(crate) async fn run() -> Result<()> {
 
     tracing::info!("started scan task, sending register keys message with zecpages key to start scanning for a new key",);
 
-    scan_task.register_keys(parsed_keys)?;
+    scan_task.register_keys(
+        parsed_keys
+            .into_iter()
+            .map(|(key, (_, _, Height(h)))| (key, Some(h)))
+            .collect(),
+    )?;
 
     tracing::info!(
         ?WAIT_FOR_RESULTS_DURATION,