Disallow the zero address as the delegatee or beneficiary

uniswapfoundation · Jan 24, 2024 · 186e279 · 186e279
1 parent d6a72cd
commit 186e279
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 2 deletions.
diff --git a/src/UniStaker.sol b/src/UniStaker.sol
@@ -13,6 +13,7 @@ contract UniStaker is ReentrancyGuard {
   error UniStaker__Unauthorized(bytes32 reason, address caller);
   error UniStaker__InvalidRewardRate();
   error UniStaker__InsufficientRewardBalance();
+  error UniStaker__InvalidAddress();
 
   struct Deposit {
     uint256 balance;
@@ -102,6 +103,7 @@ contract UniStaker is ReentrancyGuard {
   }
 
   function alterDelegatee(DepositIdentifier _depositId, address _newDelegatee) public nonReentrant {
+    _revertIfAddressZero(_newDelegatee);
     Deposit storage deposit = deposits[_depositId];
     _revertIfNotDepositOwner(deposit);
 
@@ -115,6 +117,7 @@ contract UniStaker is ReentrancyGuard {
     public
     nonReentrant
   {
+    _revertIfAddressZero(_newBeneficiary);
     Deposit storage deposit = deposits[_depositId];
     _revertIfNotDepositOwner(deposit);
 
@@ -188,6 +191,9 @@ contract UniStaker is ReentrancyGuard {
     internal
     returns (DepositIdentifier _depositId)
   {
+    _revertIfAddressZero(_delegatee);
+    _revertIfAddressZero(_beneficiary);
+
     _updateReward(_beneficiary);
 
     DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee);
@@ -220,4 +226,8 @@ contract UniStaker is ReentrancyGuard {
   function _revertIfNotDepositOwner(Deposit storage deposit) internal view {
     if (msg.sender != deposit.owner) revert UniStaker__Unauthorized("not owner", msg.sender);
   }
+
+  function _revertIfAddressZero(address _account) internal pure {
+    if (_account == address(0)) revert UniStaker__InvalidAddress();
+  }
 }
diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol
@@ -55,6 +55,8 @@ contract UniStakerTest is Test {
     internal
     returns (UniStaker.DepositIdentifier _depositId)
   {
+    vm.assume(_delegatee != address(0));
+
     vm.startPrank(_depositor);
     govToken.approve(address(uniStaker), _amount);
     _depositId = uniStaker.stake(_amount, _delegatee);
@@ -65,6 +67,8 @@ contract UniStakerTest is Test {
     internal
     returns (UniStaker.DepositIdentifier _depositId)
   {
+    vm.assume(_delegatee != address(0) && _beneficiary != address(0));
+
     vm.startPrank(_depositor);
     govToken.approve(address(uniStaker), _amount);
     _depositId = uniStaker.stake(_amount, _delegatee, _beneficiary);
@@ -474,6 +478,32 @@ contract Stake is UniStakerTest {
       _delegatee = address(uint160(uint256(keccak256(abi.encode(_delegatee)))));
     }
   }
+
+  function testFuzz_RevertIf_DelegateeIsTheZeroAddress(address _depositor, uint256 _amount) public {
+    _amount = _boundMintAmount(_amount);
+    _mintGovToken(_depositor, _amount);
+    govToken.approve(address(uniStaker), _amount);
+
+    vm.prank(_depositor);
+    vm.expectRevert(UniStaker.UniStaker__InvalidAddress.selector);
+    uniStaker.stake(_amount, address(0));
+  }
+
+  function testFuzz_RevertIf_BeneficiaryIsTheZeroAddress(
+    address _depositor,
+    uint256 _amount,
+    address _delegatee
+  ) public {
+    vm.assume(_delegatee != address(0));
+
+    _amount = _boundMintAmount(_amount);
+    _mintGovToken(_depositor, _amount);
+    govToken.approve(address(uniStaker), _amount);
+
+    vm.prank(_depositor);
+    vm.expectRevert(UniStaker.UniStaker__InvalidAddress.selector);
+    uniStaker.stake(_amount, _delegatee, address(0));
+  }
 }
 
 contract StakeMore is UniStakerTest {
@@ -721,7 +751,7 @@ contract AlterDelegatee is UniStakerTest {
     UniStaker.DepositIdentifier _depositId,
     address _newDelegatee
   ) public {
-    vm.assume(_depositor != address(0));
+    vm.assume(_depositor != address(0) && _newDelegatee != address(0));
 
     // Since no deposits have been made yet, all DepositIdentifiers are invalid, and any call to
     // alter one should revert. We rely on the default owner of any uninitialized deposit being
@@ -734,6 +764,19 @@ contract AlterDelegatee is UniStakerTest {
     );
     uniStaker.alterDelegatee(_depositId, _newDelegatee);
   }
+
+  function testFuzz_RevertIf_DelegateeIsTheZeroAddress(
+    address _depositor,
+    uint256 _depositAmount,
+    address _delegatee
+  ) public {
+    UniStaker.DepositIdentifier _depositId;
+    (_depositAmount, _depositId) = _boundMintAndStake(_depositor, _depositAmount, _delegatee);
+
+    vm.prank(_depositor);
+    vm.expectRevert(UniStaker.UniStaker__InvalidAddress.selector);
+    uniStaker.alterDelegatee(_depositId, address(0));
+  }
 }
 
 contract AlterBeneficiary is UniStakerTest {
@@ -812,7 +855,7 @@ contract AlterBeneficiary is UniStakerTest {
     UniStaker.DepositIdentifier _depositId,
     address _newBeneficiary
   ) public {
-    vm.assume(_depositor != address(0));
+    vm.assume(_depositor != address(0) && _newBeneficiary != address(0));
 
     // Since no deposits have been made yet, all DepositIdentifiers are invalid, and any call to
     // alter one should revert. We rely on the default owner of any uninitialized deposit being
@@ -825,6 +868,19 @@ contract AlterBeneficiary is UniStakerTest {
     );
     uniStaker.alterBeneficiary(_depositId, _newBeneficiary);
   }
+
+  function testFuzz_RevertIf_BeneficiaryIsTheZeroAddress(
+    address _depositor,
+    uint256 _depositAmount,
+    address _delegatee
+  ) public {
+    UniStaker.DepositIdentifier _depositId;
+    (_depositAmount, _depositId) = _boundMintAndStake(_depositor, _depositAmount, _delegatee);
+
+    vm.prank(_depositor);
+    vm.expectRevert(UniStaker.UniStaker__InvalidAddress.selector);
+    uniStaker.alterBeneficiary(_depositId, address(0));
+  }
 }
 
 contract Withdraw is UniStakerTest {