zecfoundation
/
zebra
mirror of https://github.com/ZcashFoundation/zebra.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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>
This commit is contained in:
Marek 2024-02-09 16:23:19 +01:00 committed by GitHub
parent 0c2c421bec
commit 6b8cbf9904
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 166 additions and 139 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
zebra-chain/src/parameters/network.rs
View File

@ -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 {

4
zebra-node-services/src/scan_service/request.rs
View File

@ -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>),

15
zebra-node-services/src/scan_service/response.rs
View File

@ -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>>),
}

15
zebra-scan/src/service.rs
View File

@ -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) => {

69
zebra-scan/src/service/scan_task/commands.rs
View File

@ -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)
}
}

4
zebra-scan/src/service/scan_task/scan.rs
View File

@ -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() {

2
zebra-scan/src/service/scan_task/scan/scan_range.rs
View File

@ -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,

89
zebra-scan/src/service/scan_task/tests/vectors.rs
View File

@ -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 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(
(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(),

20
zebra-scan/src/storage.rs
View File

@ -2,10 +2,7 @@
use std::collections::{BTreeMap, HashMap};
use zebra_chain::{
block::Height,
parameters::{Network, NetworkUpgrade},
};
use zebra_chain::{block::Height, parameters::Network};
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))
}
}

2
zebra-scan/src/storage/db/sapling.rs
View File

@ -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

39
zebra-scan/src/tests.rs
View File

@ -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::{
ChainMetadata, CompactBlock, CompactSaplingOutput, CompactSaplingSpend, CompactTx,
use zcash_client_backend::{
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>,

32
zebra-scan/src/tests/vectors.rs
View File

@ -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(&[]);
// TODO: find out how to do it with `to_diversifiable_full_viewing_key` as `to_extended_full_viewing_key` is deprecated.
let extfvk = extsk.to_extended_full_viewing_key();
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);
let efvk = mock_sapling_efvk(&[]);
let dfvk = efvk.to_diversifiable_full_viewing_key();
let key_to_be_stored = encode_extended_full_viewing_key("zxviews", &efvk);
// 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])
);

7
zebrad/tests/common/shielded_scan/scans_for_new_key.rs
View File

@ -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,