Add gen-input script, integration tests for contracts, fixed JWT Registry

packages/circuits/jwt-verifier-test.circom

@@ -2,4 +2,4 @@ pragma circom 2.1.6;
 
 include "./jwt-verifier.circom";
 
-component main  = JWTVerifier(121, 17, 1024, 128, 896, 14, 605);
+component main  = JWTVerifier(121, 17, 1024, 128, 896, 72, 605);

packages/circuits/package.json

@@ -6,7 +6,8 @@
         "build": "mkdir -p build && circom jwt-verifier-test.circom --r1cs --wasm --sym -l ../../node_modules -o ./build",
         "publish": "yarn npm publish --access=public",
         "test": "NODE_OPTIONS=--max_old_space_size=8192 jest --runInBand --detectOpenHandles --forceExit --verbose tests",
-        "dev-setup": "NODE_OPTIONS=--max_old_space_size=16384 npx ts-node scripts/dev-setup.ts --output ./build"
+        "dev-setup": "NODE_OPTIONS=--max_old_space_size=16384 npx ts-node scripts/dev-setup.ts --output ./build",
+        "gen-input": "NODE_OPTIONS=--max_old_space_size=8192 npx ts-node scripts/gen-input.ts"
     },
     "dependencies": {
         "@zk-email/circuits": "6.1.5-nightly.2024-09-16",
@@ -33,6 +34,7 @@
         "/helpers",
         "/lib",
         "/utils",
+        "/scripts",
         "./jwt-verifier.circom"
     ],
     "babel": {

packages/circuits/scripts/gen-input.ts

@@ -0,0 +1,181 @@
+// @ts-ignore
+import { program } from 'commander';
+import { generateJWT } from '../../helpers/src/jwt'; // Specify the path to the function that generates JWT
+import { generateJWTVerifierInputs } from '../../helpers/src/input-generators'; // Specify the path to the function that generates inputs
+import { splitJWT } from "../../helpers/src/utils";
+import fs from "fs";
+const snarkjs = require("snarkjs");
+import { promisify } from "util";
+import path from "path";
+const relayerUtils = require("@zk-email/relayer-utils");
+import https from 'https';
+
+program
+  .requiredOption(
+    "--input-file <string>",
+    "Path of a json file to write the generated input"
+  )
+  .requiredOption('-a, --accountCode <string>', 'Account code as bigint string')
+  .option('-h, --header <string>', 'JWT header as JSON string')
+  .option('-p, --payload <string>', 'JWT payload as JSON string')
+  .option('-m, --maxMessageLength <number>', 'Maximum message length', '1024')
+  .option("--silent", "No console logs")
+  .option("--prove", "Also generate proof");
+
+program.parse(process.argv);
+
+const options = program.opts();
+
+function log(...message: any) {
+  if (!options.silent) {
+    console.log(...message);
+  }
+}
+
+async function main() {
+  const kid = BigInt("0x5aaff47c21d06e266cce395b2145c7c6d4730ea5");
+  const issuer = "random.website.com";
+  const timestamp = 1694989812;
+  const azp = "demo-client-id";
+  const email = "dummy@gmail.com";
+
+  const defaultHeader = {
+    alg: "RS256",
+    typ: "JWT",
+    kid: kid.toString(16),
+  };
+  const header = defaultHeader;
+  const defaultPayload = {
+    email,
+    iat: timestamp,
+    azp,
+    iss: issuer,
+  };
+  const payload = defaultPayload;
+  const accountCode = BigInt(options.accountCode);
+  const maxMessageLength = parseInt(options.maxMessageLength, 1024);
+
+
+  const { rawJWT, publicKey } = generateJWT(header, {
+    ...payload,
+    nonce: "Send 0.1 ETH to alice@gmail.com",
+  });
+
+  const jwtVerifierInputs = await generateJWTVerifierInputs(
+    rawJWT,
+    publicKey,
+    accountCode,
+    { maxMessageLength }
+  );
+
+  console.log('JWT Verifier Inputs:', jwtVerifierInputs);
+
+  const publicKeyHash = relayerUtils.publicKeyHash(
+    "0x" + Buffer.from(publicKey.n, "base64").toString("hex")
+  );
+  console.log("publicKeyHash");
+  console.log(publicKeyHash);
+  const [, , signature] = splitJWT(rawJWT);
+  const expectedJwtNullifier = relayerUtils.emailNullifier(
+    "0x" + Buffer.from(signature, "base64").toString("hex")
+  );
+  console.log("expectedJwtNullifier");
+  console.log(expectedJwtNullifier);
+
+  if (!options.inputFile.endsWith(".json")) {
+    throw new Error("--input file path arg must end with .json");
+  }
+
+  // log("Generating Inputs for:", options);
+
+  // const circuitInputs = await genEmailCircuitInput(args.emailFile, args.accountCode, {
+  //     maxHeaderLength: 1024,
+  //     ignoreBodyHashCheck: true
+  // });
+  // log("\n\nGenerated Inputs:", circuitInputs, "\n\n");
+  const processedInputs = convertBigIntFieldsToString(jwtVerifierInputs);
+
+  await promisify(fs.writeFile)(options.inputFile, JSON.stringify(processedInputs, null, 2));
+
+  log("Inputs written to", options.inputFile);
+
+  if (options.prove) {
+    console.log("generate pub signal");
+    const fileContent = fs.readFileSync(options.inputFile as string, 'utf-8');
+    const jsonData = JSON.parse(fileContent);
+    const payload = JSON.stringify({ input: jsonData });
+    const urlObject = new URL("https://zkemail--jwt-prover-v0-1-0-flask-app.modal.run/prove/jwt");
+    const reqOptions = {
+      hostname: urlObject.hostname,
+      path: urlObject.pathname,
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Content-Length': Buffer.byteLength(payload)
+      }
+    };
+    await new Promise<void>((resolve, reject) => {
+      const req = https.request(reqOptions, (res) => {
+        let data = '';
+
+        res.on('data', (chunk) => {
+          data += chunk;
+        });
+
+        res.on('end', async () => {
+          try {
+            const dir = path.dirname(options.inputFile);
+            const responseJson = JSON.parse(data);
+            const proof = responseJson.proof;
+            // console.log(proof);
+            const publicSignals = responseJson.pub_signals;
+
+            await fs.promises.writeFile(
+              path.join(dir, "proof.json"),
+              JSON.stringify(convertBigIntFieldsToString(proof), null, 2)
+            );
+
+            await fs.promises.writeFile(
+              path.join(dir, "public.json"),
+              JSON.stringify(convertBigIntFieldsToString(publicSignals), null, 2)
+            );
+            console.log('Files written successfully');
+            resolve();
+          } catch (error) {
+            console.error('Error processing response:', error);
+            reject(error);
+          }
+        });
+      });
+
+      req.on('error', (error) => {
+        console.error('Error posting JSON data:', error);
+        reject(error);
+      });
+
+      req.write(payload);
+      req.end();
+    });
+
+  };
+  // Create the request
+
+  process.exit(0);
+}
+
+function convertBigIntFieldsToString(obj: any): any {
+  if (typeof obj === 'object' && obj !== null) {
+    return Object.fromEntries(
+      Object.entries(obj).map(([key, value]) => [
+        key,
+        typeof value === 'bigint' ? value.toString() : value
+      ])
+    );
+  }
+  return obj;
+}
+
+main().catch((err) => {
+  console.error("Error generating inputs", err);
+  process.exit(1);
+});

packages/contracts/foundry.toml

@@ -4,6 +4,9 @@ out = "artifacts"
 libs = ["../../node_modules", "lib"]
 optimizer = true
 optimizer-runs = 20_000
+fs_permissions = [
+    { access = "read", path = "./test/jwt_proofs" },
+]
 
 solc = "0.8.26"
 

packages/contracts/src/interfaces/IJwtGroth16Verifier.sol

@@ -6,6 +6,6 @@ interface IJwtGroth16Verifier {
         uint[2] calldata _pA, 
         uint[2][2] calldata _pB, 
         uint[2] calldata _pC, 
-        uint[29] calldata _pubSignals
+        uint[31] calldata _pubSignals
     ) external view returns (bool);
 }

packages/contracts/src/interfaces/IVerifier.sol

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.9;
+
+struct EmailProof {
+    string domainName; // Domain name of the sender's email
+    bytes32 publicKeyHash; // Hash of the DKIM public key used in email/proof
+    uint timestamp; // Timestamp of the email
+    string maskedCommand; // Masked command of the email
+    bytes32 emailNullifier; // Nullifier of the email to prevent its reuse.
+    bytes32 accountSalt; // Create2 salt of the account
+    bool isCodeExist; // Check if the account code is exist
+    bytes proof; // ZK Proof of Email
+}
+
+interface IVerifier {
+
+    /**
+     * @notice Verifies the provided email proof.
+     * @param proof The email proof to be verified.
+     * @return bool indicating whether the proof is valid.
+     */    
+    function verifyEmailProof(
+        EmailProof memory proof
+    ) external view returns (bool);
+
+    /**
+     * @notice Returns a constant value representing command bytes.
+     * @return uint256 The constant value of command bytes.
+     */
+    function getCommandBytes() external pure returns (uint256);
+}

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();

packages/contracts/src/utils/JwtVerifier.sol

@@ -5,19 +5,9 @@ import "../interfaces/IJwtGroth16Verifier.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import { strings } from "solidity-stringutils/src/strings.sol";
+import {IVerifier, EmailProof} from "../interfaces/IVerifier.sol";
 
-struct EmailProof {
-    string domainName; // Domain name of the sender's email
-    bytes32 publicKeyHash; // Hash of the DKIM public key used in email/proof
-    uint timestamp; // Timestamp of the email
-    string maskedCommand; // Masked command of the email
-    bytes32 emailNullifier; // Nullifier of the email to prevent its reuse.
-    bytes32 accountSalt; // Create2 salt of the account
-    bool isCodeExist; // Check if the account code is exist
-    bytes proof; // ZK Proof of Email
-}
-
-contract JwtVerifier is OwnableUpgradeable, UUPSUpgradeable {
+contract JwtVerifier is IVerifier, OwnableUpgradeable, UUPSUpgradeable {
     using strings for *;
 
     IJwtGroth16Verifier groth16Verifier;
@@ -26,8 +16,8 @@ contract JwtVerifier is OwnableUpgradeable, UUPSUpgradeable {
     uint256 public constant ISS_BYTES = 32;
     uint256 public constant COMMAND_FIELDS = 20;
     uint256 public constant COMMAND_BYTES = 605;
-    uint256 public constant AZP_FIELDS = 1;
-    uint256 public constant AZP_BYTES = 14;
+    uint256 public constant AZP_FIELDS = 3;
+    uint256 public constant AZP_BYTES = 72;
 
     constructor() {}
 

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"
+        );
+    }
+}

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);
     }
 }