add(scan): Implement RegisterKeys service request to register new k…

…eys (#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 <[email protected]> * 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 <[email protected]>
ZcashFoundation · Feb 9, 2024 · 6b8cbf9 · 6b8cbf9
1 parent 0c2c421
commit 6b8cbf9
Show file tree

Hide file tree

Showing 13 changed files with 166 additions and 139 deletions.
diff --git a/zebra-chain/src/parameters/network.rs b/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 {

diff --git a/zebra-node-services/src/scan_service/request.rs b/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
-    RegisterKeys(Vec<()>),
+    /// Submits viewing keys with their optional birth-heights for scanning.
+    RegisterKeys(Vec<(String, Option<u32>)>),
 
     /// Deletes viewing keys and their results from the database.
     DeleteKeys(Vec<String>),

diff --git a/zebra-node-services/src/scan_service/response.rs b/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>>),
 }
diff --git a/zebra-scan/src/service.rs b/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) => {
-                // TODO:
-                //  - add these keys as entries in db
-                //  - send new keys to scan task
+            Request::RegisterKeys(keys) => {
+                let mut scan_task = self.scan_task.clone();
+
+                return async move {
+                    Ok(Response::RegisteredKeys(
+                        scan_task.register_keys(keys)?.await?,
+                    ))
+                }
+                .boxed();
             }
 
             Request::DeleteKeys(keys) => {

diff --git a/zebra-scan/src/service/scan_task/commands.rs b/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<
-            SaplingScanningKey,
-            (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>, Height),
-        >,
+        keys: Vec<(String, Option<u32>)>,
+        /// Returns the set of keys the scanner accepted.
+        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, _)| {
-                            !parsed_keys.contains_key(key) || new_keys.contains_key(key)
+                        .filter_map(|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() {
-                        new_keys.extend(keys.clone());
+                    // Send the accepted keys back.
+                    let _ = rsp_tx.send(keys.iter().map(|key| key.0.clone()).collect());
 
-                        let keys =
-                            keys.into_iter()
-                                .map(|(key, (decoded_dfvks, decoded_ivks, _h))| {
-                                    (key, (decoded_dfvks, decoded_ivks))
-                                });
+                    new_keys.extend(keys.clone());
 
-                        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<
-            SaplingScanningKey,
-            (Vec<DiversifiableFullViewingKey>, Vec<SaplingIvk>, Height),
-        >,
-    ) -> Result<(), mpsc::SendError<ScanTaskCommand>> {
-        self.send(ScanTaskCommand::RegisterKeys { keys })
+        keys: Vec<(String, Option<u32>)>,
+    ) -> Result<oneshot::Receiver<Vec<String>>, mpsc::SendError<ScanTaskCommand>> {
+        let (rsp_tx, rsp_rx) = oneshot::channel();
+
+        self.send(ScanTaskCommand::RegisterKeys { keys, rsp_tx })?;
+
+        Ok(rsp_rx)
     }
 }
diff --git a/zebra-scan/src/service/scan_task/scan.rs b/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() {

diff --git a/zebra-scan/src/service/scan_task/scan/scan_range.rs b/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,

diff --git a/zebra-scan/src/service/scan_task/tests/vectors.rs b/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;
+use crate::{service::ScanTask, tests::mock_sapling_scanning_keys};
 
 /// 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(
-        (0..num_keys)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
-            .collect(),
-    )?;
+    let sapling_keys = mock_sapling_scanning_keys(num_keys.try_into().expect("should fit in u8"));
+    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())?;
 
-    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(
-        (0..num_keys)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
-            .collect(),
-    )?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
 
     // 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(
-        (10..20)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MIN)))
-            .collect(),
-    )?;
+    let sapling_keys = mock_sapling_scanning_keys(20);
+    let sapling_keys_with_birth_heights: Vec<(String, Option<u32>)> = sapling_keys
+        .clone()
+        .into_iter()
+        .map(|key| (key, Some(0)))
+        .collect();
 
-    mock_scan_task.register_keys(
-        (10..15)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[10..20].to_vec())?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[10..15].to_vec())?;
 
-    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 =
-        mock_scan_task.remove_keys(&(0..200).map(|i| i.to_string()).collect::<Vec<_>>())?;
+    let sapling_keys = mock_sapling_scanning_keys(30);
+
+    let done_rx = 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)?;
 
     // 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(
-        (0..200)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
 
-    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(
-        (0..200)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights.clone())?;
 
-    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(
-        (0..2)
-            .map(|i| (i.to_string(), (vec![], vec![], Height::MAX)))
-            .collect(),
-    )?;
+    mock_scan_task.register_keys(sapling_keys_with_birth_heights[..2].to_vec())?;
 
-    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(),