Add test for JwtRegistry

packages/contracts/src/utils/JwtRegistry.sol

@@ -18,26 +18,30 @@ contract JwtRegistry is IDKIMRegistry, Ownable {
 
     DKIMRegistry public dkimRegistry;
 
+    // Check if azp is registered
+    mapping(string => bool) public whitelistedClients;
+
     constructor(address _owner) Ownable(_owner) {
         dkimRegistry = new DKIMRegistry(address(this));
     }
 
-    /// @notice Checks if a DKIM public key hash is valid and not revoked for a given domain name.
-    /// @param domainName The domain name to check the DKIM public key hash for.
-    /// @param publicKeyHash The DKIM public key hash to validate.
+    /// @notice Checks if a public key hash is valid and not revoked for a given kis and iss.
+    /// @param domainName The domain name contains kis, iss and azp fields.
+    /// @param publicKeyHash The public key hash to validate.
     /// @return bool Returns true if the public key hash is valid and not revoked, false otherwise.
     function isDKIMPublicKeyHashValid(
         string memory domainName,
         bytes32 publicKeyHash
     ) public view returns (bool) {
         string[] memory parts = this.stringToArray(domainName);
         string memory kidAndIss = string(abi.encode(parts[0], "|", parts[1]));
-        return dkimRegistry.isDKIMPublicKeyHashValid(kidAndIss, publicKeyHash);
+        return dkimRegistry.isDKIMPublicKeyHashValid(kidAndIss, publicKeyHash) 
+          && whitelistedClients[parts[2]];
     }
 
-    /// @notice Sets a DKIM public key hash for a domain name after validating the provided signature.
-    /// @param domainName The domain name to set the DKIM public key hash for.
-    /// @param publicKeyHash The DKIM public key hash to set.
+    /// @notice Sets a public key hash for a `kis|iss` string  after validating the provided signature.
+    /// @param domainName The domain name contains kis, iss and azp fields.
+    /// @param publicKeyHash The public key hash to set.
     /// @dev This function requires that the public key hash is not already set or revoked.
     function setDKIMPublicKeyHash(
         string memory domainName,
@@ -57,11 +61,13 @@ contract JwtRegistry is IDKIMRegistry, Ownable {
         );
 
         dkimRegistry.setDKIMPublicKeyHash(kidAndIss, publicKeyHash);
+        // Register azp
+        whitelistedClients[parts[2]] = true;
     }
 
-    /// @notice Revokes a DKIM public key hash for a domain name after validating the provided signature.
-    /// @param domainName The domain name to revoke the DKIM public key hash for.
-    /// @param publicKeyHash The DKIM public key hash to revoke.
+    /// @notice Revokes a public key hash for `kis|iss` string after validating the provided signature.
+    /// @param domainName The domain name contains kis, iss and azp fields.
+    /// @param publicKeyHash The public key hash to revoke.
     /// @dev This function requires that the public key hash is currently set and not already revoked.
     function revokeDKIMPublicKeyHash(
         string memory domainName,
@@ -79,6 +85,9 @@ contract JwtRegistry is IDKIMRegistry, Ownable {
         );
 
         dkimRegistry.revokeDKIMPublicKeyHash(publicKeyHash);
+        // Disable azp
+        string[] memory parts = this.stringToArray(domainName);
+        whitelistedClients[parts[2]] = false;
     }
 
     function stringToArray(string memory _strings) external pure returns (string[] memory) {

packages/contracts/test/utils/JwtRegistry/JwtRegistry_isDKIMPublicKeyHashValid.t.sol

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import "forge-std/Test.sol";
+import "forge-std/console.sol";
+import {EmailAuth, EmailAuthMsg} from "../../../src/EmailAuth.sol";
+import {RecoveryController} from "../../helpers/RecoveryController.sol";
+import {StructHelper} from "../../helpers/StructHelper.sol";
+import {SimpleWallet} from "../../helpers/SimpleWallet.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+
+contract JwtRegistryTest_isDKIMPublicKeyHashValid is StructHelper {
+    constructor() {}
+
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function testFail_isDKIMPublicKeyHashValid_invalidKid() public {
+        string memory domainName = "54321|https://example.com|client-id-12345";
+        bool res = jwtRegistry.isDKIMPublicKeyHashValid(
+            domainName,
+            publicKeyHash
+        );
+        assertEq(res, true);
+    }
+
+    function testFail_isDKIMPublicKeyHashValid_invalidIss() public {
+        string memory domainName = "12345|https://example.xyz|client-id-12345";
+        bool res = jwtRegistry.isDKIMPublicKeyHashValid(
+            domainName,
+            publicKeyHash
+        );
+        assertEq(res, true);
+    }
+
+    function testFail_isDKIMPublicKeyHashValid_invalidAzp() public {
+        string memory domainName = "12345|https://example.com|client-id-54321";
+        bool res = jwtRegistry.isDKIMPublicKeyHashValid(
+            domainName,
+            publicKeyHash
+        );
+        assertEq(res, true);
+    }
+
+    function test_isDKIMPublicKeyHashValid() public {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+        bool res = jwtRegistry.isDKIMPublicKeyHashValid(
+            domainName,
+            publicKeyHash
+        );
+        assertEq(res, true);
+    }
+}

packages/contracts/test/utils/JwtRegistry/JwtRegistry_revokeDKIMPublicKeyHash.t.sol

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import "forge-std/Test.sol";
+import "forge-std/console.sol";
+import {EmailAuth, EmailAuthMsg} from "../../../src/EmailAuth.sol";
+import {RecoveryController} from "../../helpers/RecoveryController.sol";
+import {StructHelper} from "../../helpers/StructHelper.sol";
+import {SimpleWallet} from "../../helpers/SimpleWallet.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@zk-email/contracts/DKIMRegistry.sol";
+
+contract JwtRegistryTest_revokeDKIMPublicKeyHash is StructHelper {
+    constructor() {}
+
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function testRevert_revokeDKIMPublicKeyHash_invalidDomainName() public {
+        string memory domainName = "";
+        vm.expectRevert(bytes("Invalid domain name"));
+        jwtRegistry.revokeDKIMPublicKeyHash(domainName, publicKeyHash);
+    }
+
+    function testRevert_revokeDKIMPublicKeyHash_invalidPublicKeyHash() public {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+        vm.expectRevert(bytes("Invalid public key hash"));
+        jwtRegistry.revokeDKIMPublicKeyHash(domainName, bytes32(0));
+    }
+
+    function testRevert_revokeDKIMPublicKeyHash_publicKeyHashIsNotSet() public {
+        string memory domainName = "54321|https://example.com|client-id-12345";
+        vm.expectRevert(bytes("publicKeyHash is not set"));
+        jwtRegistry.revokeDKIMPublicKeyHash(domainName, publicKeyHash);
+    }
+
+    function test_revokeDKIMPublicKeyHash() public {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+        jwtRegistry.revokeDKIMPublicKeyHash(domainName, publicKeyHash);
+        assertEq(jwtRegistry.whitelistedClients("client-id-12345"), false);
+    }
+}

packages/contracts/test/utils/JwtRegistry/JwtRegistry_setDKIMPublicKeyHash.t.sol

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import "forge-std/Test.sol";
+import "forge-std/console.sol";
+import {EmailAuth, EmailAuthMsg} from "../../../src/EmailAuth.sol";
+import {RecoveryController} from "../../helpers/RecoveryController.sol";
+import {StructHelper} from "../../helpers/StructHelper.sol";
+import {SimpleWallet} from "../../helpers/SimpleWallet.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@zk-email/contracts/DKIMRegistry.sol";
+
+contract JwtRegistryTest_setDKIMPublicKeyHash is StructHelper {
+    constructor() {}
+
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function testRevert_setDKIMPublicKeyHash_publicKeyHashIsAlreadySet()
+        public
+    {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+        vm.expectRevert(bytes("publicKeyHash is already set"));
+        jwtRegistry.setDKIMPublicKeyHash(domainName, publicKeyHash);
+    }
+
+    function testRevert_setDKIMPublicKeyHash_publicKeyHashIsRevoked() public {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+        jwtRegistry.revokeDKIMPublicKeyHash(domainName, publicKeyHash);
+        vm.expectRevert(bytes("publicKeyHash is revoked"));
+        jwtRegistry.setDKIMPublicKeyHash(domainName, publicKeyHash);
+    }
+
+    function test_setDKIMPublicKeyHash() public {
+        string memory domainName = "12345|https://example.xyz|client-id-12345";
+        jwtRegistry.setDKIMPublicKeyHash(domainName, publicKeyHash);
+        assertEq(jwtRegistry.whitelistedClients("client-id-12345"), true);
+    }
+}