set operator claim address function

towns-protocol · May 20, 2024 · cf7536e · cf7536e
1 parent f48f45a
commit cf7536e
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 34 deletions.
diff --git a/contracts/scripts/deployments/facets/DeployNodeOperator.s.sol b/contracts/scripts/deployments/facets/DeployNodeOperator.s.sol
@@ -18,8 +18,8 @@ contract DeployNodeOperator is Deployer, FacetHelper {
     addSelector(NodeOperatorFacet.getOperatorStatus.selector);
     addSelector(NodeOperatorFacet.setCommissionRate.selector);
     addSelector(NodeOperatorFacet.getCommissionRate.selector);
-    addSelector(NodeOperatorFacet.setClaimAddress.selector);
-    addSelector(NodeOperatorFacet.getClaimAddress.selector);
+    addSelector(NodeOperatorFacet.setClaimAddressForOperator.selector);
+    addSelector(NodeOperatorFacet.getClaimAddressForOperator.selector);
   }
 
   function initializer() public pure override returns (bytes4) {

diff --git a/contracts/src/base/registry/facets/distribution/RewardsDistribution.sol b/contracts/src/base/registry/facets/distribution/RewardsDistribution.sol
@@ -96,7 +96,7 @@ contract RewardsDistribution is
     uint256 operatorClaimAmount = (commission * amountPerOperator) / 100;
 
     //set that amount to the operator
-    address operatorClaimAddress = nos.claimAddressByOperator[operator];
+    address operatorClaimAddress = nos.claimerByOperator[operator];
     ds.distributionByOperator[operatorClaimAddress] += operatorClaimAmount;
     emit RewardsDistributed(operator, operatorClaimAmount);
 

diff --git a/contracts/src/base/registry/facets/operator/INodeOperator.sol b/contracts/src/base/registry/facets/operator/INodeOperator.sol
@@ -24,6 +24,7 @@ interface INodeOperatorBase {
   error NodeOperator__InvalidStakeRequirement();
   error NodeOperator__ClaimAddressNotChanged();
   error NodeOperator__InvalidCommissionRate();
+  error NodeOperator__NotClaimer();
   // =============================================================
   //                           Events
   // =============================================================
@@ -78,9 +79,14 @@ interface INodeOperator is INodeOperatorBase {
   // =============================================================
   //                           Operator Information
   // =============================================================
-  function setClaimAddress(address claimAddress) external;
+  function setClaimAddressForOperator(
+    address claimer,
+    address operator
+  ) external;
 
-  function getClaimAddress(address operator) external view returns (address);
+  function getClaimAddressForOperator(
+    address operator
+  ) external view returns (address);
 
   // =============================================================
   //                           Commission

diff --git a/contracts/src/base/registry/facets/operator/NodeOperatorFacet.sol b/contracts/src/base/registry/facets/operator/NodeOperatorFacet.sol
@@ -36,6 +36,8 @@ contract NodeOperatorFacet is INodeOperator, OwnableBase, ERC721ABase, Facet {
 
     ds.operators.add(msg.sender);
     ds.statusByOperator[msg.sender] = NodeOperatorStatus.Standby;
+    ds.claimerByOperator[msg.sender] = msg.sender;
+    ds.operatorsByClaimer[msg.sender].add(msg.sender);
 
     emit OperatorRegistered(msg.sender);
   }
@@ -104,34 +106,42 @@ contract NodeOperatorFacet is INodeOperator, OwnableBase, ERC721ABase, Facet {
   // =============================================================
 
   /// @inheritdoc INodeOperator
-  function setClaimAddress(
-    address claimAddress
-  ) external onlyValidOperator(msg.sender) {
+  function setClaimAddressForOperator(
+    address claimer,
+    address operator
+  ) external onlyClaimer(msg.sender, operator) {
     NodeOperatorStorage.Layout storage ds = NodeOperatorStorage.layout();
 
-    address currentClaimAddress = ds.claimAddressByOperator[msg.sender];
-    if (currentClaimAddress == claimAddress) {
+    if (!ds.operators.contains(operator)) revert NodeOperator__NotRegistered();
+
+    address currentClaimer = ds.claimerByOperator[operator];
+
+    if (currentClaimer == claimer) {
       revert NodeOperator__ClaimAddressNotChanged();
     }
-    ds.claimAddressByOperator[msg.sender] = claimAddress;
 
-    emit OperatorClaimAddressChanged(msg.sender, claimAddress);
+    if (ds.operatorsByClaimer[currentClaimer].contains(operator)) {
+      ds.operatorsByClaimer[currentClaimer].remove(operator);
+    }
+
+    ds.claimerByOperator[operator] = claimer;
+    ds.operatorsByClaimer[claimer].add(operator);
+
+    emit OperatorClaimAddressChanged(operator, msg.sender);
   }
 
   /// @inheritdoc INodeOperator
-  function getClaimAddress(address operator) external view returns (address) {
-    NodeOperatorStorage.Layout storage ds = NodeOperatorStorage.layout();
-    return ds.claimAddressByOperator[operator];
+  function getClaimAddressForOperator(
+    address operator
+  ) external view returns (address) {
+    return NodeOperatorStorage.layout().claimerByOperator[operator];
   }
 
   // =============================================================
   //                           Commission
   // =============================================================
-  function setCommissionRate(
-    uint256 rate
-  ) external onlyValidOperator(msg.sender) {
+  function setCommissionRate(uint256 rate) external {
     NodeOperatorStorage.Layout storage ds = NodeOperatorStorage.layout();
-
     if (!ds.operators.contains(msg.sender))
       revert NodeOperator__NotRegistered();
     if (rate > 100) revert NodeOperator__InvalidCommissionRate();
@@ -147,11 +157,13 @@ contract NodeOperatorFacet is INodeOperator, OwnableBase, ERC721ABase, Facet {
   // =============================================================
   //                           Modifiers
   // =============================================================
-  modifier onlyValidOperator(address operator) {
+
+  // only an existing claimer for that operator can call this function
+  modifier onlyClaimer(address claimer, address operator) {
     NodeOperatorStorage.Layout storage ds = NodeOperatorStorage.layout();
 
-    if (ds.statusByOperator[operator] == NodeOperatorStatus.Exiting) {
-      revert NodeOperator__NotRegistered();
+    if (!ds.operatorsByClaimer[claimer].contains(operator)) {
+      revert NodeOperator__NotClaimer();
     }
     _;
   }

diff --git a/contracts/src/base/registry/facets/operator/NodeOperatorStorage.sol b/contracts/src/base/registry/facets/operator/NodeOperatorStorage.sol
@@ -24,7 +24,8 @@ library NodeOperatorStorage {
     EnumerableSet.AddressSet operators;
     mapping(address => NodeOperatorStatus) statusByOperator;
     mapping(address => uint256) commissionByOperator;
-    mapping(address => address) claimAddressByOperator;
+    mapping(address => address) claimerByOperator;
+    mapping(address => EnumerableSet.AddressSet) operatorsByClaimer;
   }
 
   function layout() internal pure returns (Layout storage l) {

diff --git a/contracts/test/base/registry/NodeOperator.t.sol b/contracts/test/base/registry/NodeOperator.t.sol
@@ -189,8 +189,10 @@ contract NodeOperatorFacetTest is
     address _claimAddress
   ) {
     vm.assume(_claimAddress != address(0));
+    vm.assume(_operator != address(0));
+    vm.assume(_operator != _claimAddress);
     vm.prank(_operator);
-    operator.setClaimAddress(_claimAddress);
+    operator.setClaimAddressForOperator(_claimAddress, _operator);
     _;
   }
 
@@ -369,22 +371,25 @@ contract NodeOperatorFacetTest is
 
   function test_revertWhen_setClaimAddressIsCalledByInvalidOperator(
     address randomOperator,
-    address randomClaimAddress
+    address randomClaimer
   ) public {
-    vm.expectRevert(NodeOperator__NotRegistered.selector);
-    vm.prank(randomOperator);
-    operator.setClaimAddress(randomClaimAddress);
+    vm.expectRevert(NodeOperator__NotClaimer.selector);
+    vm.prank(randomClaimer);
+    operator.setClaimAddressForOperator(randomClaimer, randomOperator);
   }
 
   function test_setClaimAddress(
     address randomOperator,
-    address randomClaimAddress
+    address randomClaimer
   )
     public
     givenOperatorIsRegistered(randomOperator)
-    givenOperatorHasSetClaimAddress(randomOperator, randomClaimAddress)
+    givenOperatorHasSetClaimAddress(randomOperator, randomClaimer)
   {
-    assertEq(operator.getClaimAddress(randomOperator), randomClaimAddress);
+    assertEq(
+      operator.getClaimAddressForOperator(randomOperator),
+      randomClaimer
+    );
   }
 
   // =============================================================

diff --git a/contracts/test/base/registry/RewardsDistribution.t.sol b/contracts/test/base/registry/RewardsDistribution.t.sol
@@ -563,7 +563,7 @@ contract RewardsDistributionTest is
   {
     for (uint256 i = 0; i < operators.length; i++) {
       uint256 reward = rewardsDistributionFacet.getClaimableAmount(
-        operator.getClaimAddress(operators[i].addr)
+        operator.getClaimAddressForOperator(operators[i].addr)
       );
 
       uint256 expectedReward = _calculateExpectedOperatorReward(
@@ -601,7 +601,7 @@ contract RewardsDistributionTest is
   ) internal {
     for (uint256 i = 0; i < operators.length; i++) {
       uint256 reward = rewardsDistributionFacet.getClaimableAmount(
-        operator.getClaimAddress(operators[i].addr)
+        operator.getClaimAddressForOperator(operators[i].addr)
       );
 
       assertEq(
@@ -1157,7 +1157,7 @@ contract RewardsDistributionTest is
     vm.expectEmit();
     emit INodeOperatorBase.OperatorClaimAddressChanged(operatorAddr, claimAddr);
     vm.prank(operatorAddr);
-    operator.setClaimAddress(claimAddr);
+    operator.setClaimAddressForOperator(claimAddr, operatorAddr);
   }
 
   function delegateToOperator(address user, address operatorAddr) internal {