Add POL token to ValidatorRegistry contract

stader-labs · Sep 12, 2024 · 1ceff49 · 1ceff49
1 parent 2865055
commit 1ceff49
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 25 deletions.
diff --git a/contracts/ValidatorRegistry.sol b/contracts/ValidatorRegistry.sol
@@ -16,8 +16,10 @@ contract ValidatorRegistry is
 	AccessControlUpgradeable,
 	ReentrancyGuardUpgradeable
 {
+	bytes32 public constant BOT = keccak256("BOT");
+
 	address private stakeManager;
-	address private polygonERC20;
+	address private maticToken;
 	address private maticX;
 
 	string public override version;
@@ -26,27 +28,26 @@ contract ValidatorRegistry is
 	mapping(uint256 => bool) public override validatorIdExists;
 
 	uint256[] private validators;
-
-	bytes32 public constant BOT = keccak256("BOT");
+	address private polToken;
 
 	/// -------------------------- initialize ----------------------------------
 
 	/// @notice Initialize the ValidatorRegistry contract.
 	/// @param _stakeManager address of the polygon stake manager.
-	/// @param _polygonERC20 address of the polygon ERC20 contract.
+	/// @param _maticToken address of the polygon ERC20 contract.
 	/// @param _maticX address of the MaticX contract.
 	/// @param _manager address of the manager.
 	function initialize(
 		address _stakeManager,
-		address _polygonERC20,
+		address _maticToken,
 		address _maticX,
 		address _manager
 	) external initializer {
 		__AccessControl_init();
 		__Pausable_init();
 
 		stakeManager = _stakeManager;
-		polygonERC20 = _polygonERC20;
+		maticToken = _maticToken;
 		maticX = _maticX;
 
 		_setupRole(DEFAULT_ADMIN_ROLE, _manager);
@@ -179,26 +180,46 @@ contract ValidatorRegistry is
 		emit SetVersion(_version);
 	}
 
+	/**
+	 * @dev Sets the address of the POL token. Callable by the admin only.
+	 * @param _address - Address of the POL token
+	 */
+	function setPOLToken(
+		address _address
+	) external override onlyRole(DEFAULT_ADMIN_ROLE) {
+		require(_address != address(0), "Zero POL token address");
+
+		polToken = _address;
+		emit SetPOLToken(_address);
+	}
+
 	/// @notice Allows to pause the contract.
 	function togglePause() external override onlyRole(DEFAULT_ADMIN_ROLE) {
 		paused() ? _unpause() : _pause();
 	}
 
 	/// -------------------------------Getters-----------------------------------
 
-	/// @notice Get the maticX contract addresses
-	/// @return _stakeManager address of the polygon stake manager.
-	/// @return _polygonERC20 address of the polygon ERC20 contract.
-	/// @return _maticX address of the MaticX contract.
+	/// @notice Get contract addresses.
+	/// @return _stakeManager address of the stake manager.
+	/// @return _maticToken address of the Matic token.
+	/// @return _maticX address of MaticX.
+	/// @return _polToken address of the POL token.
 	function getContracts()
 		external
 		view
 		override
-		returns (address _stakeManager, address _polygonERC20, address _maticX)
+		returns (
+			address _stakeManager,
+			address _maticToken,
+			address _maticX,
+			address _polToken
+		)
 	{
 		_stakeManager = stakeManager;
-		_polygonERC20 = polygonERC20;
+		_maticToken = maticToken;
 		_maticX = maticX;
+		_polToken = polToken;
 	}
 
 	/// @notice Get validator id by its index.

diff --git a/contracts/interfaces/IValidatorRegistry.sol b/contracts/interfaces/IValidatorRegistry.sol
@@ -4,6 +4,20 @@ pragma solidity 0.8.7;
 /// @title IValidatorRegistry
 /// @notice Node validator registry interface
 interface IValidatorRegistry {
+	event AddValidator(uint256 indexed _validatorId);
+
+	event RemoveValidator(uint256 indexed _validatorId);
+
+	event SetPreferredDepositValidatorId(uint256 indexed _validatorId);
+
+	event SetPreferredWithdrawalValidatorId(uint256 indexed _validatorId);
+
+	event SetMaticX(address _address);
+
+	event SetVersion(string _version);
+
+	event SetPOLToken(address _address);
+
 	function addValidator(uint256 _validatorId) external;
 
 	function removeValidator(uint256 _validatorId) external;
@@ -24,28 +38,23 @@ interface IValidatorRegistry {
 
 	function preferredWithdrawalValidatorId() external view returns (uint256);
 
-	function validatorIdExists(uint256 _validatorId)
-		external
-		view
-		returns (bool);
+	function validatorIdExists(
+		uint256 _validatorId
+	) external view returns (bool);
 
 	function getContracts()
 		external
 		view
 		returns (
 			address _stakeManager,
-			address _polygonERC20,
-			address _maticX
+			address _maticToken,
+			address _maticX,
+			address _polToken
 		);
 
 	function getValidatorId(uint256 _index) external view returns (uint256);
 
 	function getValidators() external view returns (uint256[] memory);
 
-	event AddValidator(uint256 indexed _validatorId);
-	event RemoveValidator(uint256 indexed _validatorId);
-	event SetPreferredDepositValidatorId(uint256 indexed _validatorId);
-	event SetPreferredWithdrawalValidatorId(uint256 indexed _validatorId);
-	event SetMaticX(address _address);
-	event SetVersion(string _version);
+	function setPOLToken(address _address) external;
 }
diff --git a/test/ValidatorRegistry.spec.ts b/test/ValidatorRegistry.spec.ts
@@ -128,7 +128,6 @@ describe("ValidatorRegistry", function () {
 
 		await validatorRegistry.setMaticX(maticX.address);
 
-		// add bot role for deployer
 		await validatorRegistry.grantRole(
 			await validatorRegistry.BOT(),
 			manager.address