Add disableAzp

zkemail · Oct 11, 2024 · 69e5fff · 69e5fff
1 parent a8ff317
commit 69e5fff
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 4 deletions.
diff --git a/packages/contracts/src/utils/JwtRegistry.sol b/packages/contracts/src/utils/JwtRegistry.sol
@@ -85,11 +85,17 @@ contract JwtRegistry is IDKIMRegistry, Ownable {
         );
 
         dkimRegistry.revokeDKIMPublicKeyHash(publicKeyHash);
-        // Disable azp
-        string[] memory parts = this.stringToArray(domainName);
-        whitelistedClients[parts[2]] = false;
     }
 
+    /// @notice Disables the azp (authorized party) associated with the given domain name
+    /// @param domainName The domain name containing kis, iss, and azp fields
+    /// @dev This function removes the azp from the whitelisted clients
+    function disableAzp(string memory domainName) public {
+        string[] memory parts = this.stringToArray(domainName);
+        string memory azp = parts[2];
+        whitelistedClients[azp] = false;
+    }
+
     function stringToArray(string memory _strings) external pure returns (string[] memory) {
         strings.slice memory slicee = _strings.toSlice();
         strings.slice memory delim = "|".toSlice();

diff --git a/packages/contracts/test/JwtRegistry/JwtRegistry_disableAzp.t.sol b/packages/contracts/test/JwtRegistry/JwtRegistry_disableAzp.t.sol
@@ -0,0 +1,58 @@
+// 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";
+import {JwtRegistryTestBase} from "./JwtRegistryBase.t.sol";
+
+contract JwtRegistryTest_disableAzp is JwtRegistryTestBase {
+    constructor() {}
+
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function testRevert_disableAzp_invalidDomainNameFormat() public {
+        string memory invalidDomainName = "12345|https://example.com";
+        vm.expectRevert(bytes("Invalid kid|iss|azp strings"));
+        jwtRegistry.disableAzp(invalidDomainName);
+    }
+
+    function testRevert_disableAzp_tooManyParts() public {
+        string
+            memory invalidDomainName = "12345|https://example.com|client-id-12345|extra";
+        vm.expectRevert(bytes("Invalid kid|iss|azp strings"));
+        jwtRegistry.disableAzp(invalidDomainName);
+    }
+
+    function testRevert_disableAzp_emptyString() public {
+        string memory invalidDomainName = "";
+        vm.expectRevert(bytes("Invalid kid|iss|azp strings"));
+        jwtRegistry.disableAzp(invalidDomainName);
+    }
+
+    function test_disableAzp() public {
+        string memory domainName = "12345|https://example.com|client-id-12345";
+
+        // Verify that client-id-12345 is whitelisted
+        assertTrue(
+            jwtRegistry.whitelistedClients("client-id-12345"),
+            "Client should be whitelisted initially"
+        );
+
+        // Call disableAzp
+        jwtRegistry.disableAzp(domainName);
+
+        // Verify that client-id-12345 is no longer whitelisted
+        assertFalse(
+            jwtRegistry.whitelistedClients("client-id-12345"),
+            "Client should not be whitelisted after disableAzp"
+        );
+    }
+}
diff --git a/packages/contracts/test/JwtRegistry/JwtRegistry_revokeDKIMPublicKeyHash.t.sol b/packages/contracts/test/JwtRegistry/JwtRegistry_revokeDKIMPublicKeyHash.t.sol
@@ -39,6 +39,7 @@ contract JwtRegistryTest_revokeDKIMPublicKeyHash is JwtRegistryTestBase {
     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);
+        // revokeDKIMPublicKeyHash does not set azp to false
+        assertEq(jwtRegistry.whitelistedClients("client-id-12345"), true);
     }
 }