diff --git a/packages/contracts/src/utils/JwtRegistry.sol b/packages/contracts/src/utils/JwtRegistry.sol index d2545fb5..0ed7d8f9 100644 --- a/packages/contracts/src/utils/JwtRegistry.sol +++ b/packages/contracts/src/utils/JwtRegistry.sol @@ -18,13 +18,16 @@ 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, @@ -32,12 +35,13 @@ contract JwtRegistry is IDKIMRegistry, Ownable { ) 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) { diff --git a/packages/contracts/test/utils/JwtRegistry/JwtRegistry_isDKIMPublicKeyHashValid.t.sol b/packages/contracts/test/utils/JwtRegistry/JwtRegistry_isDKIMPublicKeyHashValid.t.sol new file mode 100644 index 00000000..f7980730 --- /dev/null +++ b/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); + } +} diff --git a/packages/contracts/test/utils/JwtRegistry/JwtRegistry_revokeDKIMPublicKeyHash.t.sol b/packages/contracts/test/utils/JwtRegistry/JwtRegistry_revokeDKIMPublicKeyHash.t.sol new file mode 100644 index 00000000..b9dfe176 --- /dev/null +++ b/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); + } +} diff --git a/packages/contracts/test/utils/JwtRegistry/JwtRegistry_setDKIMPublicKeyHash.t.sol b/packages/contracts/test/utils/JwtRegistry/JwtRegistry_setDKIMPublicKeyHash.t.sol new file mode 100644 index 00000000..07c92f07 --- /dev/null +++ b/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); + } +}