fix(account): rename deposit to add balance and add overflow check

jstz-dev · Aug 29, 2024 · c8fcaed · c8fcaed
1 parent 6155f07
commit c8fcaed
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/crates/jstz_core/tests/transaction_test.rs b/crates/jstz_core/tests/transaction_test.rs
@@ -62,7 +62,7 @@ mod test {
 
         tx.begin();
 
-        let _ = Account::deposit(hrt, &mut tx, &pkh2, 25);
+        let _ = Account::add_balance(hrt, &mut tx, &pkh2, 25);
 
         verify_account_balance(hrt, &mut tx, &pkh1, 0);
         verify_account_balance(hrt, &mut tx, &pkh2, 25);
@@ -73,7 +73,7 @@ mod test {
 
         verify_account_balance(hrt, &mut tx, &pkh2, 25);
 
-        let _ = Account::deposit(hrt, &mut tx, &pkh1, 57);
+        let _ = Account::add_balance(hrt, &mut tx, &pkh1, 57);
 
         verify_account_balance(hrt, &mut tx, &pkh1, 57);
 
@@ -83,7 +83,7 @@ mod test {
 
         verify_account_balance(hrt, &mut tx, &pkh2, 25);
 
-        let _ = Account::deposit(hrt, &mut tx, &pkh1, 57);
+        let _ = Account::add_balance(hrt, &mut tx, &pkh1, 57);
 
         verify_account_balance(hrt, &mut tx, &pkh1, 2 * 57);
 
@@ -93,7 +93,7 @@ mod test {
 
         verify_account_balance(hrt, &mut tx, &pkh1, 2 * 57);
 
-        let _ = Account::deposit(hrt, &mut tx, &pkh1, 57);
+        let _ = Account::add_balance(hrt, &mut tx, &pkh1, 57);
 
         verify_account_balance(hrt, &mut tx, &pkh1, 3 * 57);
 

diff --git a/crates/jstz_proto/src/context/account.rs b/crates/jstz_proto/src/context/account.rs
@@ -148,15 +148,19 @@ impl Account {
         Ok(account.amount)
     }
 
-    pub fn deposit(
+    pub fn add_balance(
         hrt: &impl HostRuntime,
         tx: &mut Transaction,
         addr: &Address,
         amount: Amount,
     ) -> Result<()> {
         let account = Self::get_mut(hrt, tx, addr)?;
+        let checked_balance = account
+            .amount
+            .checked_add(amount)
+            .ok_or(crate::error::Error::BalanceOverflow)?;
 
-        account.amount += amount;
+        account.amount = checked_balance;
         Ok(())
     }
 

diff --git a/crates/jstz_proto/src/executor/deposit.rs b/crates/jstz_proto/src/executor/deposit.rs
@@ -12,7 +12,7 @@ pub fn execute(
         amount, receiver, ..
     } = deposit;
 
-    let result = Account::deposit(hrt, tx, &receiver, amount);
+    let result = Account::add_balance(hrt, tx, &receiver, amount);
     let hash = Blake2b::from(deposit.inbox_id.to_be_bytes().as_slice());
     Receipt::new(hash, result.map(|_| crate::receipt::Content::Deposit))
 }