Fix incorrect unchecked math operation in burn

rocket-pool · Jul 30, 2024 · 64a9909 · 64a9909
1 parent c2605b5
commit 64a9909
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/src/WRETH.sol b/src/WRETH.sol
@@ -102,10 +102,10 @@ contract WRETH is IWRETH {
         // Transfer that number of token to this contract
         require(rETH.transfer(msg.sender, amountTokens), "Transfer failed");
         // Burn wrETH
-        tokenTotalSupply -= amountTokens;
-        // Cannot overflow because the sum of all user balances can't exceed the max uint256 value
+        tokenBalanceOf[msg.sender] -= amountTokens;
+        // Cannot underflow because a user's balance will never be larger than the total supply.
         unchecked {
-            tokenBalanceOf[msg.sender] -= amountTokens;
+            tokenTotalSupply -= amountTokens;
         }
         // Emit event
         emit Transfer(msg.sender, address(0), _amountWreth);
@@ -120,10 +120,10 @@ contract WRETH is IWRETH {
         // Transfer that number of token to this contract
         require(rETH.transfer(msg.sender, _amountTokens), "Transfer failed");
         // Burn wrETH
-        tokenTotalSupply -= _amountTokens;
-        // Cannot overflow because the sum of all user balances can't exceed the max uint256 value
+        tokenBalanceOf[msg.sender] -= _amountTokens;
+        // Cannot underflow because a user's balance will never be larger than the total supply.
         unchecked {
-            tokenBalanceOf[msg.sender] -= _amountTokens;
+            tokenTotalSupply -= _amountTokens;
         }
         // Emit event
         emit Transfer(msg.sender, address(0), amountWreth);

diff --git a/test/wrETH.t.sol b/test/wrETH.t.sol
@@ -47,7 +47,7 @@ contract wrETHTest is Test {
     }
 
     //
-    // Rebasing logic
+    // Misc
     //
 
     function testFail_MintZero() public {
@@ -58,6 +58,22 @@ contract wrETHTest is Test {
         wrETH.burn(0);
     }
 
+    function testFail_BurnFromOther() public {
+        mint(alice, 10000);
+        vm.prank(bob);
+        wrETH.burnTokens(1);
+    }
+
+    function testFail_BurnExceedingBalance() public {
+        mint(alice, 10000);
+        vm.prank(alice);
+        wrETH.burnTokens(10001);
+    }
+
+    //
+    // Rebasing logic
+    //
+
     function test_MintAndBurn() public {
         // Mint some wrETH
         rETH.mint(address(this), 1 ether);