Fix masp height overflow in SDK

.changelog/unreleased/bug-fixes/4725-fix-masp-height-overflow.md

@@ -0,0 +1,2 @@
+- Fixed a possible overflow of the masp expiration height in the SDK.
+  ([\#4725](https://github.com/anoma/namada/pull/4725))

crates/node/src/dry_run_tx.rs

@@ -33,7 +33,7 @@ where
     tx.validate_tx().into_storage_result()?;
 
     let gas_scale = parameters::get_gas_scale(&state)?;
-    let height = state.in_mem().get_last_block_height();
+    let height = state.in_mem().get_block_height().0;
 
     // Wrapper dry run to allow estimating the entire gas cost of a transaction
     let (wrapper_hash, tx_result, tx_gas_meter) = match tx.header().tx_type {

crates/node/src/shell/testing/node.rs

@@ -371,6 +371,15 @@ impl MockNode {
             .unwrap_or_default()
     }
 
+    pub fn last_block_height(&self) -> BlockHeight {
+        self.shell
+            .lock()
+            .unwrap()
+            .state
+            .in_mem()
+            .get_last_block_height()
+    }
+
     pub fn current_epoch(&self) -> Epoch {
         self.shell.lock().unwrap().state.in_mem().last_epoch
     }
@@ -484,9 +493,8 @@ impl MockNode {
     pub fn finalize_and_commit(&self, header_time: Option<DateTimeUtc>) {
         let (proposer_address, votes) = self.prepare_request();
 
+        let height = self.last_block_height().next_height();
         let mut locked = self.shell.lock().unwrap();
-        let height =
-            locked.state.in_mem().get_last_block_height().next_height();
 
         // check if we have protocol txs to be included
         // in the finalize block request
@@ -614,9 +622,8 @@ impl MockNode {
             }),
             ..Default::default()
         };
+        let height = self.last_block_height().next_height();
         let mut locked = self.shell.lock().unwrap();
-        let height =
-            locked.state.in_mem().get_last_block_height().next_height();
         let (result, tx_results) = locked.process_proposal(req);
 
         let mut errors: Vec<_> = tx_results

crates/node/src/utils.rs

@@ -189,7 +189,7 @@ pub fn dry_run_proposal(
 
     let gas_scale = parameters::get_gas_scale(&state)
         .expect("Failed to get gas scale from parameters");
-    let height = state.in_mem().get_last_block_height();
+    let height = state.in_mem().get_block_height().0;
 
     let mut tx = Tx::from_type(TxType::Raw);
     tx.header.chain_id = chain_id.clone();

crates/shielded_token/src/masp/shielded_wallet.rs

@@ -1156,10 +1156,13 @@ pub trait ShieldedApi<U: ShieldedUtils + MaybeSend + MaybeSync>:
             Some(expiration) => {
                 // Try to match a DateTime expiration with a plausible
                 // corresponding block height
-                let last_block_height = Self::query_block(context.client())
-                    .await
-                    .map_err(|e| TransferErr::General(e.to_string()))?
-                    .unwrap_or(1);
+                let last_block_height = u32::try_from(
+                    Self::query_block(context.client())
+                        .await
+                        .map_err(|e| TransferErr::General(e.to_string()))?
+                        .unwrap_or(1),
+                )
+                .map_err(|e| TransferErr::General(e.to_string()))?;
                 let max_block_time =
                     Self::query_max_block_time_estimate(context.client())
                         .await
@@ -1175,14 +1178,20 @@ pub trait ShieldedApi<U: ShieldedUtils + MaybeSend + MaybeSync>:
                         / i64::try_from(max_block_time.0).unwrap(),
                 )
                 .map_err(|e| TransferErr::General(e.to_string()))?;
-                u32::try_from(last_block_height)
-                    .map_err(|e| TransferErr::General(e.to_string()))?
-                    + delta_blocks
+                match checked!(last_block_height + delta_blocks) {
+                    Ok(height) if height <= u32::MAX - 20 => height,
+                    _ => {
+                        return Err(TransferErr::General(
+                            "The provided expiration exceeds the maximum \
+                             allowed"
+                                .to_string(),
+                        ));
+                    }
+                }
             }
             None => {
-                // NOTE: The masp library doesn't support optional
-                // expiration so we set the max to mimic
-                // a never-expiring tx. We also need to
+                // NOTE: The masp library doesn't support optional expiration so
+                // we set the max to mimic a never-expiring tx. We also need to
                 // remove 20 which is going to be added back by the builder
                 u32::MAX - 20
             }

crates/state/src/in_memory.rs

@@ -216,12 +216,23 @@ where
         )
     }
 
-    /// Get the block height
+    /// Get the block height. The height is that of the block to which the
+    /// current transaction is being applied if we are in between the
+    /// `FinalizeBlock` and the `Commit` phases. For all the other phases we
+    /// return the block height of next block that the consensus process
+    /// will decide upon (i.e. the block height of the last committed block
+    /// + 1)
     pub fn get_block_height(&self) -> (BlockHeight, Gas) {
+        let height = match self.header {
+            Some(_) => self.block.height,
+            // When not finalizing a decided block, increase the block height to
+            // match that of the next block that will be proposed
+            None => self.block.height.next_height(),
+        };
         // Adding consts that cannot overflow
         #[allow(clippy::arithmetic_side_effects)]
         (
-            self.block.height,
+            height,
             (BLOCK_HEIGHT_LENGTH as u64 * MEMORY_ACCESS_GAS_PER_BYTE).into(),
         )
     }

crates/state/src/wl_state.rs

@@ -1198,10 +1198,10 @@ where
         Ok(tree)
     }
 
-    /// Get the timestamp of the last committed block, or the current timestamp
-    /// if no blocks have been produced yet
+    /// Get the timestamp of the last committed block, or the current local
+    /// timestamp if no blocks have been produced yet
     pub fn get_last_block_timestamp(&self) -> Result<DateTimeUtc> {
-        let last_block_height = self.in_mem.get_block_height().0;
+        let last_block_height = self.in_mem.get_last_block_height();
 
         Ok(self.db.read_block_header(last_block_height)?.map_or_else(
             #[allow(clippy::disallowed_methods)]

crates/storage/src/lib.rs

@@ -89,7 +89,11 @@ pub trait StorageRead {
     fn get_chain_id(&self) -> Result<ChainId>;
 
     /// Getting the block height. The height is that of the block to which the
-    /// current transaction is being applied.
+    /// current transaction is being applied if we are in between the
+    /// `FinalizeBlock` and the `Commit` phases. For all the other phases we
+    /// return the block height of next block that the consensus process
+    /// will decide upon (i.e. the block height of the last committed block
+    /// + 1)
     fn get_block_height(&self) -> Result<BlockHeight>;
 
     /// Getting the block header.

crates/tests/src/integration/ledger_tests.rs

@@ -2325,7 +2325,7 @@ fn scheduled_migration() -> Result<()> {
         locked.scheduled_migration = Some(scheduled_migration);
     }
 
-    while node.block_height().0 != 4 {
+    while node.last_block_height().0 != 4 {
         node.finalize_and_commit(None)
     }
     // check that the key doesn't exist before the scheduled block