From 69e5fff4a30bcbfbb1140698be12990d943ba3ab Mon Sep 17 00:00:00 2001
From: wshino <wshinohara@gmail.com>
Date: Fri, 11 Oct 2024 14:59:56 +0900
Subject: [PATCH] Add disableAzp

---
 packages/contracts/src/utils/JwtRegistry.sol  | 12 +++-
 .../JwtRegistry/JwtRegistry_disableAzp.t.sol  | 58 +++++++++++++++++++
 .../JwtRegistry_revokeDKIMPublicKeyHash.t.sol |  3 +-
 3 files changed, 69 insertions(+), 4 deletions(-)
 create mode 100644 packages/contracts/test/JwtRegistry/JwtRegistry_disableAzp.t.sol

diff --git a/packages/contracts/src/utils/JwtRegistry.sol b/packages/contracts/src/utils/JwtRegistry.sol
index 0ed7d8f..4f3511e 100644
--- 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
new file mode 100644
index 0000000..578cc35
--- /dev/null
+++ 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
index 8e3f74e..ab15ad6 100644
--- 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);
     }
 }