improve vm.assume

blockful-io · Aug 3, 2024 · cccbe73 · cccbe73
1 parent 895ea8e
commit cccbe73
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 77 deletions.
diff --git a/src/BaseStaking.sol b/src/BaseStaking.sol
@@ -181,6 +181,7 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
         uint256 assets
     ) public view virtual returns (uint256) {
         // sum + 1 on both sides to prevent donation attack
+        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
         return assets.mulDivDown(totalSupply() + 1, _totalAssets() + 1);
     }
 
@@ -190,6 +191,7 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
         uint256 shares
     ) public view virtual returns (uint256) {
         // sum + 1 on both sides to prevent donation attack
+        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
         return shares.mulDivDown(_totalAssets() + 1, totalSupply() + 1);
     }
 
@@ -215,6 +217,10 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
         // Calculate the amount of shares to mint
         uint256 shares = convertToShares(amount);
 
+        // Shares may be 0 if a first deposit donation attack occurs. Even it will not profitable for the attacker, as he will spend more tokens than he will get back
+        // this attack can still be employed to make a DDOS attack against a specific user. Although the targeted user will still be able to withdraw the SHU, this will only hold true if someone mints to increase the total supply.
+        require(shares > 0, "Shares must be greater than 0");
+
         // Update the total locked amount
         totalLocked[user] += amount;
 
@@ -253,6 +259,7 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     /// @param assets The amount of assets
     function _previewWithdraw(uint256 assets) internal view returns (uint256) {
         // sum + 1 on both sides to prevent donation attack
+        // this is the same as OZ ERC4626 prevetion to inflation attack with decimal offset = 0
         return assets.mulDivUp(totalSupply() + 1, _totalAssets() + 1);
     }
 

diff --git a/test/DelegateStaking.t.sol b/test/DelegateStaking.t.sol
@@ -138,7 +138,8 @@ contract DelegateStakingTest is Test {
                 _user != address(this) &&
                 _user != address(delegate) &&
                 _user != ProxyUtils.getAdminAddress(address(delegate)) &&
-                _user != address(rewardsDistributor)
+                _user != address(rewardsDistributor) &&
+                _user != address(staking)
         );
 
         vm.startPrank(_user);

diff --git a/test/Staking.t.sol b/test/Staking.t.sol
@@ -657,81 +657,6 @@ contract Stake is StakingTest {
         vm.stopPrank();
     }
 
-    //   function test_DonationAttack(address bob, address alice) public {
-    //       uint256 initialStake = MIN_STAKE;
-    //       uint256 donationAmount = MIN_STAKE * 100;
-    //       uint256 bobStake = MIN_STAKE * 100;
-
-    //       // first alice mints
-    //       _mintGovToken(alice, initialStake);
-    //       _setKeyper(alice, true);
-    //       _stake(alice, initialStake);
-
-    //       assertEq(staking.totalSupply(), initialStake);
-
-    //       // simulate donation
-    //       govToken.mint(address(staking), donationAmount);
-
-    //       assertEq(staking.totalSupply(), initialStake);
-    //       assertEq(
-    //           govToken.balanceOf(address(staking)),
-    //           initialStake + donationAmount
-    //       );
-
-    //       // bob mints
-    //       _mintGovToken(bob, bobStake);
-    //       _setKeyper(bob, true);
-    //       uint256 bobStakeId = _stake(bob, bobStake);
-
-    //       // bob shares
-    //       uint256 bobShares = staking.balanceOf(bob);
-    //       console.log("bob shares", bobShares);
-
-    //       //vm.prank(alice);
-    //       // uint256 aliceUnstake = staking.unstake(alice, 1, 0);
-    //       // assertEq(aliceUnstake, initialStake);
-
-    //       // alice claim rewards withdrawing donation
-    //       vm.prank(alice);
-    //       uint256 aliceRewards = staking.claimRewards(0);
-
-    //       // attacker cost is greater than expected gains
-    //       assertGt(
-    //           donationAmount,
-    //           aliceRewards,
-    //           "Alice receive more than expend for the attack"
-    //       );
-
-    //       _jumpAhead(vm.getBlockTimestamp() + LOCK_PERIOD);
-    //       // bob unstake maximum he can unstake
-    //       uint256 maxBobCanWithdraw = staking.exposed_maxWithdraw(bob, bobStake);
-    //       vm.prank(bob);
-    //       staking.unstake(bob, bobStakeId, maxBobCanWithdraw);
-
-    //       uint256 bobBalance = govToken.balanceOf(bob);
-    //       uint256 aliceBalance = govToken.balanceOf(alice);
-
-    //       vm.prank(bob);
-    //       uint256 bobRewards = staking.claimRewards(0);
-    //       console.log("bob rewards", bobRewards);
-
-    //       // bob lost a small amount maximum lost is 1%
-    //       assertApproxEqRel(
-    //           bobBalance,
-    //           bobStake + bobRewards,
-    //           0.01e18,
-    //           "Bob lost more than 1%"
-    //       );
-
-    //       // at the end Alice still lost more than bob
-    //       assertGtDecimal(
-    //           donationAmount - aliceRewards,
-    //           bobStake - bobBalance,
-    //           1e18,
-    //           "Alice receive more than bob"
-    //       );
-    //   }
-
     function testFuzz_DonationAttackNoRewards(
         address bob,
         address alice,
@@ -1588,7 +1513,6 @@ contract OwnableFunctions is StakingTest {
                 _newRewardsDistributor != address(govToken)
         );
 
-        vm.expectEmit();
         staking.setRewardsDistributor(_newRewardsDistributor);
 
         assertEq(