RACE #31 Of The Secureum Bootcamp Epoch∞

use dep::std;

global DEPTH: Field = 4;

/// Compute the Merkle tree root
fn compute_merkle_root(key: [u1; DEPTH], leaf: Field, nodes: [Field; DEPTH]) -> Field {
    // Start with the `leaf` node.
    let mut node: Field = leaf;

    // Iterates from `0` until `DEPTH - 1`
    for i in 0..DEPTH {
        // Hash current node `node` with provided node `nodes[i]`
        // to the left or right with `nodes[i]` depending on `key`s i-th bit.
        node = std::hash::poseidon::bn254::hash_2(
            if (key[i] == 0) {
                [node, nodes[i]]
            } else {
                [nodes[i], node]
            });
    }

    node
}

/// Main circuit that verifies a zero knowledge proof, demonstrating:
///   - Knowledge of pre-image of leaf: `leaf == hash(secret + 1)`
///   - `nullifier` is derived correctly: `nullifier == hash(secret + 2)`
///   - The leaf is contained in a Merkle tree with root `root`
///   - The proof is generated for `receiver`
fn main(
    receiver: pub Field,
    key: [u1; DEPTH],
    secret: Field,
    nullifier: pub Field,
    nodes: [Field; DEPTH],
    root: pub Field
) {
    // Compute `leaf` using `secret`.
    let leaf: Field = std::hash::poseidon::bn254::hash_1([secret + 1]);

    // Assert given `nullifier` is derived from `secret`.
    assert(nullifier == std::hash::poseidon::bn254::hash_1([secret + 2]));

    // Ensure `leaf` is included in the Merkle tree.
    assert(compute_merkle_root(key, leaf, nodes) == root);

    // Tie `receiver` into proof.
    assert(receiver + leaf + nullifier != 0);
}

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import {UltraVerifier} from "../circuits/contract/ZeroLink/plonk_vk.sol";
import {PoseidonT3} from "./utils/Poseidon.sol";

uint256 constant DEPTH = 4;
uint256 constant PRIME_FIELD_ORDER = 21888242871839275222246405745257275088548364400416034343698204186575808495617;

library MerkleLib {
    function hash(uint256 left, uint256 right) internal pure returns (uint256 result) {
        require(left < PRIME_FIELD_ORDER && right < PRIME_FIELD_ORDER, "Invalid field element");
        return PoseidonT3.hash(left, right);
    }

    /// @notice Returns pre-computed zero sub-tree root of depth `level`.
    ///         Each sub-tree root is computed by:
    ///             root(0) := hash(0)
    ///             root(N) := hash(root(N-1), root(N-1))
    function zeros(uint256 level) internal pure returns (uint256) {
        if (level == 0) return 0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864;
        if (level == 1) return 0x1069673dcdb12263df301a6ff584a7ec261a44cb9dc68df067a4774460b1f1e1;
        if (level == 2) return 0x18f43331537ee2af2e3d758d50f72106467c6eea50371dd528d57eb2b856d238;
        if (level == 3) return 0x07f9d837cb17b0d36320ffe93ba52345f1b728571a568265caac97559dbc952a;
        if (level == 4) return 0x2b94cf5e8746b3f5c9631f4c5df32907a699c58c94b2ad4d7b5cec1639183f55;
        revert("Invalid level");
    }

    /// @notice Computes the Merkle root starting with `leaf` at given `key`.
    ///         Internal nodes are computed by hashing the current node with the
    ///         pre-computed zero subtrees to the right, or with the provided
    ///         nodes `nodes[i]` to the left, depending on the path given by `key`s i-th bit.
    function appendLeaf(uint256 key, uint256 leaf, uint256[DEPTH] memory nodes)
        internal
        pure
        returns (uint256 root, uint256[DEPTH] memory newNodes)
    {
        unchecked {
            require(key >> DEPTH == 0, "Invalid key");

            uint256 node = leaf;

            for (uint256 i; i < DEPTH; ++i) {
                newNodes[i] = node;
                node = ((key >> i) & 1 == 0) // Read `key`s i-th least-significant bit.
                    ? hash(node, zeros(i))
                    : hash(nodes[i], node);
            }

            root = node;
        }
    }
}

contract ZeroLink is UltraVerifier {
    uint256 public constant DEPOSIT_AMOUNT = 1 ether;

    uint256 public key;
    uint256 public root;
    uint256[DEPTH] public proofNodes;

    mapping(uint256 root => bool) public validRoot;
    mapping(uint256 nullifier => bool) public nullifierSpent;
    mapping(uint256 leaf => bool) public leafCommitted;

    /// @notice Create a deposit by committing a leaf node to an
    ///         append-only, fixed size Merkle tree. Every new leaf is
    ///         appended at `key`----the next available position in the Merkle tree.
    ///         `leaf` must be correctly derived from `secret`: `leaf == hash(secret + 1)`.
    function deposit(uint256 leaf) public payable {
        unchecked {
            require(msg.value == DEPOSIT_AMOUNT, "Invalid deposit amount");
            require(!leafCommitted[leaf], "Leaf already committed");

            // Mark the leaf as committed.
            leafCommitted[leaf] = true;

            // Append `leaf` at index `key` of Merkle tree.
            // Update Merkle root and internal nodes inserting `leaf` at index `key`.
            // Increment the Merkle tree index `key`.
            // Throws if `leaf` or any of `nodes` is not a field element.
            (root, proofNodes) = MerkleLib.appendLeaf(key++, leaf, proofNodes);

            // Mark new `root` as valid.
            validRoot[root] = true;
        }
    }

    function withdraw(address receiver, uint256 nullifier, uint256 root_, bytes calldata proof) public {
        // Verify the zero knowledge proof.
        verifyProof(receiver, nullifier, root_, proof);

        // Refund `receiver`.
        (bool success,) = receiver.call{value: DEPOSIT_AMOUNT}("");
        require(success, "Transfer failed");

        // Withdrawer's proof must relate to a previously committed root.
        require(validRoot[root], "Invalid root");

        // Check `nullifier` to prevent replay.
        require(!nullifierSpent[nullifier], "Nullifier already spent");

        // Mark `nullifier` as spent.
        nullifierSpent[nullifier] = true;
    }

    function verifyProof(address receiver, uint256 nullifier, uint256 root_, bytes calldata proof) public view {
        // Set up public inputs for `proof` verification.
        // The Noir circuit expects 3 public inputs.
        bytes32[] memory publicInputs = new bytes32[](3);

        // All inputs must be prime field elements.
        publicInputs[0] = bytes32(uint256(uint160(receiver)));
        publicInputs[1] = bytes32(nullifier % PRIME_FIELD_ORDER);
        publicInputs[2] = bytes32(root_ % PRIME_FIELD_ORDER);

        // Verify the zero knowledge proof.
        this.verify(proof, publicInputs);
    }
}

   root
    /
    \
     \
     /
  leaf