Add ZKJWTUtils

.github/CODEOWNERS

@@ -18,5 +18,6 @@
 ## ZKEmail
 
 /contracts/utils/cryptography/ZKEmailUtils.sol @zkfriendly @benceharomi
+/contracts/utils/cryptography/ZKJWTUtils.sol @zkfriendly @benceharomi
 /contracts/utils/cryptography/signers/SignerZKEmail.sol @zkfriendly @benceharomi
 /contracts/utils/cryptography/verifiers/ERC7913ZKEmailVerifier.sol @zkfriendly @benceharomi

.gitmodules

@@ -20,3 +20,6 @@
 	path = lib/zk-email-verify
 	branch = v6.3.2
 	url = https://github.com/zkemail/zk-email-verify
+[submodule "lib/zk-jwt"]
+	path = lib/zk-jwt
+	url = https://github.com/zkemail/zk-jwt

contracts/mocks/import.sol

@@ -3,6 +3,7 @@
 pragma solidity ^0.8.20;
 
 import {ECDSAOwnedDKIMRegistry} from "@zk-email/email-tx-builder/src/utils/ECDSAOwnedDKIMRegistry.sol";
+import {JwtRegistry} from "@zk-email/zk-jwt/src/utils/JwtRegistry.sol";
 import {UpgradeableBeacon} from "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol";
 import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
 import {ERC721Enumerable} from "@openzeppelin/contracts/token/ERC721/extensions/ERC721Enumerable.sol";

contracts/mocks/utils/cryptography/ZKJWTVerifierMock.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.20;
+
+import {IVerifier, EmailProof} from "@zk-email/zk-jwt/src/interfaces/IVerifier.sol";
+
+contract ZKJWTVerifierMock is IVerifier {
+    function getCommandBytes() external pure returns (uint256) {
+        // Same as in https://github.com/zkemail/zk-jwt/blob/27436a2f23e78e89cf624f649ec1d125f13772dd/packages/contracts/src/utils/JwtVerifier.sol#L20
+        return 605;
+    }
+
+    function verifyEmailProof(EmailProof memory proof) external pure returns (bool) {
+        return proof.proof.length > 0 && bytes1(proof.proof[0]) == 0x01; // boolean true
+    }
+}

contracts/utils/cryptography/ZKJWTUtils.sol

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {Bytes} from "@openzeppelin/contracts/utils/Bytes.sol";
+import {Strings} from "@openzeppelin/contracts/utils/Strings.sol";
+import {IDKIMRegistry} from "@zk-email/contracts/DKIMRegistry.sol";
+import {IVerifier, EmailProof} from "@zk-email/zk-jwt/src/interfaces/IVerifier.sol";
+import {CommandUtils} from "@zk-email/email-tx-builder/src/libraries/CommandUtils.sol";
+
+/**
+ * @dev Library for https://docs.zk.email[ZK JWT] validation utilities.
+ *
+ * ZK JWT is a protocol that enables JWT-based authentication and authorization for smart contracts
+ * using zero-knowledge proofs. It allows users to prove ownership of a JWT token without revealing
+ * the token content or private keys. See https://datatracker.ietf.org/doc/html/rfc7519[RFC-7519] for
+ * details on the JWT verification process.
+ *
+ * The validation process involves several key components:
+ *
+ * * A https://docs.zk.email/jwt-tx-builder/architecture[JWT Registry] verification mechanism to ensure
+ * the JWT was issued from a valid issuer with valid public key. The registry validates the `kid|iss|azp` format
+ * used in JWT verification (i.e. key id, issuer, and authorized party).
+ * * A https://docs.zk.email/email-tx-builder/architecture/command-templates[command template] validation
+ * mechanism to ensure the JWT command matches the expected format and parameters.
+ * * A https://docs.zk.email/jwt-tx-builder/architecture[zero-knowledge proof] verification mechanism to ensure
+ * the JWT was actually issued and received without revealing its contents through RSA signature validation
+ * and selective claim disclosure.
+ *
+ * NOTE: This library adapts the email authentication infrastructure for JWT verification,
+ * reusing the EmailProof structure which contains JWT-specific data encoded in the domainName field
+ * as "kid|iss|azp" format.
+ */
+library ZKJWTUtils {
+    using CommandUtils for bytes[];
+    using Bytes for bytes;
+    using Strings for string;
+
+    /// @dev Enumeration of possible JWT proof validation errors.
+    /// See https://docs.zk.email/jwt-tx-builder/architecture[ZK JWT Architecture] for validation flow details.
+    enum JWTProofError {
+        NoError,
+        JWTPublicKeyHash, // The JWT public key hash verification fails
+        MaskedCommandLength, // The masked command length exceeds the maximum allowed by the circuit
+        MismatchedCommand, // The command does not match the proof command
+        JWTProof // The JWT proof verification fails
+    }
+
+    /// @dev Enumeration of possible string cases used to compare the command with the expected proven command.
+    enum Case {
+        CHECKSUM, // Computes a checksum of the command.
+        LOWERCASE, // Converts the command to hex lowercase.
+        UPPERCASE, // Converts the command to hex uppercase.
+        ANY
+    }
+
+    /// @dev Validates a ZK JWT proof with default "signHash" command template.
+    function isValidZKJWT(
+        EmailProof memory jwtProof,
+        IDKIMRegistry jwtRegistry,
+        IVerifier verifier,
+        bytes32 hash
+    ) internal view returns (JWTProofError) {
+        string[] memory signHashTemplate = new string[](1);
+        signHashTemplate[0] = CommandUtils.UINT_MATCHER;
+        bytes[] memory signHashParams = new bytes[](1);
+        signHashParams[0] = abi.encode(hash);
+        return isValidZKJWT(jwtProof, jwtRegistry, verifier, signHashTemplate, signHashParams, Case.LOWERCASE); // UINT_MATCHER is always lowercase
+    }
+
+    /**
+     * @dev Validates a ZK JWT proof against a command template.
+     *
+     * This function takes a JWT proof, a JWT registry contract, and a verifier contract
+     * as inputs. It performs several validation checks and returns a {JWTProofError} indicating the result.
+     * Returns {JWTProofError-NoError} if all validations pass, or a specific {JWTProofError} indicating
+     * which validation check failed.
+     *
+     * NOTE: Attempts to validate the command for all possible string {Case} values.
+     */
+    function isValidZKJWT(
+        EmailProof memory jwtProof,
+        IDKIMRegistry jwtRegistry,
+        IVerifier verifier,
+        string[] memory template,
+        bytes[] memory templateParams
+    ) internal view returns (JWTProofError) {
+        return isValidZKJWT(jwtProof, jwtRegistry, verifier, template, templateParams, Case.ANY);
+    }
+
+    /**
+     * @dev Validates a ZK JWT proof against a template with a specific string {Case}.
+     *
+     * Useful for templates with Ethereum address matchers (i.e. `{ethAddr}`), which are case-sensitive
+     * (e.g., `["someCommand", "{address}"]`).
+     */
+    function isValidZKJWT(
+        EmailProof memory jwtProof,
+        IDKIMRegistry jwtRegistry,
+        IVerifier verifier,
+        string[] memory template,
+        bytes[] memory templateParams,
+        Case stringCase
+    ) internal view returns (JWTProofError) {
+        if (bytes(jwtProof.maskedCommand).length > verifier.getCommandBytes()) {
+            return JWTProofError.MaskedCommandLength;
+        } else if (!_commandMatch(jwtProof, template, templateParams, stringCase)) {
+            return JWTProofError.MismatchedCommand;
+        } else if (!jwtRegistry.isDKIMPublicKeyHashValid(jwtProof.domainName, jwtProof.publicKeyHash)) {
+            // Validate JWT public key and authorized party through registry
+            // The domainName contains "kid|iss|azp" format for JWT validation
+            return JWTProofError.JWTPublicKeyHash;
+        }
+
+        // TODO: Can we remove the try catch? Or add it to ZKEmailUtils.sol?
+        // Verify the zero-knowledge proof of JWT signature
+        try verifier.verifyEmailProof(jwtProof) returns (bool isValid) {
+            return isValid ? JWTProofError.NoError : JWTProofError.JWTProof;
+        } catch {
+            return JWTProofError.JWTProof;
+        }
+    }
+
+    /// @dev Compares the command in the JWT proof with the expected command template.
+    function _commandMatch(
+        EmailProof memory jwtProof,
+        string[] memory template,
+        bytes[] memory templateParams,
+        Case stringCase
+    ) private pure returns (bool) {
+        // Convert template to expected command format
+        uint256 commandPrefixLength = bytes(jwtProof.maskedCommand).indexOf(bytes1(" "));
+        string memory command = string(bytes(jwtProof.maskedCommand).slice(commandPrefixLength + 1));
+
+        if (stringCase != Case.ANY)
+            return templateParams.computeExpectedCommand(template, uint8(stringCase)).equal(command);
+        return
+            templateParams.computeExpectedCommand(template, uint8(Case.LOWERCASE)).equal(command) ||
+            templateParams.computeExpectedCommand(template, uint8(Case.UPPERCASE)).equal(command) ||
+            templateParams.computeExpectedCommand(template, uint8(Case.CHECKSUM)).equal(command);
+    }
+}

remappings.txt

@@ -3,4 +3,5 @@
 @openzeppelin/community-contracts/=contracts/
 @axelar-network/axelar-gmp-sdk-solidity/=lib/axelar-gmp-sdk-solidity/
 @zk-email/email-tx-builder/=lib/email-tx-builder/packages/contracts/
+@zk-email/zk-jwt/=lib/zk-jwt/packages/contracts/
 @zk-email/contracts/=lib/zk-email-verify/packages/contracts/

test/helpers/enums.js

@@ -11,6 +11,7 @@ module.exports = {
     'InvalidFieldPoint',
     'EmailProof',
   ),
+  JWTProofError: enums.Enum('NoError', 'JWTPublicKeyHash', 'MaskedCommandLength', 'MismatchedCommand', 'JWTProof'),
   Case: enums.EnumTyped('CHECKSUM', 'LOWERCASE', 'UPPERCASE', 'ANY'),
   OperationState: enums.Enum('Unknown', 'Scheduled', 'Ready', 'Expired', 'Executed', 'Canceled'),
 };

test/utils/cryptography/ZKEmailUtils.t.sol

@@ -133,11 +133,7 @@ contract ZKEmailUtilsTest is Test {
         _mockVerifyEmailProof();
 
         // Test validation
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.NoError));
     }
@@ -180,7 +176,7 @@ contract ZKEmailUtilsTest is Test {
 
         ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
             emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
+            _dkimRegistry,
             _verifier,
             template
         );
@@ -206,7 +202,7 @@ contract ZKEmailUtilsTest is Test {
         commandParams[0] = abi.encode(addr);
 
         // Test with different cases
-        for (uint256 i = 0; i < uint8(type(ZKEmailUtils.Case).max) - 1; i++) {
+        for (uint256 i = 0; i < uint8(type(ZKEmailUtils.Case).max); i++) {
             EmailAuthMsg memory emailAuthMsg = _buildEmailAuthMsgMock(
                 string.concat(commandPrefix, " ", CommandUtils.addressToHexString(addr, i)),
                 commandParams,
@@ -228,7 +224,7 @@ contract ZKEmailUtilsTest is Test {
 
             ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
                 emailAuthMsg,
-                IDKIMRegistry(_dkimRegistry),
+                _dkimRegistry,
                 _verifier,
                 template,
                 ZKEmailUtils.Case(i)
@@ -297,11 +293,7 @@ contract ZKEmailUtilsTest is Test {
         emailAuthMsg.proof.domainName = domainName;
         emailAuthMsg.proof.publicKeyHash = publicKeyHash;
 
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.DKIMPublicKeyHash));
     }
@@ -314,11 +306,7 @@ contract ZKEmailUtilsTest is Test {
 
         EmailAuthMsg memory emailAuthMsg = _buildEmailAuthMsgMock(string(new bytes(length)), commandParams, 0);
 
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.MaskedCommandLength));
     }
@@ -336,11 +324,7 @@ contract ZKEmailUtilsTest is Test {
             skippedPrefix
         );
 
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.SkippedCommandPrefixSize));
     }
@@ -351,11 +335,7 @@ contract ZKEmailUtilsTest is Test {
 
         EmailAuthMsg memory emailAuthMsg = _buildEmailAuthMsgMock(invalidCommand, commandParams, 0);
 
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.MismatchedCommand));
     }
@@ -379,11 +359,7 @@ contract ZKEmailUtilsTest is Test {
 
         emailAuthMsg.proof.proof = abi.encode(pA, pB, pC);
 
-        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(
-            emailAuthMsg,
-            IDKIMRegistry(_dkimRegistry),
-            _verifier
-        );
+        ZKEmailUtils.EmailProofError err = ZKEmailUtils.isValidZKEmail(emailAuthMsg, _dkimRegistry, _verifier);
 
         assertEq(uint256(err), uint256(ZKEmailUtils.EmailProofError.EmailProof));
     }