Merge pull request #4019 from anoma/grarco/improve-speculative-shield…

…ed-ctx

Improve speculative shielded ctx

.changelog/unreleased/improvements/4019-improve-speculative-shielded-ctx.md

@@ -0,0 +1,3 @@
+- The speculative shielded context now avoids updating its
+  state if the transaction failed. Added a test for it.
+  ([\#4019](https://github.com/anoma/namada/pull/4019))

crates/apps_lib/src/client/rpc.rs

@@ -4,6 +4,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::io;
 
 use borsh::BorshDeserialize;
+use color_eyre::owo_colors::OwoColorize;
 use data_encoding::HEXLOWER;
 use either::Either;
 use masp_primitives::asset_type::AssetType;
@@ -395,6 +396,15 @@ async fn query_shielded_balance(
     context: &impl Namada,
     args: args::QueryBalance,
 ) {
+    display_line!(
+        context.io(),
+        "{}: {}\n",
+        "WARNING".bold().underline().yellow(),
+        "The resulting balance could be outdated, make sure to run `namadac \
+         shielded-sync` before querying the balance to get the most recent \
+         value."
+    );
+
     let args::QueryBalance {
         // Token owner (needs to be a viewing key)
         owner,

crates/apps_lib/src/client/tx.rs

@@ -3,7 +3,9 @@ use std::io::Write;
 
 use borsh::BorshDeserialize;
 use borsh_ext::BorshSerializeExt;
+use color_eyre::owo_colors::OwoColorize;
 use ledger_namada_rs::{BIP44Path, NamadaApp};
+use namada_core::masp::MaspTransaction;
 use namada_sdk::address::{Address, ImplicitAddress};
 use namada_sdk::args::TxBecomeValidator;
 use namada_sdk::collections::HashSet;
@@ -829,13 +831,33 @@ pub async fn submit_shielded_transfer(
     namada: &impl Namada,
     args: args::TxShieldedTransfer,
 ) -> Result<(), error::Error> {
+    display_line!(
+        namada.io(),
+        "{}: {}\n",
+        "WARNING".bold().underline().yellow(),
+        "Some information might be leaked if your shielded wallet is not up \
+         to date, make sure to run `namadac shielded-sync` before running \
+         this command.",
+    );
+
     let (mut tx, signing_data) = args.clone().build(namada).await?;
 
+    let masp_section = tx
+        .sections
+        .iter()
+        .find_map(|section| section.masp_tx())
+        .ok_or_else(|| {
+            error::Error::Other(
+                "Missing MASP section in shielded transaction".to_string(),
+            )
+        })?;
     if args.tx.dump_tx || args.tx.dump_wrapper_tx {
         tx::dump_tx(namada.io(), &args.tx, tx)?;
+        pre_cache_masp_data(namada, &masp_section).await;
     } else {
         sign(namada, &mut tx, &args.tx, signing_data).await?;
-        namada.submit(tx, &args.tx).await?;
+        let res = namada.submit(tx, &args.tx).await?;
+        pre_cache_masp_data_on_tx_result(namada, &res, &masp_section).await;
     }
     Ok(())
 }
@@ -900,13 +922,33 @@ pub async fn submit_unshielding_transfer(
     namada: &impl Namada,
     args: args::TxUnshieldingTransfer,
 ) -> Result<(), error::Error> {
+    display_line!(
+        namada.io(),
+        "{}: {}\n",
+        "WARNING".bold().underline().yellow(),
+        "Some information might be leaked if your shielded wallet is not up \
+         to date, make sure to run `namadac shielded-sync` before running \
+         this command.",
+    );
+
     let (mut tx, signing_data) = args.clone().build(namada).await?;
 
+    let masp_section = tx
+        .sections
+        .iter()
+        .find_map(|section| section.masp_tx())
+        .ok_or_else(|| {
+            error::Error::Other(
+                "Missing MASP section in shielded transaction".to_string(),
+            )
+        })?;
     if args.tx.dump_tx || args.tx.dump_wrapper_tx {
         tx::dump_tx(namada.io(), &args.tx, tx)?;
+        pre_cache_masp_data(namada, &masp_section).await;
     } else {
         sign(namada, &mut tx, &args.tx, signing_data).await?;
-        namada.submit(tx, &args.tx).await?;
+        let res = namada.submit(tx, &args.tx).await?;
+        pre_cache_masp_data_on_tx_result(namada, &res, &masp_section).await;
     }
     Ok(())
 }
@@ -920,16 +962,25 @@ where
 {
     let (tx, signing_data, _) = args.build(namada).await?;
 
+    let opt_masp_section =
+        tx.sections.iter().find_map(|section| section.masp_tx());
     if args.tx.dump_tx || args.tx.dump_wrapper_tx {
         tx::dump_tx(namada.io(), &args.tx, tx)?;
+        if let Some(masp_section) = opt_masp_section {
+            pre_cache_masp_data(namada, &masp_section).await;
+        }
     } else {
-        batch_opt_reveal_pk_and_submit(
+        let res = batch_opt_reveal_pk_and_submit(
             namada,
             &args.tx,
             &[&args.source.effective_address()],
             (tx, signing_data),
         )
         .await?;
+
+        if let Some(masp_section) = opt_masp_section {
+            pre_cache_masp_data_on_tx_result(namada, &res, &masp_section).await;
+        }
     }
     // NOTE that the tx could fail when its submission epoch doesn't match
     // construction epoch
@@ -1482,3 +1533,42 @@ pub async fn gen_ibc_shielding_transfer(
     }
     Ok(())
 }
+
+// Pre-cache the data for the provided MASP transaction. Log an error on
+// failure.
+async fn pre_cache_masp_data(namada: &impl Namada, masp_tx: &MaspTransaction) {
+    if let Err(e) = namada
+        .shielded_mut()
+        .await
+        .pre_cache_transaction(masp_tx)
+        .await
+    {
+        // Just display the error but do not propagate it
+        edisplay_line!(namada.io(), "Failed to pre-cache masp data: {}.", e);
+    }
+}
+
+// Check the result of a transaction and pre-cache the masp data accordingly
+async fn pre_cache_masp_data_on_tx_result(
+    namada: &impl Namada,
+    tx_result: &ProcessTxResponse,
+    masp_tx: &MaspTransaction,
+) {
+    match tx_result {
+        ProcessTxResponse::Applied(resp) => {
+            if let Some(InnerTxResult::Success(_)) =
+                // If we have the masp data in an ibc transfer it
+                // means we are unshielding, so there's no reveal pk
+                // tx in the batch which contains only the ibc tx
+                resp.batch_result().first().map(|(_, res)| res)
+            {
+                pre_cache_masp_data(namada, masp_tx).await;
+            }
+        }
+        ProcessTxResponse::Broadcast(_) => {
+            pre_cache_masp_data(namada, masp_tx).await;
+        }
+        // Do not pre-cache when dry-running
+        ProcessTxResponse::DryRun(_) => {}
+    }
+}

crates/node/src/bench_utils.rs

@@ -1223,7 +1223,6 @@ impl BenchShieldedCtx {
                         vec![masp_transfer_data],
                         None,
                         expiration,
-                        true,
                     )
                     .await
             })

crates/sdk/src/tx.rs

@@ -2648,7 +2648,6 @@ pub async fn build_ibc_transfer(
         context,
         masp_transfer_data,
         masp_fee_data,
-        !(args.tx.dry_run || args.tx.dry_run_wrapper),
         args.tx.expiration.to_datetime(),
     )
     .await?;
@@ -3113,7 +3112,6 @@ pub async fn build_shielded_transfer<N: Namada>(
         context,
         transfer_data,
         masp_fee_data,
-        !(args.tx.dry_run || args.tx.dry_run_wrapper),
         args.tx.expiration.to_datetime(),
     )
     .await?
@@ -3280,7 +3278,6 @@ pub async fn build_shielding_transfer<N: Namada>(
         context,
         transfer_data,
         None,
-        !(args.tx.dry_run || args.tx.dry_run_wrapper),
         args.tx.expiration.to_datetime(),
     )
     .await?
@@ -3402,7 +3399,6 @@ pub async fn build_unshielding_transfer<N: Namada>(
         context,
         transfer_data,
         masp_fee_data,
-        !(args.tx.dry_run || args.tx.dry_run_wrapper),
         args.tx.expiration.to_datetime(),
     )
     .await?
@@ -3455,7 +3451,6 @@ async fn construct_shielded_parts<N: Namada>(
     context: &N,
     data: Vec<MaspTransferData>,
     fee_data: Option<MaspFeeData>,
-    update_ctx: bool,
     expiration: Option<DateTimeUtc>,
 ) -> Result<Option<(ShieldedTransfer, HashSet<AssetData>)>> {
     // Precompute asset types to increase chances of success in decoding
@@ -3469,9 +3464,7 @@ async fn construct_shielded_parts<N: Namada>(
             .await;
 
         shielded
-            .gen_shielded_transfer(
-                context, data, fee_data, expiration, update_ctx,
-            )
+            .gen_shielded_transfer(context, data, fee_data, expiration)
             .await
     };
 
@@ -3852,7 +3845,6 @@ pub async fn gen_ibc_shielding_transfer<N: Namada>(
                 // Fees are paid from the transparent balance of the relayer
                 None,
                 args.expiration.to_datetime(),
-                true,
             )
             .await
             .map_err(|err| TxSubmitError::MaspError(err.to_string()))?

crates/shielded_token/src/masp.rs

@@ -320,10 +320,10 @@ pub type WitnessMap = HashMap<usize, IncrementalWitness<Node>>;
 #[derive(BorshSerialize, BorshDeserialize, Debug)]
 /// The possible sync states of the shielded context
 pub enum ContextSyncStatus {
-    /// The context contains only data that has been confirmed by the protocol
+    /// The context contains data that has been confirmed by the protocol
     Confirmed,
-    /// The context contains that that has not yet been confirmed by the
-    /// protocol and could end up being invalid
+    /// The context possibly contains that that has not yet been confirmed by
+    /// the protocol and could be incomplete or invalid
     Speculative,
 }
 

crates/shielded_token/src/masp/shielded_wallet.rs

@@ -359,8 +359,9 @@ impl<U: ShieldedUtils + MaybeSend + MaybeSync> ShieldedWallet<U> {
     /// Updates the internal state with the data of the newly generated
     /// transaction. More specifically invalidate the spent notes, but do not
     /// cache the newly produced output descriptions and therefore the merkle
-    /// tree
-    async fn pre_cache_transaction(
+    /// tree (this is because we don't know the exact position in the tree where
+    /// the new notes will be appended)
+    pub async fn pre_cache_transaction(
         &mut self,
         masp_tx: &Transaction,
     ) -> Result<(), eyre::Error> {
@@ -1018,7 +1019,6 @@ pub trait ShieldedApi<U: ShieldedUtils + MaybeSend + MaybeSync>:
         data: Vec<MaspTransferData>,
         fee_data: Option<MaspFeeData>,
         expiration: Option<DateTimeUtc>,
-        update_ctx: bool,
     ) -> Result<Option<ShieldedTransfer>, TransferErr> {
         // Determine epoch in which to submit potential shielded transaction
         let epoch = Self::query_masp_epoch(context.client())
@@ -1212,12 +1212,6 @@ pub trait ShieldedApi<U: ShieldedUtils + MaybeSend + MaybeSync>:
             )
             .map_err(|error| TransferErr::Build { error })?;
 
-        if update_ctx {
-            self.pre_cache_transaction(&masp_tx)
-                .await
-                .map_err(|e| TransferErr::General(e.to_string()))?;
-        }
-
         Ok(Some(ShieldedTransfer {
             builder: builder_clone,
             masp_tx,