From f038366751711ede46afe6300d264bf9d06ae47f Mon Sep 17 00:00:00 2001
From: jpbogle <jpbogle22@gmail.com>
Date: Tue, 24 May 2022 18:14:43 -0700
Subject: [PATCH] Soteria unsafe math operations

---
 .../src/processors/transfer_shares/transfer_shares.rs  | 10 ++++++++--
 programs/hydra/src/utils/logic/calculation.rs          | 10 ++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/programs/hydra/src/processors/transfer_shares/transfer_shares.rs b/programs/hydra/src/processors/transfer_shares/transfer_shares.rs
index 476fb56..2a940d7 100644
--- 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
index 1eeffb5..61bc8f4 100644
--- 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(())