Staking: Addressing internal audit feedback

contracts/staking/README.md

@@ -40,7 +40,7 @@ To stake, any account should call `stake()`, passing in the amount to be staked
 
 To unstake, the account that previously staked should call, `unstake(uint256 _amountToUnstake)`.
 
-Accounts that wish to distribute rewards should call, `distributeRewards(AccountAmount[] calldata _recipientsAndAmounts)`. The `AccountAmount` structure consists of recipient address and amount to distribute pairs. Distributions can only be made to accounts that have previously or are currently staking. The amount to be distributed must be passed in as msg.value and must equal to the sum of the amounts specified in the `_recipientsAndAmounts` array.
+Accounts that have DISTRIBUTE_ROLE that wish to distribute rewards should call, `distributeRewards(AccountAmount[] calldata _recipientsAndAmounts)`. The `AccountAmount` structure consists of recipient address and amount to distribute pairs. Distributions can only be made to accounts that have previously or are currently staking. The amount to be distributed must be passed in as msg.value and must equal to the sum of the amounts specified in the `_recipientsAndAmounts` array.
 
 The `stakers` array needs to be analysed to determine which accounts have staked and how much. The following functions provide access to this data structure:
 

contracts/staking/StakeHolder.sol

@@ -1,4 +1,4 @@
-// Copyright (c) Immutable Pty Ltd 2018 - 2024
+// Copyright (c) Immutable Pty Ltd 2018 - 2025
 // SPDX-License-Identifier: Apache 2
 pragma solidity >=0.8.19 <0.8.29;
 
@@ -49,6 +49,9 @@ contract StakeHolder is AccessControlEnumerableUpgradeable, UUPSUpgradeable {
     /// @notice Only UPGRADE_ROLE can upgrade the contract
     bytes32 public constant UPGRADE_ROLE = bytes32("UPGRADE_ROLE");
 
+    /// @notice Only DISTRIBUTE_ROLE can call the distribute function
+    bytes32 public constant DISTRIBUTE_ROLE = bytes32("DISTRIBUTE_ROLE");
+
     /// @notice Version 0 version number
     uint256 private constant _VERSION0 = 0;
 
@@ -78,12 +81,14 @@ contract StakeHolder is AccessControlEnumerableUpgradeable, UUPSUpgradeable {
      * @notice Initialises the upgradeable contract, setting up admin accounts.
      * @param _roleAdmin the address to grant `DEFAULT_ADMIN_ROLE` to
      * @param _upgradeAdmin the address to grant `UPGRADE_ROLE` to
+     * @param _distributeAdmin the address to grant `DISTRIBUTE_ROLE` to
      */
-    function initialize(address _roleAdmin, address _upgradeAdmin) public initializer {
+    function initialize(address _roleAdmin, address _upgradeAdmin, address _distributeAdmin) public initializer {
         __UUPSUpgradeable_init();
         __AccessControl_init();
         _grantRole(DEFAULT_ADMIN_ROLE, _roleAdmin);
         _grantRole(UPGRADE_ROLE, _upgradeAdmin);
+        _grantRole(DISTRIBUTE_ROLE, _distributeAdmin);
         version = _VERSION0;
     }
 
@@ -136,14 +141,16 @@ contract StakeHolder is AccessControlEnumerableUpgradeable, UUPSUpgradeable {
     }
 
     /**
-     * @notice Any account can distribute tokens to any set of accounts.
+     * @notice Accounts with DISTRIBUTE_ROLE can distribute tokens to any set of accounts.
      * @dev The total amount to distribute must match msg.value.
      *  This function does not need re-entrancy guard as the distribution mechanism
      *  does not call out to another contract.
      * @param _recipientsAndAmounts An array of recipients to distribute value to and
      *          amounts to be distributed to each recipient.
      */
-    function distributeRewards(AccountAmount[] calldata _recipientsAndAmounts) external payable {
+    function distributeRewards(
+        AccountAmount[] calldata _recipientsAndAmounts
+    ) external payable onlyRole(DISTRIBUTE_ROLE) {
         // Initial validity checks
         if (msg.value == 0) {
             revert MustDistributeMoreThanZero();

script/staking/DeployStakeHolder.sol

@@ -34,6 +34,7 @@ struct DeploymentArgs {
 struct StakeHolderContractArgs {
     address roleAdmin;
     address upgradeAdmin;
+    address distributeAdmin;
 }
 
 /**
@@ -58,7 +59,8 @@ contract DeployStakeHolder is Test {
 
         StakeHolderContractArgs memory stakeHolderContractArgs = StakeHolderContractArgs({
             roleAdmin: makeAddr("role"),
-            upgradeAdmin: makeAddr("upgrade")
+            upgradeAdmin: makeAddr("upgrade"),
+            distributeAdmin: makeAddr("distribute")
         });
 
         // Run deployment against forked testnet
@@ -75,13 +77,14 @@ contract DeployStakeHolder is Test {
         address factory = vm.envAddress("OWNABLE_CREATE3_FACTORY_ADDRESS");
         address roleAdmin = vm.envAddress("ROLE_ADMIN");
         address upgradeAdmin = vm.envAddress("UPGRADE_ADMIN");
+        address distributeAdmin = vm.envAddress("DISTRIBUTE_ADMIN");
         string memory salt1 = vm.envString("IMPL_SALT");
         string memory salt2 = vm.envString("PROXY_SALT");
 
         DeploymentArgs memory deploymentArgs = DeploymentArgs({signer: signer, factory: factory, salt1: salt1, salt2: salt2});
 
         StakeHolderContractArgs memory stakeHolderContractArgs =
-            StakeHolderContractArgs({roleAdmin: roleAdmin, upgradeAdmin: upgradeAdmin});
+            StakeHolderContractArgs({roleAdmin: roleAdmin, upgradeAdmin: upgradeAdmin, distributeAdmin: distributeAdmin});
 
         _deploy(deploymentArgs, stakeHolderContractArgs);
     }
@@ -107,7 +110,8 @@ contract DeployStakeHolder is Test {
 
         // Create init data for teh ERC1967 Proxy
         bytes memory initData = abi.encodeWithSelector(
-            StakeHolder.initialize.selector, stakeHolderContractArgs.roleAdmin, stakeHolderContractArgs.upgradeAdmin
+            StakeHolder.initialize.selector, stakeHolderContractArgs.roleAdmin, 
+            stakeHolderContractArgs.upgradeAdmin, stakeHolderContractArgs.distributeAdmin
         );
 
         // Deploy ERC1967Proxy via the Ownable Create3 factory.

test/staking/README.md

@@ -48,4 +48,5 @@ Operational tests (in [StakeHolderOperational.t.sol](../../contracts/staking/Sta
 | testDistributeMismatch         | Fail if the total to distribute does not equal msg.value.   | No         | Yes         |
 | testDistributeToEmptyAccount   | Stake, unstake, distribute rewards.                         | Yes        | Yes         |
 | testDistributeToUnusedAccount  | Attempt to distribute rewards to an account that has never staked. | No  | Yes         |
+| testDistributeBadAuth  | Attempt to distribute rewards using an unauthorised account.        | No  | Yes         |
 

test/staking/StakeHolderBase.t.sol

@@ -15,12 +15,14 @@ contract StakeHolderBaseTest is Test {
 
     bytes32 public defaultAdminRole;
     bytes32 public upgradeRole;
+    bytes32 public distributeRole;
 
     ERC1967Proxy public proxy;
     StakeHolder public stakeHolder;
 
     address public roleAdmin;
     address public upgradeAdmin;
+    address public distributeAdmin;
 
     address public staker1;
     address public staker2;
@@ -30,6 +32,7 @@ contract StakeHolderBaseTest is Test {
     function setUp() public {
         roleAdmin = makeAddr("RoleAdmin");
         upgradeAdmin = makeAddr("UpgradeAdmin");
+        distributeAdmin = makeAddr("DistributeAdmin");
 
         staker1 = makeAddr("Staker1");
         staker2 = makeAddr("Staker2");
@@ -39,13 +42,14 @@ contract StakeHolderBaseTest is Test {
         StakeHolder impl = new StakeHolder();
 
         bytes memory initData = abi.encodeWithSelector(
-            StakeHolder.initialize.selector, roleAdmin, upgradeAdmin
+            StakeHolder.initialize.selector, roleAdmin, upgradeAdmin, distributeAdmin
         );
 
         proxy = new ERC1967Proxy(address(impl), initData);
         stakeHolder = StakeHolder(address(proxy));
 
         defaultAdminRole = stakeHolder.DEFAULT_ADMIN_ROLE();
         upgradeRole = stakeHolder.UPGRADE_ROLE();
+        distributeRole = stakeHolder.DISTRIBUTE_ROLE();
     }
 }

test/staking/StakeHolderInit.t.sol

@@ -20,7 +20,9 @@ contract StakeHolderInitTest is StakeHolderBaseTest {
     function testAdmins() public {
         assertEq(stakeHolder.getRoleMemberCount(defaultAdminRole), 1, "Expect one role admin");
         assertEq(stakeHolder.getRoleMemberCount(upgradeRole), 1, "Expect one upgrade admin");
+        assertEq(stakeHolder.getRoleMemberCount(distributeRole), 1, "Expect one distribute admin");
         assertTrue(stakeHolder.hasRole(defaultAdminRole, roleAdmin), "Expect roleAdmin is role admin");
         assertTrue(stakeHolder.hasRole(upgradeRole, upgradeAdmin), "Expect upgradeAdmin is upgrade admin");
+        assertTrue(stakeHolder.hasRole(distributeRole, distributeAdmin), "Expect distributeAdmin is distribute admin");
     }
 }

test/staking/StakeHolderOperational.t.sol

@@ -199,7 +199,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         vm.deal(staker1, 100 ether);
         vm.deal(staker2, 100 ether);
         vm.deal(staker3, 100 ether);
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         vm.prank(staker1);
         stakeHolder.stake{value: 10 ether}();
@@ -211,7 +211,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         // Distribute rewards to staker2 only.
         AccountAmount[] memory accountsAmounts = new AccountAmount[](1);
         accountsAmounts[0] = AccountAmount(staker2, 0.5 ether);
-        vm.prank(bank);
+        vm.prank(distributeAdmin);
         stakeHolder.distributeRewards{value: 0.5 ether}(accountsAmounts);
 
         assertEq(stakeHolder.getBalance(staker1), 10 ether, "Incorrect balance1");
@@ -223,7 +223,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         vm.deal(staker1, 100 ether);
         vm.deal(staker2, 100 ether);
         vm.deal(staker3, 100 ether);
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         vm.prank(staker1);
         stakeHolder.stake{value: 10 ether}();
@@ -236,7 +236,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         AccountAmount[] memory accountsAmounts = new AccountAmount[](2);
         accountsAmounts[0] = AccountAmount(staker2, 0.5 ether);
         accountsAmounts[1] = AccountAmount(staker3, 1 ether);
-        vm.prank(bank);
+        vm.prank(distributeAdmin);
         stakeHolder.distributeRewards{value: 1.5 ether}(accountsAmounts);
 
         assertEq(stakeHolder.getBalance(staker1), 10 ether, "Incorrect balance1");
@@ -246,7 +246,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
 
     function testDistributeZeroReward() public {
         vm.deal(staker1, 100 ether);
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         vm.prank(staker1);
         stakeHolder.stake{value: 10 ether}();
@@ -255,15 +255,15 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         AccountAmount[] memory accountsAmounts = new AccountAmount[](1);
         accountsAmounts[0] = AccountAmount(staker2, 0 ether);
         vm.expectRevert(abi.encodeWithSelector(StakeHolder.MustDistributeMoreThanZero.selector));
-        vm.prank(bank);
+        vm.prank(distributeAdmin);
         stakeHolder.distributeRewards{value: 0 ether}(accountsAmounts);
     }
 
     function testDistributeMismatch() public {
         vm.deal(staker1, 100 ether);
         vm.deal(staker2, 100 ether);
         vm.deal(staker3, 100 ether);
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         vm.prank(staker1);
         stakeHolder.stake{value: 10 ether}();
@@ -277,13 +277,13 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         accountsAmounts[0] = AccountAmount(staker2, 0.5 ether);
         accountsAmounts[1] = AccountAmount(staker3, 1 ether);
         vm.expectRevert(abi.encodeWithSelector(StakeHolder.DistributionAmountsDoNotMatchTotal.selector, 1 ether, 1.5 ether));
-        vm.prank(bank);
+        vm.prank(distributeAdmin);
         stakeHolder.distributeRewards{value: 1 ether}(accountsAmounts);
     }
 
     function testDistributeToEmptyAccount() public {
         vm.deal(staker1, 100 ether);
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         vm.prank(staker1);
         stakeHolder.stake{value: 10 ether}();
@@ -293,7 +293,7 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
         // Distribute rewards to staker2 only.
         AccountAmount[] memory accountsAmounts = new AccountAmount[](1);
         accountsAmounts[0] = AccountAmount(staker1, 0.5 ether);
-        vm.prank(bank);
+        vm.prank(distributeAdmin);
         stakeHolder.distributeRewards{value: 0.5 ether}(accountsAmounts);
 
         assertEq(stakeHolder.getBalance(staker1), 0.5 ether, "Incorrect balance1");
@@ -302,13 +302,30 @@ contract StakeHolderOperationalTest is StakeHolderBaseTest {
     }
 
     function testDistributeToUnusedAccount() public {
-        vm.deal(bank, 100 ether);
+        vm.deal(distributeAdmin, 100 ether);
 
         // Distribute rewards to staker2 only.
         AccountAmount[] memory accountsAmounts = new AccountAmount[](1);
         accountsAmounts[0] = AccountAmount(staker1, 0.5 ether);
         vm.expectRevert(abi.encodeWithSelector(StakeHolder.AttemptToDistributeToNewAccount.selector, staker1, 0.5 ether));
+        vm.prank(distributeAdmin);
+        stakeHolder.distributeRewards{value: 0.5 ether}(accountsAmounts);
+    }
+
+    function testDistributeBadAuth() public {
+        vm.deal(staker1, 100 ether);
+        vm.deal(bank, 100 ether);
+
+        vm.prank(staker1);
+        stakeHolder.stake{value: 10 ether}();
+
+        // Distribute rewards to staker1 only, but not from distributeAdmin
+        AccountAmount[] memory accountsAmounts = new AccountAmount[](1);
+        accountsAmounts[0] = AccountAmount(staker1, 0.5 ether);
         vm.prank(bank);
+        // Error will be of the form: 
+        // AccessControl: account 0x7fa9385be102ac3eac297483dd6233d62b3e1496 is missing role 0x555047524144455f524f4c450000000000000000000000000000000000000000
+        vm.expectRevert();
         stakeHolder.distributeRewards{value: 0.5 ether}(accountsAmounts);
     }
 }