Add circuits for KYC enforcement with anonymity

Signed-off-by: Jim Zhang <[email protected]>

solidity/contracts/lib/registry.sol

@@ -16,35 +16,83 @@
 pragma solidity ^0.8.20;
 
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {SmtLib} from "@iden3/contracts/lib/SmtLib.sol";
+import {PoseidonUnit2L, PoseidonUnit3L} from "@iden3/contracts/lib/Poseidon.sol";
+import {Commonlib} from "./common.sol";
+import {Groth16Verifier_CheckSMTProof} from "./verifier_check_smt_proof.sol";
+import "hardhat/console.sol";
 
-/// @title A sample on-chain implementation of an account mapping between Ethereum addresses and BabyJubjub public keys
+uint256 constant MAX_SMT_DEPTH = 64;
+
+/// @title A sample on-chain implementation of a KYC registry
+/// @dev The registry uses a Sparse Merkle Tree to store the
+///   Zeto accounts (Babyjubjub keys), so that transaction
+///   submitters can generate proofs of membership for the
+///   accounts in a privacy-preserving manner.
 /// @author Kaleido, Inc.
-contract Registry is Ownable {
-    mapping(address => uint256[2]) private publicKeys;
+abstract contract Registry is Ownable {
+    SmtLib.Data internal _publicKeysTree;
+    using SmtLib for SmtLib.Data;
+
+    Groth16Verifier_CheckSMTProof private verifier;
+
+    error AlreadyRegistered(uint256[2]);
+
+    constructor(Groth16Verifier_CheckSMTProof _verifier) Ownable(msg.sender) {
+        verifier = _verifier;
+        _publicKeysTree.initialize(MAX_SMT_DEPTH);
+    }
 
-    error AlreadyRegistered(address addr);
+    modifier onlyRegistered(Commonlib.Proof calldata proof) {
+        uint256 root = _publicKeysTree.getRoot();
+        uint256[1] memory publicInputs;
+        publicInputs[0] = root;
 
-    constructor() Ownable(msg.sender) {}
+        // // Check the proof
+        require(
+            verifier.verifyProof(proof.pA, proof.pB, proof.pC, publicInputs),
+            "Identity not registered"
+        );
+        _;
+    }
 
-    /// @dev Register a new public key for the calling Ethereum address
-    /// @param ethAddress The Ethereum address to register
+    /// @dev Register a new Zeto account
     /// @param publicKey The public Babyjubjub key to register
-    function register(
-        address ethAddress,
-        uint256[2] memory publicKey
-    ) public onlyOwner {
-        if (publicKeys[ethAddress][0] != 0 || publicKeys[ethAddress][1] != 0) {
-            revert AlreadyRegistered(ethAddress);
+    function register(uint256[2] memory publicKey) public onlyOwner {
+        uint256 nodeHash = _getLeafNodeHash(publicKey);
+        SmtLib.Node memory node = _publicKeysTree.getNode(nodeHash);
+        if (node.nodeType != SmtLib.NodeType.EMPTY) {
+            revert AlreadyRegistered(publicKey);
         }
-        publicKeys[ethAddress] = publicKey;
+        _publicKeysTree.addLeaf(nodeHash, nodeHash);
+    }
+
+    /// @dev returns whether the given public key is registered
+    /// @param publicKey The Babyjubjub public key to check
+    /// @return bool whether the given public key is included in the registry
+    function isRegistered(
+        uint256[2] memory publicKey
+    ) public view returns (bool) {
+        uint256 nodeKey = _getLeafNodeKey(publicKey);
+        SmtLib.Node memory node = _publicKeysTree.getNode(nodeKey);
+        return node.nodeType != SmtLib.NodeType.EMPTY;
+    }
+
+    function getRoot() public view returns (uint256) {
+        return _publicKeysTree.getRoot();
     }
 
-    /// @dev Query the public key for a given Ethereum address
-    /// @param addr The Ethereum address to query
-    /// @return publicKey The public key for the given address
-    function getPublicKey(
-        address addr
-    ) public view returns (uint256[2] memory publicKey) {
-        return publicKeys[addr];
+    function _getLeafNodeHash(
+        uint256[2] memory publicKey
+    ) internal pure returns (uint256) {
+        return PoseidonUnit2L.poseidon(publicKey);
+    }
+
+    function _getLeafNodeKey(
+        uint256[2] memory publicKey
+    ) internal pure returns (uint256) {
+        uint256 nodeHash = PoseidonUnit2L.poseidon(publicKey);
+        uint256[3] memory params = [nodeHash, nodeHash, uint256(1)];
+        return PoseidonUnit3L.poseidon(params);
     }
 }

solidity/contracts/lib/test/MyContract.sol

@@ -0,0 +1,31 @@
+// Copyright © 2024 Kaleido, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pragma solidity ^0.8.20;
+
+import {Registry} from "../../lib/registry.sol";
+import {Commonlib} from "../common.sol";
+import {Groth16Verifier_CheckSMTProof} from "../verifier_check_smt_proof.sol";
+
+contract MyContract is Registry {
+    constructor(Groth16Verifier_CheckSMTProof _verifier) Registry(_verifier) {}
+
+    function myMethod(
+        Commonlib.Proof calldata kycProof
+    ) public onlyRegistered(kycProof) returns (bool) {
+        return true;
+    }
+}

solidity/contracts/lib/verifier_check_smt_proof.sol

@@ -0,0 +1,170 @@
+// SPDX-License-Identifier: GPL-3.0
+/*
+    Copyright 2021 0KIMS association.
+
+    This file is generated with [snarkJS](https://github.com/iden3/snarkjs).
+
+    snarkJS is a free software: you can redistribute it and/or modify it
+    under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    snarkJS is distributed in the hope that it will be useful, but WITHOUT
+    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+    or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
+    License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with snarkJS. If not, see <https://www.gnu.org/licenses/>.
+*/
+
+pragma solidity >=0.7.0 <0.9.0;
+
+contract Groth16Verifier_CheckSMTProof {
+    // Scalar field size
+    uint256 constant r    = 21888242871839275222246405745257275088548364400416034343698204186575808495617;
+    // Base field size
+    uint256 constant q   = 21888242871839275222246405745257275088696311157297823662689037894645226208583;
+
+    // Verification Key data
+    uint256 constant alphax  = 20491192805390485299153009773594534940189261866228447918068658471970481763042;
+    uint256 constant alphay  = 9383485363053290200918347156157836566562967994039712273449902621266178545958;
+    uint256 constant betax1  = 4252822878758300859123897981450591353533073413197771768651442665752259397132;
+    uint256 constant betax2  = 6375614351688725206403948262868962793625744043794305715222011528459656738731;
+    uint256 constant betay1  = 21847035105528745403288232691147584728191162732299865338377159692350059136679;
+    uint256 constant betay2  = 10505242626370262277552901082094356697409835680220590971873171140371331206856;
+    uint256 constant gammax1 = 11559732032986387107991004021392285783925812861821192530917403151452391805634;
+    uint256 constant gammax2 = 10857046999023057135944570762232829481370756359578518086990519993285655852781;
+    uint256 constant gammay1 = 4082367875863433681332203403145435568316851327593401208105741076214120093531;
+    uint256 constant gammay2 = 8495653923123431417604973247489272438418190587263600148770280649306958101930;
+    uint256 constant deltax1 = 11559732032986387107991004021392285783925812861821192530917403151452391805634;
+    uint256 constant deltax2 = 10857046999023057135944570762232829481370756359578518086990519993285655852781;
+    uint256 constant deltay1 = 4082367875863433681332203403145435568316851327593401208105741076214120093531;
+    uint256 constant deltay2 = 8495653923123431417604973247489272438418190587263600148770280649306958101930;
+
+
+    uint256 constant IC0x = 11441386565062873444786639656911805357241936257964157933969994950801496354882;
+    uint256 constant IC0y = 20022101915260480572828340426231221361899658863656101570036295881015826037670;
+
+    uint256 constant IC1x = 3308190451273467014980204609916919444233682137507309328633800984472104564895;
+    uint256 constant IC1y = 7638520442680558409849691682434554839963936218727440693299165844678822853364;
+
+
+    // Memory data
+    uint16 constant pVk = 0;
+    uint16 constant pPairing = 128;
+
+    uint16 constant pLastMem = 896;
+
+    function verifyProof(uint[2] calldata _pA, uint[2][2] calldata _pB, uint[2] calldata _pC, uint[1] calldata _pubSignals) public view returns (bool) {
+        assembly {
+            function checkField(v) {
+                if iszero(lt(v, q)) {
+                    mstore(0, 0)
+                    return(0, 0x20)
+                }
+            }
+
+            // G1 function to multiply a G1 value(x,y) to value in an address
+            function g1_mulAccC(pR, x, y, s) {
+                let success
+                let mIn := mload(0x40)
+                mstore(mIn, x)
+                mstore(add(mIn, 32), y)
+                mstore(add(mIn, 64), s)
+
+                success := staticcall(sub(gas(), 2000), 7, mIn, 96, mIn, 64)
+
+                if iszero(success) {
+                    mstore(0, 0)
+                    return(0, 0x20)
+                }
+
+                mstore(add(mIn, 64), mload(pR))
+                mstore(add(mIn, 96), mload(add(pR, 32)))
+
+                success := staticcall(sub(gas(), 2000), 6, mIn, 128, pR, 64)
+
+                if iszero(success) {
+                    mstore(0, 0)
+                    return(0, 0x20)
+                }
+            }
+
+            function checkPairing(pA, pB, pC, pubSignals, pMem) -> isOk {
+                let _pPairing := add(pMem, pPairing)
+                let _pVk := add(pMem, pVk)
+
+                mstore(_pVk, IC0x)
+                mstore(add(_pVk, 32), IC0y)
+
+                // Compute the linear combination vk_x
+
+                g1_mulAccC(_pVk, IC1x, IC1y, calldataload(add(pubSignals, 0)))
+
+
+                // -A
+                mstore(_pPairing, calldataload(pA))
+                mstore(add(_pPairing, 32), mod(sub(q, calldataload(add(pA, 32))), q))
+
+                // B
+                mstore(add(_pPairing, 64), calldataload(pB))
+                mstore(add(_pPairing, 96), calldataload(add(pB, 32)))
+                mstore(add(_pPairing, 128), calldataload(add(pB, 64)))
+                mstore(add(_pPairing, 160), calldataload(add(pB, 96)))
+
+                // alpha1
+                mstore(add(_pPairing, 192), alphax)
+                mstore(add(_pPairing, 224), alphay)
+
+                // beta2
+                mstore(add(_pPairing, 256), betax1)
+                mstore(add(_pPairing, 288), betax2)
+                mstore(add(_pPairing, 320), betay1)
+                mstore(add(_pPairing, 352), betay2)
+
+                // vk_x
+                mstore(add(_pPairing, 384), mload(add(pMem, pVk)))
+                mstore(add(_pPairing, 416), mload(add(pMem, add(pVk, 32))))
+
+
+                // gamma2
+                mstore(add(_pPairing, 448), gammax1)
+                mstore(add(_pPairing, 480), gammax2)
+                mstore(add(_pPairing, 512), gammay1)
+                mstore(add(_pPairing, 544), gammay2)
+
+                // C
+                mstore(add(_pPairing, 576), calldataload(pC))
+                mstore(add(_pPairing, 608), calldataload(add(pC, 32)))
+
+                // delta2
+                mstore(add(_pPairing, 640), deltax1)
+                mstore(add(_pPairing, 672), deltax2)
+                mstore(add(_pPairing, 704), deltay1)
+                mstore(add(_pPairing, 736), deltay2)
+
+
+                let success := staticcall(sub(gas(), 2000), 8, _pPairing, 768, _pPairing, 0x20)
+
+                isOk := and(success, mload(_pPairing))
+            }
+
+            let pMem := mload(0x40)
+            mstore(0x40, add(pMem, pLastMem))
+
+            // Validate that all evaluations ∈ F
+
+            checkField(calldataload(add(_pubSignals, 0)))
+
+            checkField(calldataload(add(_pubSignals, 32)))
+
+
+            // Validate all evaluations
+            let isValid := checkPairing(_pA, _pB, _pC, _pubSignals, pMem)
+
+            mstore(0, isValid)
+             return(0, 0x20)
+         }
+     }
+ }

solidity/contracts/lib/zeto_base.sol

@@ -35,7 +35,7 @@ abstract contract ZetoBase is ZetoCommon {
     // maintains all the UTXOs
     mapping(uint256 => UTXOStatus) internal _utxos;
 
-    constructor(Registry _registry) ZetoCommon(_registry) {}
+    constructor() ZetoCommon() {}
 
     /// @dev query whether a UTXO is currently spent
     /// @return bool whether the UTXO is spent

solidity/contracts/lib/zeto_common.sol

@@ -16,7 +16,6 @@
 pragma solidity ^0.8.20;
 
 import {Commonlib} from "./common.sol";
-import {Registry} from "./registry.sol";
 import {Groth16Verifier_CheckValue} from "./verifier_check_hashes_value.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
@@ -46,20 +45,14 @@ abstract contract ZetoCommon is Ownable {
     error UTXODuplicate(uint256 utxo);
     error IdentityNotRegistered(address addr);
 
-    // Registry contract to lookup public keys of the token owners
-    // This is a sample onchain KYC solution
-    Registry internal registry;
-
     // used for multi-step transaction flows that require counterparties
     // to upload proofs. To protect the party that uploads their proof first,
     // and prevent the other party from utilizing the uploaded proof to execute
     // a transaction, the proof can be locked and only usable by the same party
     // that did the locking.
     mapping(bytes32 => address) internal lockedProofs;
 
-    constructor(Registry _registry) Ownable(msg.sender) {
-        registry = _registry;
-    }
+    constructor() Ownable(msg.sender) {}
 
     // should be called by escrow contracts that will use uploaded proofs
     // to execute transactions, in order to prevent the proof from being used
@@ -69,13 +62,6 @@ abstract contract ZetoCommon is Ownable {
         lockedProofs[proofHash] = msg.sender;
     }
 
-    function validateIdentity(address receiver) internal view {
-        uint256[2] memory ownerPublicKey = registry.getPublicKey(receiver);
-        if (ownerPublicKey[0] == 0 || ownerPublicKey[1] == 0) {
-            revert IdentityNotRegistered(receiver);
-        }
-    }
-
     function sortInputsAndOutputs(
         uint256[2] memory inputs,
         uint256[2] memory outputs

solidity/contracts/lib/zeto_nullifier.sol

@@ -34,7 +34,7 @@ abstract contract ZetoNullifier is ZetoCommon {
 
     error UTXORootNotFound(uint256 root);
 
-    constructor(Registry _registry) ZetoCommon(_registry) {
+    constructor() ZetoCommon() {
         _commitmentsTree.initialize(MAX_SMT_DEPTH);
     }
 

solidity/contracts/zeto_anon.sol

@@ -39,12 +39,8 @@ contract Zeto_Anon is ZetoBase, ZetoFungibleWithdraw {
     constructor(
         Groth16Verifier_CheckValue _depositVerifier,
         Groth16Verifier_CheckInputsOutputsValue _withdrawVerifier,
-        Groth16Verifier_Anon _verifier,
-        Registry _registry
-    )
-        ZetoBase(_registry)
-        ZetoFungibleWithdraw(_depositVerifier, _withdrawVerifier)
-    {
+        Groth16Verifier_Anon _verifier
+    ) ZetoBase() ZetoFungibleWithdraw(_depositVerifier, _withdrawVerifier) {
         verifier = _verifier;
     }
 

solidity/contracts/zeto_anon_enc.sol

@@ -41,12 +41,8 @@ contract Zeto_AnonEnc is ZetoBase, ZetoFungibleWithdraw {
     constructor(
         Groth16Verifier_CheckValue _depositVerifier,
         Groth16Verifier_CheckInputsOutputsValue _withdrawVerifier,
-        Groth16Verifier_AnonEnc _verifier,
-        Registry _registry
-    )
-        ZetoBase(_registry)
-        ZetoFungibleWithdraw(_depositVerifier, _withdrawVerifier)
-    {
+        Groth16Verifier_AnonEnc _verifier
+    ) ZetoBase() ZetoFungibleWithdraw(_depositVerifier, _withdrawVerifier) {
         verifier = _verifier;
     }