[FEA-1063] remediate Slowmist audit suggestion [N1] (#66)

plumenetwork · Oct 23, 2024 · 70e2c22 · 70e2c22
1 parent a19b9a1
commit 70e2c22
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 15 deletions.
diff --git a/staking/src/RWAStaking.sol b/staking/src/RWAStaking.sol
@@ -191,19 +191,24 @@ contract RWAStaking is AccessControlUpgradeable, UUPSUpgradeable {
             revert NotAllowedStablecoin(stablecoin);
         }
 
+        uint256 previousBalance = stablecoin.balanceOf(address(this));
+
         stablecoin.safeTransferFrom(msg.sender, address(this), amount);
 
+        uint256 newBalance = stablecoin.balanceOf(address(this));
+        uint256 actualAmount = newBalance - previousBalance;
+
         uint256 timestamp = block.timestamp;
         UserState storage userState = $.userStates[msg.sender];
         if (userState.lastUpdate == 0) {
             $.users.push(msg.sender);
         }
         userState.amountSeconds += userState.amountStaked * (timestamp - userState.lastUpdate);
-        userState.amountStaked += amount;
+        userState.amountStaked += actualAmount;
         userState.lastUpdate = timestamp;
-        $.totalAmountStaked += amount;
+        $.totalAmountStaked += actualAmount;
 
-        emit Staked(msg.sender, stablecoin, amount, timestamp);
+        emit Staked(msg.sender, stablecoin, actualAmount, timestamp);
     }
 
     // Getter View Functions

diff --git a/staking/src/ReserveStaking.sol b/staking/src/ReserveStaking.sol
@@ -196,23 +196,34 @@ contract ReserveStaking is AccessControlUpgradeable, UUPSUpgradeable {
             $.users.push(msg.sender);
         }
 
+        uint256 actualSbtcAmount;
+        uint256 actualStoneAmount;
+
         if (sbtcAmount > 0) {
-            $.sbtc.safeTransferFrom(msg.sender, address(this), sbtcAmount);
+            IERC20 sbtc = $.sbtc;
+            uint256 previousBalance = sbtc.balanceOf(address(this));
+            sbtc.safeTransferFrom(msg.sender, address(this), sbtcAmount);
+            uint256 newBalance = sbtc.balanceOf(address(this));
+            actualSbtcAmount = newBalance - previousBalance;
             userState.sbtcAmountSeconds += userState.sbtcAmountStaked * (timestamp - userState.sbtcLastUpdate);
-            userState.sbtcAmountStaked += sbtcAmount;
+            userState.sbtcAmountStaked += actualSbtcAmount;
             userState.sbtcLastUpdate = timestamp;
-            $.sbtcTotalAmountStaked += sbtcAmount;
+            $.sbtcTotalAmountStaked += actualSbtcAmount;
         }
 
         if (stoneAmount > 0) {
-            $.stone.safeTransferFrom(msg.sender, address(this), stoneAmount);
+            IERC20 stone = $.stone;
+            uint256 previousBalance = stone.balanceOf(address(this));
+            stone.safeTransferFrom(msg.sender, address(this), stoneAmount);
+            uint256 newBalance = stone.balanceOf(address(this));
+            actualStoneAmount = newBalance - previousBalance;
             userState.stoneAmountSeconds += userState.stoneAmountStaked * (timestamp - userState.stoneLastUpdate);
             userState.stoneAmountStaked += stoneAmount;
             userState.stoneLastUpdate = timestamp;
             $.stoneTotalAmountStaked += stoneAmount;
         }
 
-        emit Staked(msg.sender, sbtcAmount, stoneAmount);
+        emit Staked(msg.sender, actualSbtcAmount, actualStoneAmount);
     }
 
     /**
@@ -234,25 +245,36 @@ contract ReserveStaking is AccessControlUpgradeable, UUPSUpgradeable {
             );
         }
 
+        uint256 actualSbtcAmount;
+        uint256 actualStoneAmount;
+
         if (sbtcAmount > 0) {
+            IERC20 sbtc = $.sbtc;
             userState.sbtcAmountSeconds += userState.sbtcAmountStaked * (timestamp - userState.sbtcLastUpdate);
-            $.sbtc.safeTransfer(msg.sender, sbtcAmount);
+            uint256 previousBalance = sbtc.balanceOf(address(this));
+            sbtc.safeTransfer(msg.sender, sbtcAmount);
+            uint256 newBalance = sbtc.balanceOf(address(this));
+            actualSbtcAmount = previousBalance - newBalance;
             userState.sbtcAmountSeconds -= userState.sbtcAmountSeconds * sbtcAmount / userState.sbtcAmountStaked;
-            userState.sbtcAmountStaked -= sbtcAmount;
+            userState.sbtcAmountStaked -= actualSbtcAmount;
             userState.sbtcLastUpdate = timestamp;
-            $.sbtcTotalAmountStaked -= sbtcAmount;
+            $.sbtcTotalAmountStaked -= actualSbtcAmount;
         }
 
         if (stoneAmount > 0) {
+            IERC20 stone = $.stone;
             userState.stoneAmountSeconds += userState.stoneAmountStaked * (timestamp - userState.stoneLastUpdate);
-            $.stone.safeTransfer(msg.sender, stoneAmount);
+            uint256 previousBalance = stone.balanceOf(address(this));
+            stone.safeTransfer(msg.sender, stoneAmount);
+            uint256 newBalance = stone.balanceOf(address(this));
+            actualStoneAmount = previousBalance - newBalance;
             userState.stoneAmountSeconds -= userState.stoneAmountSeconds * stoneAmount / userState.stoneAmountStaked;
-            userState.stoneAmountStaked -= stoneAmount;
+            userState.stoneAmountStaked -= actualStoneAmount;
             userState.stoneLastUpdate = timestamp;
-            $.stoneTotalAmountStaked -= stoneAmount;
+            $.stoneTotalAmountStaked -= actualStoneAmount;
         }
 
-        emit Withdrawn(msg.sender, sbtcAmount, stoneAmount);
+        emit Withdrawn(msg.sender, actualSbtcAmount, actualStoneAmount);
     }
 
     // Getter View Functions