feat: use ownable

.gitmodules

@@ -10,3 +10,6 @@
 [submodule "lib/aave-helpers"]
 	path = lib/aave-helpers
 	url = https://github.com/bgd-labs/aave-helpers
+[submodule "lib/openzeppelin-contracts-upgradeable"]
+	path = lib/openzeppelin-contracts-upgradeable
+	url = https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable

remappings.txt

@@ -4,6 +4,7 @@ etherscan/stkAAVE:openzeppelin-contracts/=etherscan/stkAAVE/StakedAaveV3/lib/ope
 ds-test/=lib/forge-std/lib/ds-test/src/
 forge-std/=lib/aave-helpers/lib/forge-std/src/
 openzeppelin-contracts/=lib/openzeppelin-contracts/
+openzeppelin-contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/
 aave-token-v3/=lib/aave-token-v3/src/
 aave-helpers/=lib/aave-helpers/src/
 aave-address-book/=lib/aave-helpers/lib/aave-address-book/src

src/contracts/ERC20.sol

@@ -8,11 +8,11 @@ pragma solidity ^0.8.20;
 
 import {IERC20} from 'openzeppelin-contracts/contracts/token/ERC20/IERC20.sol';
 import {IERC20Metadata} from 'openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol';
-import {Context} from 'openzeppelin-contracts/contracts/utils/Context.sol';
 import {IERC20Errors} from 'openzeppelin-contracts/contracts/interfaces/draft-IERC6093.sol';
-import {Initializable} from 'openzeppelin-contracts/contracts/proxy/utils/Initializable.sol';
+import {Initializable} from 'openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol';
 import {SafeCast} from 'openzeppelin-contracts/contracts/utils/math/SafeCast.sol';
 import {DelegationMode} from 'aave-token-v3/DelegationAwareBalance.sol';
+import {OwnableUpgradeable} from 'openzeppelin-contracts-upgradeable/contracts/access/OwnableUpgradeable.sol';
 
 /**
  * @dev Implementation of the {IERC20} interface.
@@ -37,7 +37,7 @@ import {DelegationMode} from 'aave-token-v3/DelegationAwareBalance.sol';
  * by listening to said events. Other implementations of the EIP may not emit
  * these events, as it isn't required by the specification.
  */
-abstract contract ERC20 is Context, Initializable, IERC20, IERC20Metadata, IERC20Errors {
+abstract contract ERC20 is OwnableUpgradeable, IERC20, IERC20Metadata, IERC20Errors {
   struct DelegationAwareBalance {
     uint104 balance; // maximum is 10T of 18 decimal asset
     uint72 delegatedPropositionBalance;
@@ -68,7 +68,7 @@ abstract contract ERC20 is Context, Initializable, IERC20, IERC20Metadata, IERC2
   function _initializeMetadata(
     string calldata name_,
     string calldata symbol_
-  ) internal initializer {
+  ) internal onlyInitializing {
     _name = name_;
     _symbol = symbol_;
   }

src/contracts/IStakeToken.sol

@@ -22,6 +22,7 @@ interface IStakeToken is IAaveDistributionManager {
   event ExchangeRateChanged(uint216 exchangeRate);
   event FundsReturned(uint256 amount);
   event SlashingSettled();
+  event SlashingAdminChanged(address newAdmin);
 
   /**
    * @dev Allows staking a specified amount of STAKED_TOKEN

src/contracts/StakeToken.sol

@@ -10,23 +10,19 @@ import {IERC20Permit} from 'openzeppelin-contracts/contracts/token/ERC20/extensi
 
 import {ERC20Permit} from './ERC20Permit.sol';
 import {AaveDistributionManager} from './AaveDistributionManager.sol';
-import {RoleManager} from './RoleManager.sol';
 import {IStakeToken} from './IStakeToken.sol';
 import {IAaveDistributionManager} from './IAaveDistributionManager.sol';
 import {IRewardsController} from './IRewardsController.sol';
 
 import {PercentageMath} from './lib/PercentageMath.sol';
 import {DistributionTypes} from './lib/DistributionTypes.sol';
 
-contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStakeToken {
+contract StakeToken is ERC20Permit, AaveDistributionManager, IStakeToken {
   using SafeERC20 for IERC20;
   using PercentageMath for uint256;
   using SafeCast for uint256;
   using SafeCast for uint104;
 
-  uint256 public constant SLASH_ADMIN_ROLE = 0;
-  uint256 public constant COOLDOWN_ADMIN_ROLE = 1;
-  uint256 public constant CLAIM_HELPER_ROLE = 2;
   uint216 public constant INITIAL_EXCHANGE_RATE = 1e18;
   uint256 public constant EXCHANGE_RATE_UNIT = 1e18;
 
@@ -56,18 +52,10 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
   /// @notice Flag determining if there's an ongoing slashing event that needs to be settled
   bool private DEPRECATED_inPostSlashingPeriod;
 
-  modifier onlySlashingAdmin() {
-    require(msg.sender == getAdmin(SLASH_ADMIN_ROLE), 'CALLER_NOT_SLASHING_ADMIN');
-    _;
-  }
+  address internal slashingAdmin;
 
-  modifier onlyCooldownAdmin() {
-    require(msg.sender == getAdmin(COOLDOWN_ADMIN_ROLE), 'CALLER_NOT_COOLDOWN_ADMIN');
-    _;
-  }
-
-  modifier onlyClaimHelper() {
-    require(msg.sender == getAdmin(CLAIM_HELPER_ROLE), 'CALLER_NOT_CLAIM_HELPER');
+  modifier onlySlashingAdmin() {
+    require(msg.sender == slashingAdmin, 'CALLER_NOT_SLASHING_ADMIN');
     _;
   }
 
@@ -99,14 +87,8 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
     uint256 cooldownSeconds
   ) external virtual initializer {
     _initializeMetadata(name, symbol);
-
-    InitAdmin[] memory initAdmins = new InitAdmin[](3);
-    initAdmins[0] = InitAdmin(SLASH_ADMIN_ROLE, slashingAdmin);
-    initAdmins[1] = InitAdmin(COOLDOWN_ADMIN_ROLE, cooldownPauseAdmin);
-    initAdmins[2] = InitAdmin(CLAIM_HELPER_ROLE, claimHelper);
-
-    _initAdmins(initAdmins);
-
+    _transferOwnership(slashingAdmin);
+    _setSlashingAdmin(slashingAdmin);
     _setMaxSlashablePercentage(maxSlashablePercentage);
     _setCooldownSeconds(cooldownSeconds);
     _updateExchangeRate(INITIAL_EXCHANGE_RATE);
@@ -129,14 +111,23 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
     _configureAssets(assetsConfigInput);
   }
 
-  function setDistributionEnd(uint256 newDistributionEnd) external onlyEmissionManager {
+  function setDistributionEnd(uint256 newDistributionEnd) external onlyOwner {
     require(newDistributionEnd >= block.timestamp, 'END_MUST_BE_GE_NOW');
     AssetData storage assetConfig = assets[address(this)];
     _updateAssetStateInternal(address(this), assetConfig, totalSupply());
     distributionEnd = newDistributionEnd;
     emit DistributionEndChanged(newDistributionEnd);
   }
 
+  function setSlashingAdmin(address newSlashingAdmin) external onlyOwner {
+    _setSlashingAdmin(newSlashingAdmin);
+  }
+
+  function _setSlashingAdmin(address newSlashingAdmin) internal {
+    slashingAdmin = newSlashingAdmin;
+    emit SlashingAdminChanged(newSlashingAdmin);
+  }
+
   /// @inheritdoc IStakeToken
   function previewStake(uint256 assets) public view returns (uint256) {
     return (assets * _currentExchangeRate) / EXCHANGE_RATE_UNIT;
@@ -179,7 +170,7 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
   }
 
   /// @inheritdoc IStakeToken
-  function cooldownOnBehalfOf(address from) external onlyClaimHelper {
+  function cooldownOnBehalfOf(address from) external onlyOwner {
     _cooldown(from);
   }
 
@@ -189,7 +180,7 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
   }
 
   /// @inheritdoc IStakeToken
-  function redeemOnBehalf(address from, address to, uint256 amount) external onlyClaimHelper {
+  function redeemOnBehalf(address from, address to, uint256 amount) external onlyOwner {
     _redeem(from, to, amount.toUint104());
   }
 
@@ -203,7 +194,7 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
     address from,
     address to,
     uint256 amount
-  ) external onlyClaimHelper returns (uint256) {
+  ) external onlyOwner returns (uint256) {
     return _claimRewards(from, to, amount);
   }
 
@@ -219,7 +210,7 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
     address to,
     uint256 claimAmount,
     uint256 redeemAmount
-  ) external onlyClaimHelper {
+  ) external onlyOwner {
     _claimRewards(from, to, claimAmount);
     _redeem(from, to, redeemAmount.toUint104());
   }
@@ -271,7 +262,7 @@ contract StakeToken is ERC20Permit, AaveDistributionManager, RoleManager, IStake
   }
 
   /// @inheritdoc IStakeToken
-  function setCooldownSeconds(uint256 cooldownSeconds) external onlyCooldownAdmin {
+  function setCooldownSeconds(uint256 cooldownSeconds) external onlyOwner {
     _setCooldownSeconds(cooldownSeconds);
   }
 

tests/Slashing.t.sol

@@ -7,6 +7,7 @@ import {ERC20} from 'openzeppelin-contracts/contracts/token/ERC20/ERC20.sol';
 import {ProxyAdmin} from 'openzeppelin-contracts/contracts/proxy/transparent/ProxyAdmin.sol';
 import {TransparentUpgradeableProxy} from 'openzeppelin-contracts/contracts/proxy/transparent/TransparentUpgradeableProxy.sol';
 import {StkTestUtils} from './StkTestUtils.t.sol';
+import {OwnableUpgradeable} from 'openzeppelin-contracts-upgradeable/contracts/access/OwnableUpgradeable.sol';
 
 contract Slashing is StkTestUtils {
   function setUp() public {
@@ -167,14 +168,11 @@ contract Slashing is StkTestUtils {
 
   function test_changeSlashingAdmin() public {
     address newUser = vm.addr(1000);
-    try stakeToken.setPendingAdmin(stakeToken.SLASH_ADMIN_ROLE(), newUser) {} catch Error(
-      string memory reason
-    ) {
-      require(keccak256(bytes(reason)) == keccak256(bytes('CALLER_NOT_ROLE_ADMIN')));
-    }
+    vm.expectRevert(
+      abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, address(this))
+    );
+    stakeToken.setSlashingAdmin(newUser);
     vm.startPrank(admin);
-    stakeToken.setPendingAdmin(stakeToken.SLASH_ADMIN_ROLE(), newUser);
-    vm.startPrank(newUser);
-    stakeToken.claimRoleAdmin(stakeToken.SLASH_ADMIN_ROLE());
+    stakeToken.setSlashingAdmin(newUser);
   }
 }