Fix shard events

Analog-Labs · Dec 22, 2024 · 6baabf4 · 6baabf4
1 parent 707c762
commit 6baabf4
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 30 deletions.
diff --git a/src/Gateway.sol b/src/Gateway.sol
@@ -101,8 +101,7 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
         ShardStore.getMainStorage().registerTssKeys(keys);
 
         // emit event
-        TssKey[] memory revoked = new TssKey[](0);
-        emit KeySetChanged(bytes32(0), revoked, keys);
+        emit ShardsRegistered(keys);
     }
 
     function nonceOf(address account) external view returns (uint256) {
@@ -164,7 +163,13 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
         store.registerTssKeys(message.register);
 
         // Emit event
-        emit KeySetChanged(messageHash, message.revoke, message.register);
+        if (message.revoke.length > 0) {
+            emit ShardsUnregistered(message.revoke);
+        }
+
+        if (message.register.length > 0) {
+            emit ShardsRegistered(message.register);
+        }
     }
 
     // Execute GMP message
@@ -425,14 +430,25 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
     function setShard(TssKey calldata publicKey) external {
         require(msg.sender == ERC1967.getAdmin(), "unauthorized");
         ShardStore.getMainStorage().register(publicKey);
+        TssKey[] memory keys = new TssKey[](1);
+        keys[0] = publicKey;
+        emit ShardsRegistered(keys);
     }
 
     /**
      * @dev Register Shards in batch.
      */
     function setShards(TssKey[] calldata publicKeys) external {
         require(msg.sender == ERC1967.getAdmin(), "unauthorized");
-        ShardStore.getMainStorage().replaceTssKeys(publicKeys);
+        (TssKey[] memory created, TssKey[] memory revoked) = ShardStore.getMainStorage().replaceTssKeys(publicKeys);
+
+        if (created.length > 0) {
+            emit ShardsRegistered(created);
+        }
+
+        if (revoked.length > 0) {
+            emit ShardsUnregistered(revoked);
+        }
     }
 
     /**
@@ -441,6 +457,9 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
     function revokeShard(TssKey calldata publicKey) external {
         require(msg.sender == ERC1967.getAdmin(), "unauthorized");
         ShardStore.getMainStorage().revoke(publicKey);
+        TssKey[] memory keys = new TssKey[](1);
+        keys[0] = publicKey;
+        emit ShardsUnregistered(keys);
     }
 
     /**
@@ -449,6 +468,7 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
     function revokeShard(TssKey[] calldata publicKeys) external {
         require(msg.sender == ERC1967.getAdmin(), "unauthorized");
         ShardStore.getMainStorage().revokeKeys(publicKeys);
+        emit ShardsUnregistered(publicKeys);
     }
 
     /*//////////////////////////////////////////////////////////////
@@ -495,19 +515,6 @@ contract Gateway is IGateway, IExecutor, IUpgradable, GatewayEIP712 {
         ERC1967.setAdmin(newAdmin);
     }
 
-    // OBS: remove != revoke (when revoked, you cannot register again)
-    function sudoRemoveShards(TssKey[] calldata revokedKeys) external payable {
-        require(msg.sender == ERC1967.getAdmin(), "unauthorized");
-        ShardStore.getMainStorage().revokeKeys(revokedKeys);
-        emit KeySetChanged(bytes32(0), revokedKeys, new TssKey[](0));
-    }
-
-    function sudoAddShards(TssKey[] calldata newKeys) external payable {
-        require(msg.sender == ERC1967.getAdmin(), "unauthorized");
-        ShardStore.getMainStorage().registerTssKeys(newKeys);
-        emit KeySetChanged(bytes32(0), new TssKey[](0), newKeys);
-    }
-
     // DANGER: This function is for migration purposes only, it allows the admin to set any storage slot.
     function sudoSetStorage(uint256[2][] calldata values) external payable {
         require(msg.sender == ERC1967.getAdmin(), "unauthorized");

diff --git a/src/interfaces/IExecutor.sol b/src/interfaces/IExecutor.sol
@@ -24,12 +24,16 @@ interface IExecutor {
     );
 
     /**
-     * @dev Emitted when `UpdateShardsMessage` is executed.
-     * @param id EIP-712 hash of the `UpdateShardsMessage`
-     * @param revoked shard's keys revoked
-     * @param registered shard's keys registered
+     * @dev Emitted when shards are registered.
+     * @param keys registered shard's keys
      */
-    event KeySetChanged(bytes32 indexed id, TssKey[] revoked, TssKey[] registered);
+    event ShardsRegistered(TssKey[] keys);
+
+    /**
+     * @dev Emitted when shards are unregistered.
+     * @param keys unregistered shard's keys
+     */
+    event ShardsUnregistered(TssKey[] keys);
 
     /**
      * @dev List all shards currently registered in the gateway.

diff --git a/src/storage/Shards.sol b/src/storage/Shards.sol
@@ -245,7 +245,7 @@ library ShardStore {
     /**
      * @dev Replace TSS keys in batch.
      * Requirements:
-     * - The `keys` should not be already registered.
+     * - The `keys` may or may not be registered.
      */
     function replaceTssKeys(MainStorage storage store, TssKey[] calldata keys)
         internal

diff --git a/src/utils/GasUtils.sol b/src/utils/GasUtils.sol
@@ -15,7 +15,7 @@ library GasUtils {
     /**
      * @dev How much gas is used until the first `gasleft()` instruction is executed.
      */
-    uint256 internal constant EXECUTION_SELECTOR_OVERHEAD = 429 + 67;
+    uint256 internal constant EXECUTION_SELECTOR_OVERHEAD = 429 + 67 - 22;
 
     /**
      * @dev Base cost of the `IExecutor.execute` method.
@@ -25,7 +25,7 @@ library GasUtils {
     /**
      * @dev Base cost of the `IGateway.submitMessage` method.
      */
-    uint256 internal constant SUBMIT_BASE_COST = 23525;
+    uint256 internal constant SUBMIT_BASE_COST = 23525 - 22;
 
     /**
      * @dev Extra gas cost of the first `IGateway.submitMessage` method.

diff --git a/test/Gateway.t.sol b/test/Gateway.t.sol
@@ -186,7 +186,7 @@ contract GatewayBase is Test {
         return Signature({xCoord: signer.xCoord(), e: e, s: s});
     }
 
-    function _shortTssKeys(TssKey[] memory keys) private pure {
+    function _sortTssKeys(TssKey[] memory keys) private pure {
         // sort keys by xCoord
         for (uint256 i = 0; i < keys.length; i++) {
             for (uint256 j = i + 1; j < keys.length; j++) {
@@ -208,7 +208,7 @@ contract GatewayBase is Test {
             signer = TestUtils.signerFromEntropy(bytes32(i));
             keys[i] = TssKey({yParity: signer.yParity() == 28 ? 3 : 2, xCoord: signer.xCoord()});
         }
-        _shortTssKeys(keys);
+        _sortTssKeys(keys);
 
         // Only admin can set shards keys
         vm.expectRevert("unauthorized");
@@ -220,7 +220,7 @@ contract GatewayBase is Test {
 
         // Check shards keys
         TssKey[] memory shards = gateway.shards();
-        _shortTssKeys(shards);
+        _sortTssKeys(shards);
         for (uint256 i = 0; i < shards.length; i++) {
             assertEq(shards[i].xCoord, keys[i].xCoord);
             assertEq(shards[i].yParity, keys[i].yParity);
@@ -230,13 +230,13 @@ contract GatewayBase is Test {
         signer = TestUtils.signerFromEntropy(bytes32(uint256(12345)));
         keys[0].xCoord = signer.xCoord();
         keys[0].yParity = signer.yParity() == 28 ? 3 : 2;
-        _shortTssKeys(keys);
+        _sortTssKeys(keys);
         vm.prank(ADMIN, ADMIN);
         gateway.setShards(keys);
 
         // Check shards keys
         shards = gateway.shards();
-        _shortTssKeys(shards);
+        _sortTssKeys(shards);
         for (uint256 i = 0; i < shards.length; i++) {
             assertEq(shards[i].xCoord, keys[i].xCoord);
             assertEq(shards[i].yParity, keys[i].yParity);