update node operator address change tracking

DMDcoin · Nov 21, 2024 · c742bd4 · c742bd4
1 parent 9ce99c8
commit c742bd4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/contracts/StakingHbbft.sol b/contracts/StakingHbbft.sol
@@ -148,8 +148,7 @@ contract StakingHbbft is Initializable, OwnableUpgradeable, ReentrancyGuardUpgra
     mapping(address => uint256) public poolNodeOperatorShare;
 
     /// @dev The epoch number in which the operator's address can be changed.
-    /// Next epoch number used to avoid issues in epoch #0.
-    mapping(address => uint256) internal _poolNodeOperatorNextAllowedChangeEpoch;
+    mapping(address => uint256) internal _poolNodeOperatorLastChangeEpoch;
 
     // ============================================== Constants =======================================================
 
@@ -1492,16 +1491,15 @@ contract StakingHbbft is Initializable, OwnableUpgradeable, ReentrancyGuardUpgra
             revert InvalidNodeOperatorConfiguration(_operatorAddress, _operatorSharePercent);
         }
 
-        uint256 nextChangeEpoch = _poolNodeOperatorNextAllowedChangeEpoch[_stakingAddress];
-
-        if (stakingEpoch < nextChangeEpoch) {
+        uint256 lastChangeEpoch = _poolNodeOperatorLastChangeEpoch[_stakingAddress];
+        if (lastChangeEpoch != 0 && lastChangeEpoch == stakingEpoch) {
             revert OnlyOncePerEpoch(stakingEpoch);
         }
 
         poolNodeOperator[_stakingAddress] = _operatorAddress;
         poolNodeOperatorShare[_stakingAddress] = _operatorSharePercent;
 
-        _poolNodeOperatorNextAllowedChangeEpoch[_stakingAddress] = stakingEpoch + 1;
+        _poolNodeOperatorLastChangeEpoch[_stakingAddress] = stakingEpoch;
 
         emit SetNodeOperator(_stakingAddress, _operatorAddress, _operatorSharePercent);
     }

diff --git a/test/StakingHbbft.ts b/test/StakingHbbft.ts
@@ -624,7 +624,7 @@ describe('StakingHbbft', () => {
         });
 
         it('should not allow to change node operator twice within same epoch', async () => {
-            const { stakingHbbft } = await helpers.loadFixture(deployContractsFixture);
+            const { stakingHbbft, validatorSetHbbft } = await helpers.loadFixture(deployContractsFixture);
 
             await stakingHbbft.connect(candidateStakingAddress).addPool(
                 candidateMiningAddress.address,
@@ -638,9 +638,19 @@ describe('StakingHbbft', () => {
             const operator = ethers.Wallet.createRandom().address;
             const share = 1000;
 
+            await stakingHbbft.setValidatorMockSetAddress(accounts[7].address);
+            await stakingHbbft.connect(accounts[7]).incrementStakingEpoch();
+            await stakingHbbft.setValidatorMockSetAddress(await validatorSetHbbft.getAddress());
+
+
+            await stakingHbbft.connect(candidateStakingAddress).setNodeOperator(operator, share);
+            expect(await stakingHbbft.poolNodeOperator(candidateStakingAddress.address)).to.equal(operator);
+            expect(await stakingHbbft.poolNodeOperatorShare(candidateStakingAddress.address)).to.equal(share);
+
+            const newOperator = ethers.Wallet.createRandom().address;
             const stakingEpoch = await stakingHbbft.stakingEpoch();
 
-            await expect(stakingHbbft.connect(candidateStakingAddress).setNodeOperator(operator, share))
+            await expect(stakingHbbft.connect(candidateStakingAddress).setNodeOperator(newOperator, share))
                 .to.be.revertedWithCustomError(stakingHbbft, "OnlyOncePerEpoch")
                 .withArgs(stakingEpoch);
         });