Make shielded syncing a separate command

crates/apps/src/lib/bench_utils.rs

@@ -685,12 +685,9 @@ impl ShieldedUtils for BenchShieldedUtils {
         // Atomicity is required to prevent other client instances from reading
         // corrupt data.
         std::fs::rename(
-            tmp_path.clone(),
+            tmp_path,
             self.context_dir.0.path().to_path_buf().join(FILE_NAME),
         )?;
-        // Finally, remove our temporary file to allow future saving of shielded
-        // contexts.
-        std::fs::remove_file(tmp_path)?;
         Ok(())
     }
 }
@@ -957,9 +954,11 @@ impl BenchShieldedCtx {
             .wallet
             .find_spending_key(ALBERT_SPENDING_KEY, None)
             .unwrap();
-        async_runtime
-            .block_on(self.shielded.fetch(
+        self.shielded = async_runtime
+            .block_on(self.shielded.syncing(
                 &self.shell,
+                &StdIo,
+                None,
                 &[spending_key.into()],
                 &[],
             ))

crates/apps/src/lib/cli.rs

@@ -267,6 +267,7 @@ pub mod cmds {
                 .subcommand(QueryMetaData::def().display_order(5))
                 // Actions
                 .subcommand(SignTx::def().display_order(6))
+                .subcommand(ShieldedSync::def().display_order(6))
                 .subcommand(GenIbcShieldedTransafer::def().display_order(6))
                 // Utils
                 .subcommand(Utils::def().display_order(7))
@@ -346,6 +347,7 @@ pub mod cmds {
             let add_to_eth_bridge_pool =
                 Self::parse_with_ctx(matches, AddToEthBridgePool);
             let sign_tx = Self::parse_with_ctx(matches, SignTx);
+            let shielded_sync = Self::parse_with_ctx(matches, ShieldedSync);
             let gen_ibc_shielded =
                 Self::parse_with_ctx(matches, GenIbcShieldedTransafer);
             let utils = SubCmd::parse(matches).map(Self::WithoutContext);
@@ -397,6 +399,7 @@ pub mod cmds {
                 .or(query_metadata)
                 .or(query_account)
                 .or(sign_tx)
+                .or(shielded_sync)
                 .or(gen_ibc_shielded)
                 .or(utils)
         }
@@ -483,6 +486,7 @@ pub mod cmds {
         QueryValidatorState(QueryValidatorState),
         QueryRewards(QueryRewards),
         SignTx(SignTx),
+        ShieldedSync(ShieldedSync),
         GenIbcShieldedTransafer(GenIbcShieldedTransafer),
     }
 
@@ -1344,6 +1348,29 @@ pub mod cmds {
         }
     }
 
+    #[derive(Clone, Debug)]
+    pub struct ShieldedSync(pub args::ShieldedSync<args::CliTypes>);
+
+    impl SubCmd for ShieldedSync {
+        const CMD: &'static str = "shielded-sync";
+
+        fn parse(matches: &ArgMatches) -> Option<Self> {
+            matches
+                .subcommand_matches(Self::CMD)
+                .map(|matches| ShieldedSync(args::ShieldedSync::parse(matches)))
+        }
+
+        fn def() -> App {
+            App::new(Self::CMD)
+                .about(
+                    "Sync the local shielded context with MASP notes owned by \
+                     the provided viewing / spending keys up to an optional \
+                     specified block height.",
+                )
+                .add_args::<args::ShieldedSync<args::CliTypes>>()
+        }
+    }
+
     #[derive(Clone, Debug)]
     pub struct Bond(pub args::Bond<args::CliTypes>);
 
@@ -3056,6 +3083,8 @@ pub mod args {
     pub const SIGNATURES: ArgMulti<PathBuf, GlobStar> = arg_multi("signatures");
     pub const SOURCE: Arg<WalletAddress> = arg("source");
     pub const SOURCE_OPT: ArgOpt<WalletAddress> = SOURCE.opt();
+    pub const SPENDING_KEYS: ArgMulti<WalletSpendingKey, GlobStar> =
+        arg_multi("spending-keys");
     pub const STEWARD: Arg<WalletAddress> = arg("steward");
     pub const SOURCE_VALIDATOR: Arg<WalletAddress> = arg("source-validator");
     pub const STORAGE_KEY: Arg<storage::Key> = arg("storage-key");
@@ -3092,6 +3121,8 @@ pub mod args {
     pub const VALUE: Arg<String> = arg("value");
     pub const VOTER_OPT: ArgOpt<WalletAddress> = arg_opt("voter");
     pub const VIEWING_KEY: Arg<WalletViewingKey> = arg("key");
+    pub const VIEWING_KEYS: ArgMulti<WalletViewingKey, GlobStar> =
+        arg_multi("viewing-keys");
     pub const VP: ArgOpt<String> = arg_opt("vp");
     pub const WALLET_ALIAS_FORCE: ArgFlag = flag("wallet-alias-force");
     pub const WASM_CHECKSUMS_PATH: Arg<PathBuf> = arg("wasm-checksums-path");
@@ -5598,6 +5629,56 @@ pub mod args {
         }
     }
 
+    impl Args for ShieldedSync<CliTypes> {
+        fn parse(matches: &ArgMatches) -> Self {
+            let ledger_address = LEDGER_ADDRESS_DEFAULT.parse(matches);
+            let last_query_height = BLOCK_HEIGHT_OPT.parse(matches);
+            let spending_keys = SPENDING_KEYS.parse(matches);
+            let viewing_keys = VIEWING_KEYS.parse(matches);
+            Self {
+                ledger_address,
+                last_query_height,
+                spending_keys,
+                viewing_keys,
+            }
+        }
+
+        fn def(app: App) -> App {
+            app.arg(LEDGER_ADDRESS_DEFAULT.def().help(LEDGER_ADDRESS_ABOUT))
+                .arg(BLOCK_HEIGHT_OPT.def().help(
+                    "Option block height to sync up to. Default is latest.",
+                ))
+                .arg(SPENDING_KEYS.def().help(
+                    "List of new spending keys with which to check note \
+                     ownership. These will be added to the shielded context.",
+                ))
+                .arg(VIEWING_KEYS.def().help(
+                    "List of new viewing keys with which to check note \
+                     ownership. These will be added to the shielded context.",
+                ))
+        }
+    }
+
+    impl CliToSdk<ShieldedSync<SdkTypes>> for ShieldedSync<CliTypes> {
+        fn to_sdk(self, ctx: &mut Context) -> ShieldedSync<SdkTypes> {
+            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            ShieldedSync {
+                ledger_address: (),
+                last_query_height: self.last_query_height,
+                spending_keys: self
+                    .spending_keys
+                    .iter()
+                    .map(|sk| chain_ctx.get_cached(sk))
+                    .collect(),
+                viewing_keys: self
+                    .viewing_keys
+                    .iter()
+                    .map(|vk| chain_ctx.get_cached(vk))
+                    .collect(),
+            }
+        }
+    }
+
     impl CliToSdk<GenIbcShieldedTransafer<SdkTypes>>
         for GenIbcShieldedTransafer<CliTypes>
     {
@@ -5881,6 +5962,7 @@ pub mod args {
         type Keypair = WalletKeypair;
         type NativeAddress = ();
         type PublicKey = WalletPublicKey;
+        type SpendingKey = WalletSpendingKey;
         type TendermintAddress = TendermintAddress;
         type TransferSource = WalletTransferSource;
         type TransferTarget = WalletTransferTarget;

crates/apps/src/lib/cli/client.rs

@@ -298,6 +298,35 @@ impl CliApi {
                         tx::submit_validator_metadata_change(&namada, args)
                             .await?;
                     }
+                    Sub::ShieldedSync(ShieldedSync(mut args)) => {
+                        let client = client.unwrap_or_else(|| {
+                            C::from_tendermint_address(&mut args.ledger_address)
+                        });
+                        client.wait_until_node_is_synced(&io).await?;
+                        let args = args.to_sdk(&mut ctx);
+                        let mut chain_ctx = ctx.take_chain_or_exit();
+                        let sks = args
+                            .spending_keys
+                            .into_iter()
+                            .map(|sk| sk.into())
+                            .collect::<Vec<_>>();
+                        let vks = args
+                            .viewing_keys
+                            .into_iter()
+                            .map(|vk| vk.into())
+                            .collect::<Vec<_>>();
+                        _ = chain_ctx.shielded.load().await;
+                        chain_ctx
+                            .shielded
+                            .syncing(
+                                &client,
+                                &io,
+                                args.last_query_height,
+                                &sks,
+                                &vks,
+                            )
+                            .await?;
+                    }
                     // Eth bridge
                     Sub::AddToEthBridgePool(args) => {
                         let mut args = args.0;

crates/apps/src/lib/client/rpc.rs

@@ -133,12 +133,7 @@ pub async fn query_transfers(
     let _ = shielded.load().await;
     // Obtain the effects of all shielded and transparent transactions
     let transfers = shielded
-        .query_tx_deltas(
-            context.client(),
-            &query_owner,
-            &query_token,
-            &wallet.get_viewing_keys(),
-        )
+        .query_tx_deltas(context.client(), &query_owner, &query_token)
         .await
         .unwrap();
     // To facilitate lookups of human-readable token names
@@ -849,13 +844,6 @@ pub async fn query_shielded_balance(
     {
         let mut shielded = context.shielded_mut().await;
         let _ = shielded.load().await;
-        let fvks: Vec<_> = viewing_keys
-            .iter()
-            .map(|fvk| ExtendedFullViewingKey::from(*fvk).fvk.vk)
-            .collect();
-        shielded.fetch(context.client(), &[], &fvks).await.unwrap();
-        // Save the update state so that future fetches can be short-circuited
-        let _ = shielded.save().await;
     }
     // The epoch is required to identify timestamped tokens
     let epoch = query_and_print_epoch(context).await;

crates/core/src/types/masp.rs

@@ -147,6 +147,13 @@ impl From<masp_primitives::zip32::ExtendedFullViewingKey>
     }
 }
 
+impl From<ExtendedViewingKey> for masp_primitives::sapling::ViewingKey {
+    fn from(value: ExtendedViewingKey) -> Self {
+        let fvk = masp_primitives::zip32::ExtendedFullViewingKey::from(value);
+        fvk.fvk.vk
+    }
+}
+
 impl serde::Serialize for ExtendedViewingKey {
     fn serialize<S>(
         &self,

crates/core/src/types/storage.rs

@@ -1,4 +1,5 @@
 //! Storage types
+use std::cmp::Ordering;
 use std::collections::VecDeque;
 use std::convert::{TryFrom, TryInto};
 use std::fmt::Display;
@@ -1453,16 +1454,7 @@ impl<E> GetEventNonce for InnerEthEventsQueue<E> {
 /// Represents the pointers of an indexed tx, which are the block height and the
 /// index inside that block
 #[derive(
-    Default,
-    Debug,
-    Copy,
-    Clone,
-    BorshSerialize,
-    BorshDeserialize,
-    Eq,
-    PartialEq,
-    Ord,
-    PartialOrd,
+    Default, Debug, Copy, Clone, BorshSerialize, BorshDeserialize, Eq, PartialEq,
 )]
 pub struct IndexedTx {
     /// The block height of the indexed tx
@@ -1471,6 +1463,22 @@ pub struct IndexedTx {
     pub index: TxIndex,
 }
 
+impl PartialOrd for IndexedTx {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for IndexedTx {
+    fn cmp(&self, other: &Self) -> Ordering {
+        if self.height == other.height {
+            self.index.cmp(&other.index)
+        } else {
+            self.height.cmp(&other.height)
+        }
+    }
+}
+
 #[cfg(test)]
 /// Tests and strategies for storage
 pub mod tests {

crates/sdk/src/args.rs

@@ -11,7 +11,7 @@ use namada_core::types::ethereum_events::EthAddress;
 use namada_core::types::keccak::KeccakHash;
 use namada_core::types::key::{common, SchemeType};
 use namada_core::types::masp::PaymentAddress;
-use namada_core::types::storage::Epoch;
+use namada_core::types::storage::{BlockHeight, Epoch};
 use namada_core::types::time::DateTimeUtc;
 use namada_core::types::{storage, token};
 use namada_governance::cli::onchain::{
@@ -56,6 +56,8 @@ pub trait NamadaTypes: Clone + std::fmt::Debug {
     type EthereumAddress: Clone + std::fmt::Debug;
     /// Represents a viewing key
     type ViewingKey: Clone + std::fmt::Debug;
+    /// Represents a spending key
+    type SpendingKey: Clone + std::fmt::Debug;
     /// Represents the owner of a balance
     type BalanceOwner: Clone + std::fmt::Debug;
     /// Represents a public key
@@ -94,6 +96,7 @@ impl NamadaTypes for SdkTypes {
     type Keypair = namada_core::types::key::common::SecretKey;
     type NativeAddress = Address;
     type PublicKey = namada_core::types::key::common::PublicKey;
+    type SpendingKey = namada_core::types::masp::ExtendedSpendingKey;
     type TendermintAddress = ();
     type TransferSource = namada_core::types::masp::TransferSource;
     type TransferTarget = namada_core::types::masp::TransferTarget;
@@ -1823,6 +1826,21 @@ pub struct SignTx<C: NamadaTypes = SdkTypes> {
     pub owner: C::Address,
 }
 
+#[derive(Clone, Debug)]
+/// Sync notes from MASP owned by the provided spending /
+/// viewing keys. Syncing can be told to stop at a given
+/// block height.
+pub struct ShieldedSync<C: NamadaTypes = SdkTypes> {
+    /// The ledger address
+    pub ledger_address: C::TendermintAddress,
+    /// Height to sync up to. Defaults to most recent
+    pub last_query_height: Option<BlockHeight>,
+    /// Spending keys used to determine note ownership
+    pub spending_keys: Vec<C::SpendingKey>,
+    /// Viewing keys used to determine note ownership
+    pub viewing_keys: Vec<C::ViewingKey>,
+}
+
 /// Query PoS commission rate
 #[derive(Clone, Debug)]
 pub struct QueryCommissionRate<C: NamadaTypes = SdkTypes> {