Use custom errors for the key broadcast contract

shutter-network · Jul 4, 2023 · ef70dcf · ef70dcf
1 parent 22b1fd4
commit ef70dcf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 12 deletions.
diff --git a/src/IKeyBroadcastContract.sol b/src/IKeyBroadcastContract.sol
@@ -1,6 +1,10 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.20;
 
+error InvalidKey();
+error AlreadyHaveKey();
+error NotAllowed();
+
 interface IKeyBroadcastContract {
     function broadcastEonKey(uint64 eon, bytes memory key) external;
 

diff --git a/src/KeyBroadcastContract.sol b/src/KeyBroadcastContract.sol
@@ -13,15 +13,18 @@ contract KeyBroadcastContract is IKeyBroadcastContract {
     }
 
     function broadcastEonKey(uint64 eon, bytes memory key) external {
-        require(key.length > 0, "key is empty");
-        require(keys[eon].length == 0, "key already broadcast for this eon");
-        require(
-            msg.sender ==
-                IKeyperSetManager(keyperSetManagerAddress).getKeyperSetAddress(
-                    eon
-                ),
-            "not allowed to broadcast eon key"
-        );
+        if (key.length == 0) {
+            revert InvalidKey();
+        }
+        if (keys[eon].length > 0) {
+            revert AlreadyHaveKey();
+        }
+        if (
+            msg.sender !=
+            IKeyperSetManager(keyperSetManagerAddress).getKeyperSetAddress(eon)
+        ) {
+            revert NotAllowed();
+        }
 
         keys[eon] = key;
         emit EonKeyBroadcast(eon, key);

diff --git a/test/KeyBroadcastContract.t.sol b/test/KeyBroadcastContract.t.sol
@@ -29,14 +29,14 @@ contract KeyBroadcastTest is Test {
     }
 
     function testBroadcastEonKeyEmpty() public {
-        vm.expectRevert("key is empty");
+        vm.expectRevert(InvalidKey.selector);
         bytes memory key = bytes("");
         vm.prank(address(keyperSet));
         keyBroadcastContract.broadcastEonKey(1, key);
     }
 
     function testBroadcastEonKeyNotAllowed() public {
-        vm.expectRevert("not allowed to broadcast eon key");
+        vm.expectRevert(NotAllowed.selector);
         bytes memory key = bytes("foo bar");
         vm.prank(address(keyperSet));
         keyBroadcastContract.broadcastEonKey(0, key);
@@ -47,7 +47,7 @@ contract KeyBroadcastTest is Test {
         vm.prank(address(keyperSet));
         keyBroadcastContract.broadcastEonKey(1, key);
 
-        vm.expectRevert("key already broadcast for this eon");
+        vm.expectRevert(AlreadyHaveKey.selector);
         vm.prank(address(keyperSet));
         keyBroadcastContract.broadcastEonKey(1, key);
     }