Fix Bailsec Issue_02

stader-labs · Oct 11, 2024 · dfd39d5 · dfd39d5
1 parent 2ab3a85
commit dfd39d5
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 19 deletions.
diff --git a/contracts/ValidatorRegistry.sol b/contracts/ValidatorRegistry.sol
@@ -143,9 +143,11 @@ contract ValidatorRegistry is
 
 	/// @notice Removes a validator from the registry.
 	/// @param _validatorId - Validator id
+	/// @param _ignoreBalance - If bypass the validator balance check or not
 	// slither-disable-next-line pess-multiple-storage-read
 	function removeValidator(
-		uint256 _validatorId
+		uint256 _validatorId,
+		bool _ignoreBalance
 	)
 		external
 		override
@@ -163,12 +165,14 @@ contract ValidatorRegistry is
 			"Can't remove a preferred validator for withdrawals"
 		);
 
-		address validatorShare = stakeManager.getValidatorContract(
-			_validatorId
-		);
-		(uint256 validatorBalance, ) = IValidatorShare(validatorShare)
-			.getTotalStake(maticX);
-		require(validatorBalance == 0, "Validator has some shares left");
+		if (!_ignoreBalance) {
+			address validatorShare = stakeManager.getValidatorContract(
+				_validatorId
+			);
+			(uint256 validatorBalance, ) = IValidatorShare(validatorShare)
+				.getTotalStake(maticX);
+			require(validatorBalance == 0, "Validator has some shares left");
+		}
 
 		uint256 iterationCount = validators.length - 1;
 		for (uint256 i = 0; i < iterationCount; ) {

diff --git a/contracts/interfaces/IValidatorRegistry.sol b/contracts/interfaces/IValidatorRegistry.sol
@@ -37,7 +37,11 @@ interface IValidatorRegistry {
 
 	/// @notice Removes a validator from the registry.
 	/// @param _validatorId - Validator id
-	function removeValidator(uint256 _validatorId) external;
+	/// @param _ignoreBalance - If bypass the validator balance check or not
+	function removeValidator(
+		uint256 _validatorId,
+		bool _ignoreBalance
+	) external;
 
 	/// @notice Sets the prefered validator id for deposits.
 	/// @param _validatorId - Validator id for deposits

diff --git a/test/ValidatorRegistry.spec.ts b/test/ValidatorRegistry.spec.ts
@@ -769,7 +769,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith("Pausable: paused");
 			});
 
@@ -787,7 +787,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(executor)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith(
 					`AccessControl: account ${executor.address.toLowerCase()} is missing role ${defaultAdminRole}`
 				);
@@ -803,7 +803,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(0);
+					.removeValidator(0, false);
 				await expect(promise).to.revertedWith("Zero validator id");
 			});
 
@@ -813,7 +813,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith(
 					"Validator id doesn't exist in our registry"
 				);
@@ -833,7 +833,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith(
 					"Can't remove a preferred validator for deposits"
 				);
@@ -853,13 +853,13 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith(
 					"Can't remove a preferred validator for withdrawals"
 				);
 			});
 
-			it("Should revert with the right error if having some validator shares", async function () {
+			it("Should revert with the right error if having some validator shares and the ignore balance flag set to false", async function () {
 				const { validatorRegistry, manager } =
 					await loadFixture(deployFixture);
 
@@ -878,12 +878,35 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise).to.be.revertedWith(
 					"Validator has some shares left"
 				);
 			});
 
+			it("Shouldn't revert with an error if having some validator shares and the ignore balance flag set to true", async function () {
+				const { validatorRegistry, manager } =
+					await loadFixture(deployFixture);
+
+				const maticX = (await ethers.getContractAt(
+					"IMaticX",
+					"0xf03A7Eb46d01d9EcAA104558C732Cf82f6B6B645"
+				)) as IMaticX;
+
+				await validatorRegistry
+					.connect(manager)
+					.setMaticX(maticX.address);
+
+				await validatorRegistry
+					.connect(manager)
+					.addValidator(validatorIds[0]);
+
+				const promise = validatorRegistry
+					.connect(manager)
+					.removeValidator(validatorIds[0], true);
+				await expect(promise).not.to.be.reverted;
+			});
+
 			it("Should revert with the right error if getting a removed validator id", async function () {
 				const { validatorRegistry, manager } =
 					await loadFixture(deployFixture);
@@ -894,7 +917,7 @@ describe("ValidatorRegistry", function () {
 
 				await validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 
 				const promise = validatorRegistry.getValidatorId(0);
 				await expect(promise).to.be.revertedWith(
@@ -914,7 +937,7 @@ describe("ValidatorRegistry", function () {
 
 				const promise = validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 				await expect(promise)
 					.to.emit(validatorRegistry, "RemoveValidator")
 					.withArgs(validatorIds[0]);
@@ -935,7 +958,7 @@ describe("ValidatorRegistry", function () {
 
 				await validatorRegistry
 					.connect(manager)
-					.removeValidator(validatorIds[0]);
+					.removeValidator(validatorIds[0], false);
 
 				const currentValidatorIds =
 					await validatorRegistry.getValidators();