From 76154924a9f22033eac61f725027059d2594f57e Mon Sep 17 00:00:00 2001
From: brentstone <brentstone01@gmail.com>
Date: Mon, 29 Apr 2024 20:49:40 -0700
Subject: [PATCH] better error handling

---
 .../lib/node/ledger/shell/finalize_block.rs   |  4 +--
 crates/trans_token/src/storage.rs             | 28 +++++++++++++------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
index 4a087c8af77..5dc22ddeffe 100644
--- a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
+++ b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
@@ -850,7 +850,7 @@ mod test_finalize_block {
             &mut shell.state,
             &native_token,
             &Address::from(&keypair.ref_to()),
-            |_| Amount::native_whole(1000),
+            |_| Ok(Amount::native_whole(1000)),
         )
         .unwrap();
 
@@ -1151,7 +1151,7 @@ mod test_finalize_block {
                     &mut shell.state,
                     &native_token,
                     &bridge_pool::BRIDGE_POOL_ADDRESS,
-                    |_| amt,
+                    |_| Ok(amt),
                 )
                 .expect("Test failed");
             }
diff --git a/crates/trans_token/src/storage.rs b/crates/trans_token/src/storage.rs
index a8cce6ebe3a..4aeca8d6e9c 100644
--- a/crates/trans_token/src/storage.rs
+++ b/crates/trans_token/src/storage.rs
@@ -1,8 +1,9 @@
 use namada_core::address::{Address, InternalAddress};
 use namada_core::hints;
-use namada_core::token::{self, Amount, DenominatedAmount};
+use namada_core::token::{self, Amount, AmountError, DenominatedAmount};
 use namada_storage as storage;
 use namada_storage::{StorageRead, StorageWrite};
+use storage::ResultExt;
 
 use crate::storage_key::*;
 
@@ -41,11 +42,11 @@ pub fn update_balance<S, F>(
 ) -> storage::Result<()>
 where
     S: StorageRead + StorageWrite,
-    F: FnOnce(token::Amount) -> token::Amount,
+    F: FnOnce(token::Amount) -> storage::Result<token::Amount>,
 {
     let key = balance_key(token, owner);
     let balance = storage.read::<token::Amount>(&key)?.unwrap_or_default();
-    let new_balance = f(balance);
+    let new_balance = f(balance)?;
     storage.write(&key, new_balance)
 }
 
@@ -70,11 +71,11 @@ pub fn update_total_supply<S, F>(
 ) -> storage::Result<()>
 where
     S: StorageRead + StorageWrite,
-    F: FnOnce(token::Amount) -> token::Amount,
+    F: FnOnce(token::Amount) -> storage::Result<token::Amount>,
 {
     let key = minted_balance_key(token);
     let total_supply = storage.read::<token::Amount>(&key)?.unwrap_or_default();
-    let new_supply = f(total_supply);
+    let new_supply = f(total_supply)?;
     storage.write(&key, new_supply)
 }
 
@@ -197,10 +198,20 @@ where
     S: StorageRead + StorageWrite,
 {
     // Update the destination balance
-    update_balance(storage, token, dest, |cur_amount| cur_amount + amount)?;
+    update_balance(storage, token, dest, |cur_amount| {
+        cur_amount
+            .checked_add(amount)
+            .ok_or(AmountError::Overflow)
+            .into_storage_result()
+    })?;
 
     // Update the total supply
-    update_total_supply(storage, token, |cur_supply| cur_supply + amount)
+    update_total_supply(storage, token, |cur_supply| {
+        cur_supply
+            .checked_add(amount)
+            .ok_or(AmountError::Overflow)
+            .into_storage_result()
+    })
 }
 
 /// Burn a specified amount of tokens from some address. If the burn amount is
@@ -231,7 +242,8 @@ where
     update_total_supply(storage, token, |cur_supply| {
         cur_supply
             .checked_sub(amount_to_burn)
-            .expect("Total token supply underflowed")
+            .ok_or(AmountError::Insufficient)
+            .into_storage_result()
     })
 }