Soteria unsafe math operations

GlassEaters · May 27, 2022 · 8840eb9 · 8840eb9
1 parent bb743d1
commit 8840eb9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/programs/hydra/src/processors/transfer_shares/transfer_shares.rs b/programs/hydra/src/processors/transfer_shares/transfer_shares.rs
@@ -60,7 +60,13 @@ pub fn transfer_shares(ctx: Context<TransferShares>, shares: u64) -> Result<()>
     {
         return Err(HydraError::TransferNotSupported.into());
     }
-    from_membership_account.shares -= shares;
-    to_membership_account.shares += shares;
+    from_membership_account.shares = from_membership_account
+        .shares
+        .checked_sub(shares)
+        .expect("Sub error");
+    to_membership_account.shares = to_membership_account
+        .shares
+        .checked_add(shares)
+        .expect("Add error");
     Ok(())
 }
diff --git a/programs/hydra/src/utils/logic/calculation.rs b/programs/hydra/src/utils/logic/calculation.rs
@@ -62,7 +62,10 @@ pub fn update_inflow_for_mint(
             .or_arith_error()?
             .checked_div(tss as u128)
             .or_arith_error()? as u64;
-        fanout_for_mint.total_inflow += unstaked_correction;
+        fanout_for_mint.total_inflow = fanout_for_mint
+            .total_inflow
+            .checked_add(unstaked_correction)
+            .expect("Add error");
     }
     fanout_for_mint.last_snapshot_amount = current_snapshot;
     Ok(())
@@ -83,7 +86,10 @@ pub fn update_inflow(fanout: &mut Fanout, current_snapshot: u64) -> Result<()> {
             .or_arith_error()?
             .checked_div(tss as u128)
             .or_arith_error()? as u64;
-        fanout.total_inflow += unstaked_correction;
+        fanout.total_inflow = fanout
+            .total_inflow
+            .checked_add(unstaked_correction)
+            .expect("Add error");
     }
     fanout.last_snapshot_amount = current_snapshot;
     Ok(())