chore: Refactor ECDSA modules

Consensys · Oct 21, 2023 · 3cf7389 · 3cf7389
1 parent 2876ed7
commit 3cf7389
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 37 deletions.
diff --git a/contracts/src/examples/modules/ECDSAModule.sol b/contracts/src/examples/modules/ECDSAModule.sol
@@ -16,7 +16,7 @@ contract ECDSAModule is AbstractModule, Ownable {
   address public signer;
   mapping(uint256 => bool) public usedNonces; // nonce => used
 
-  constructor(address initialOwner, address _signer) {
+  constructor(address _signer) {
     signer = _signer;
   }
 

diff --git a/contracts/test/example/ECDSAModule.t.sol b/contracts/test/example/ECDSAModule.t.sol
@@ -11,66 +11,50 @@ contract ECDSAModuleTest is Test {
   ECDSAModule private ecdsaModule;
   address private signer;
   uint256 private signerPk;
+  AttestationPayload private attestationPayload;
 
-  function setUp() public {
-    (signer, signerPk) = makeAddrAndKey("veraxUser");
-    ecdsaModule = new ECDSAModule(signer, signer);
+  address private user = makeAddr("veraxUser");
+  uint256 private nonce = 1234567;
 
-    vm.deal(signer, 1 ether);
-  }
+  function setUp() public {
+    (signer, signerPk) = makeAddrAndKey("veraxSigner");
+    ecdsaModule = new ECDSAModule(signer);
 
-  function test_ECDSAModule_validSignature() public {
-    address user = makeAddr("user");
-    uint256 nonce = 1234567;
-    AttestationPayload memory attestationPayload = AttestationPayload(
+    attestationPayload = AttestationPayload(
       bytes32(uint256(1234)),
       0,
       abi.encode(user),
       abi.encode(nonce)
     );
+  }
 
-    bytes32 hash = ECDSA.toEthSignedMessageHash(abi.encodePacked(user, nonce));
-
-    (uint8 v, bytes32 r, bytes32 s) = vm.sign(signerPk, hash);
-    bytes memory signature = abi.encodePacked(r, s, v);
+  function test_ECDSAModule_validSignature() public {
+    bytes memory signature = makeSignature(signerPk);
 
     ecdsaModule.run(attestationPayload, signature, signer, 0);
   }
 
   function test_ECDSAModule_revertInvalidSignature() public {
-    (address user, uint256 fakeKey) = makeAddrAndKey("user");
-    uint256 nonce = 1234567;
-    AttestationPayload memory attestationPayload = AttestationPayload(
-      bytes32(uint256(1234)),
-      0,
-      abi.encode(user),
-      abi.encode(nonce)
-    );
+    Account memory fakeSigner = makeAccount("user");
 
-    bytes32 hash = ECDSA.toEthSignedMessageHash(abi.encodePacked(user, nonce));
-    (uint8 v, bytes32 r, bytes32 s) = vm.sign(fakeKey, hash);
-    bytes memory signature = abi.encodePacked(r, s, v);
+    bytes memory signature = makeSignature(fakeSigner.key);
 
     vm.expectRevert("Wrong signature");
     ecdsaModule.run(attestationPayload, signature, user, 0);
   }
 
   function test_ECDSAModule_revertNonceUsed() public {
-    address user = makeAddr("user");
-    uint256 nonce = 1234567;
-    AttestationPayload memory attestationPayload = AttestationPayload(
-      bytes32(uint256(1234)),
-      0,
-      abi.encode(user),
-      abi.encode(nonce)
-    );
-
-    bytes32 hash = ECDSA.toEthSignedMessageHash(abi.encodePacked(user, nonce));
-    (uint8 v, bytes32 r, bytes32 s) = vm.sign(signerPk, hash);
-    bytes memory signature = abi.encodePacked(r, s, v);
+    bytes memory signature = makeSignature(signerPk);
 
     ecdsaModule.run(attestationPayload, signature, user, 0);
     vm.expectRevert("Nonce already used");
     ecdsaModule.run(attestationPayload, signature, user, 0);
   }
+
+  function makeSignature(uint256 _signer) private view returns (bytes memory) {
+    bytes32 hash = ECDSA.toEthSignedMessageHash(abi.encodePacked(user, nonce));
+
+    (uint8 v, bytes32 r, bytes32 s) = vm.sign(_signer, hash);
+    return abi.encodePacked(r, s, v);
+  }
 }