Skip to content

Commit

Permalink
Move tree update in a separate function
Browse files Browse the repository at this point in the history
  • Loading branch information
AllFi committed Nov 24, 2023
1 parent 17f2adf commit 95d8750
Show file tree
Hide file tree
Showing 6 changed files with 281 additions and 114 deletions.
53 changes: 39 additions & 14 deletions src/zkbob/ZkBobPool.sol
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,15 @@ import "../utils/Ownable.sol";
import "../proxy/EIP1967Admin.sol";
import "../interfaces/IEnergyRedeemer.sol";
import "../utils/ExternalSload.sol";
import {PriorityQueue, PriorityOperation} from "./utils/PriorityQueue.sol";

/**
* @title ZkBobPool
* Shielded transactions pool
*/
abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, ExternalSload {
using SafeERC20 for IERC20;
using PriorityQueue for PriorityQueue.Queue;

uint256 internal constant MAX_POOL_ID = 0xffffff;
bytes4 internal constant MESSAGE_PREFIX_COMMON_V1 = 0x00000000;
Expand Down Expand Up @@ -61,6 +63,8 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex

mapping(address => uint256) public accumulatedFee;

PriorityQueue.Queue pendingCommitments;

event UpdateOperatorManager(address manager);
event UpdateAccounting(address accounting);
event UpdateRedeemer(address redeemer);
Expand Down Expand Up @@ -147,6 +151,12 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
return TOKEN_NUMERATOR == 1 ? TOKEN_DENOMINATOR : (1 << 255) | TOKEN_NUMERATOR;
}

function pendingCommitment() external view returns (uint256 commitment) {
PriorityOperation memory op = pendingCommitments.front();
require(op.commitment != 0, "ZkBobPool: no pending commitment");
return op.commitment;
}

/**
* @dev Updates used accounting module.
* Callable only by the contract owner / proxy admin.
Expand Down Expand Up @@ -226,14 +236,15 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex

uint256 nullifier = _transfer_nullifier();
{
require(msg.sender == _memo_prover_address(), "ZkBobPool: unauthorized");
require(nullifiers[nullifier] == 0, "ZkBobPool: doublespend detected");
require(_transfer_index() <= poolIndex, "ZkBobPool: transfer index out of bounds");
require(transfer_verifier.verifyProof(_transfer_pub(), _transfer_proof()), "ZkBobPool: bad transfer proof");
require(tree_verifier.verifyProof(_tree_pub(roots[poolIndex]), _tree_proof()), "ZkBobPool: bad tree proof");

_appendCommitment(_transfer_out_commit(), uint64(_memo_tree_update_fee()));

nullifiers[nullifier] = uint256(keccak256(abi.encodePacked(_transfer_out_commit(), _transfer_delta())));
poolIndex += 128;
roots[poolIndex] = _tree_root_after();

bytes memory message = _memo_message();
// restrict memo message prefix (items count in little endian) to be < 2**16
require(bytes4(message) & 0x0000ffff == MESSAGE_PREFIX_COMMON_V1, "ZkBobPool: bad message prefix");
Expand All @@ -244,7 +255,10 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
emit Message(poolIndex, _all_messages_hash, message);
}

uint256 fee = _memo_fee();
uint256 transactFee = _memo_transact_fee();
uint256 treeUpdateFee = _memo_tree_update_fee();
uint256 fee = transactFee + treeUpdateFee;

int256 token_amount = transfer_token_delta + int256(fee);
int256 energy_amount = _transfer_energy_amount();

Expand Down Expand Up @@ -286,25 +300,21 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
}

if (fee > 0) {
accumulatedFee[msg.sender] += fee;
accumulatedFee[msg.sender] += transactFee;
}
}

/**
* @dev Appends a batch of direct deposits into a zkBob merkle tree.
* Callable only by the current operator.
* @param _root_after new merkle tree root after append.
* @param _indices list of indices for queued pending deposits.
* @param _out_commit out commitment for output notes serialized from direct deposits.
* @param _batch_deposit_proof snark proof for batch deposit verifier.
* @param _tree_proof snark proof for tree update verifier.
*/
function appendDirectDeposits(
uint256 _root_after,
uint256[] calldata _indices,
uint256 _out_commit,
uint256[8] memory _batch_deposit_proof,
uint256[8] memory _tree_proof
uint256[8] memory _batch_deposit_proof
)
external
onlyOperator
Expand All @@ -322,11 +332,9 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
batch_deposit_verifier.verifyProof([hashsum], _batch_deposit_proof), "ZkBobPool: bad batch deposit proof"
);

uint256[3] memory tree_pub = [roots[poolIndex], _root_after, _out_commit];
require(tree_verifier.verifyProof(tree_pub, _tree_proof), "ZkBobPool: bad tree proof");
// TODO: what is about fees in this case?
_appendCommitment(_out_commit, uint64(0));

poolIndex += 128;
roots[poolIndex] = _root_after;
bytes32 message_hash = keccak256(message);
bytes32 _all_messages_hash = keccak256(abi.encodePacked(all_messages_hash, message_hash));
all_messages_hash = _all_messages_hash;
Expand All @@ -339,6 +347,16 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
emit Message(poolIndex, _all_messages_hash, message);
}

function proveTreeUpdate(uint256 _commitment, uint256[8] calldata _proof, uint256 _rootAfter) external {
PriorityOperation memory pendindCommitment = pendingCommitments.popFront();
require(pendindCommitment.commitment == _commitment, "ZkBobPool: commitment mismatch");
uint256[3] memory tree_pub = [roots[pool_index], _rootAfter, _commitment];
require(tree_verifier.verifyProof(tree_pub, _proof), "ZkBobPool: bad tree proof");
pool_index += 128;
roots[pool_index] = _rootAfter;
accumulatedFee[msg.sender] += pendindCommitment.fee;
}

/**
* @dev Commits a forced withdrawal transaction for future execution after a set delay.
* Forced exits can be executed during 23 hours after 1 hour passed since its commitment.
Expand Down Expand Up @@ -512,4 +530,11 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
function _isOwner() internal view override returns (bool) {
return super._isOwner() || _admin() == _msgSender();
}

function _appendCommitment(uint256 _commitment, uint64 _fee) internal {
pendingCommitments.pushBack(PriorityOperation({
commitment: _commitment,
fee: _fee
}));
}
}
66 changes: 32 additions & 34 deletions src/zkbob/utils/CustomABIDecoder.sol
Original file line number Diff line number Diff line change
Expand Up @@ -55,24 +55,7 @@ contract CustomABIDecoder {
}
}

uint256 constant tree_root_after_pos = transfer_proof_pos + transfer_proof_size;
uint256 constant tree_root_after_size = 32;

function _tree_root_after() internal pure returns (uint256 r) {
r = _loaduint256(tree_root_after_pos);
}

uint256 constant tree_proof_pos = tree_root_after_pos + tree_root_after_size;
uint256 constant tree_proof_size = 256;

function _tree_proof() internal pure returns (uint256[8] calldata r) {
uint256 pos = tree_proof_pos;
assembly {
r := pos
}
}

uint256 constant tx_type_pos = tree_proof_pos + tree_proof_size;
uint256 constant tx_type_pos = transfer_proof_pos + transfer_proof_size;
uint256 constant tx_type_size = 2;
uint256 constant tx_type_mask = (1 << (tx_type_size * 8)) - 1;

Expand Down Expand Up @@ -124,17 +107,17 @@ contract CustomABIDecoder {
function _memo_fixed_size() internal pure returns (uint256 r) {
uint256 t = _tx_type();
if (t == 0 || t == 1) {
// fee
// 8
r = 8;
} else if (t == 2) {
// fee + native amount + recipient
// 8 + 8 + 20
// prover address + transact fee + tree update fee
// 20 + 8 + 8
r = 36;
} else if (t == 2) {
// prover address + transact fee + tree update fee + native amount + recipient
// 20 + 8 + 8 + 8 + 20
r = 64;
} else if (t == 3) {
// fee + deadline + address
// 8 + 8 + 20
r = 36;
// prover address + transact fee + tree update fee + deadline + address
// 20 + 8 + 8 + 8 + 20
r = 64;
} else {
revert();
}
Expand All @@ -150,17 +133,32 @@ contract CustomABIDecoder {
}
}

uint256 constant memo_fee_pos = memo_data_pos;
uint256 constant memo_fee_size = 8;
uint256 constant memo_fee_mask = (1 << (memo_fee_size * 8)) - 1;
uint256 constant memo_prover_address_pos = memo_data_pos;
uint256 constant memo_prover_address_size = 20;

function _memo_prover_address() internal pure returns (address r) {
r = address(uint160(_loaduint256(memo_prover_address_pos + memo_prover_address_size - uint256_size)));
}

uint256 constant memo_transact_fee_pos = memo_prover_address_pos + memo_prover_address_size;
uint256 constant memo_transact_fee_size = 8;
uint256 constant memo_transact_fee_mask = (1 << (memo_transact_fee_size * 8)) - 1;

function _memo_transact_fee() internal pure returns (uint256 r) {
r = _loaduint256(memo_transact_fee_pos + memo_transact_fee_size - uint256_size) & memo_transact_fee_mask;
}

uint256 constant memo_tree_update_fee_pos = memo_transact_fee_pos + memo_transact_fee_size;
uint256 constant memo_tree_update_fee_size = 8;
uint256 constant memo_tree_update_fee_mask = (1 << (memo_tree_update_fee_size * 8)) - 1;

function _memo_fee() internal pure returns (uint256 r) {
r = _loaduint256(memo_fee_pos + memo_fee_size - uint256_size) & memo_fee_mask;
function _memo_tree_update_fee() internal pure returns (uint256 r) {
r = _loaduint256(memo_tree_update_fee_pos + memo_tree_update_fee_size - uint256_size) & memo_tree_update_fee_mask;
}

// Withdraw specific data

uint256 constant memo_native_amount_pos = memo_fee_pos + memo_fee_size;
uint256 constant memo_native_amount_pos = memo_tree_update_fee_pos + memo_tree_update_fee_size;
uint256 constant memo_native_amount_size = 8;
uint256 constant memo_native_amount_mask = (1 << (memo_native_amount_size * 8)) - 1;

Expand All @@ -177,7 +175,7 @@ contract CustomABIDecoder {

// Permittable token deposit specific data

uint256 constant memo_permit_deadline_pos = memo_fee_pos + memo_fee_size;
uint256 constant memo_permit_deadline_pos = memo_tree_update_fee_pos + memo_tree_update_fee_size;
uint256 constant memo_permit_deadline_size = 8;

function _memo_permit_deadline() internal pure returns (uint64 r) {
Expand Down
6 changes: 0 additions & 6 deletions src/zkbob/utils/Parameters.sol
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,6 @@ abstract contract Parameters is CustomABIDecoder {
r[4] = uint256(keccak256(_memo_data())) % R;
}

function _tree_pub(uint256 _root_before) internal view returns (uint256[3] memory r) {
r[0] = _root_before;
r[1] = _tree_root_after();
r[2] = _transfer_out_commit();
}

// NOTE only valid in the context of normal deposit (tx_type=0)
function _deposit_spender() internal pure returns (address) {
(bytes32 r, bytes32 vs) = _sign_r_vs();
Expand Down
81 changes: 81 additions & 0 deletions src/zkbob/utils/PriorityQueue.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
// SPDX-License-Identifier: MIT

pragma solidity ^0.8.13;

/// @notice The structure that contains meta information of the L2 transaction that was requested from L1
/// @dev The weird size of fields was selected specifically to minimize the structure storage size
/// @param canonicalTxHash Hashed L2 transaction data that is needed to process it
/// @param expirationTimestamp Expiration timestamp for this request (must be satisfied before)
/// @param layer2Tip Additional payment to the validator as an incentive to perform the operation
struct PriorityOperation {
uint256 commitment;
uint64 fee;
}

/// @author Matter Labs
/// @custom:security-contact security@matterlabs.dev
/// @dev The library provides the API to interact with the priority queue container
/// @dev Order of processing operations from queue - FIFO (Fist in - first out)
library PriorityQueue {
using PriorityQueue for Queue;

/// @notice Container that stores priority operations
/// @param data The inner mapping that saves priority operation by its index
/// @param head The pointer to the first unprocessed priority operation, equal to the tail if the queue is empty
/// @param tail The pointer to the free slot
struct Queue {
mapping(uint256 => PriorityOperation) data;
uint256 tail;
uint256 head;
}

/// @notice Returns zero if and only if no operations were processed from the queue
/// @return Index of the oldest priority operation that wasn't processed yet
function getFirstUnprocessedPriorityTx(Queue storage _queue) internal view returns (uint256) {
return _queue.head;
}

/// @return The total number of priority operations that were added to the priority queue, including all processed ones
function getTotalPriorityTxs(Queue storage _queue) internal view returns (uint256) {
return _queue.tail;
}

/// @return The total number of unprocessed priority operations in a priority queue
function getSize(Queue storage _queue) internal view returns (uint256) {
return uint256(_queue.tail - _queue.head);
}

/// @return Whether the priority queue contains no operations
function isEmpty(Queue storage _queue) internal view returns (bool) {
return _queue.tail == _queue.head;
}

/// @notice Add the priority operation to the end of the priority queue
function pushBack(Queue storage _queue, PriorityOperation memory _operation) internal {
// Save value into the stack to avoid double reading from the storage
uint256 tail = _queue.tail;

_queue.data[tail] = _operation;
_queue.tail = tail + 1;
}

/// @return The first unprocessed priority operation from the queue
function front(Queue storage _queue) internal view returns (PriorityOperation memory) {
require(!_queue.isEmpty(), "D"); // priority queue is empty

return _queue.data[_queue.head];
}

/// @notice Remove the first unprocessed priority operation from the queue
/// @return priorityOperation that was popped from the priority queue
function popFront(Queue storage _queue) internal returns (PriorityOperation memory priorityOperation) {
require(!_queue.isEmpty(), "s"); // priority queue is empty

// Save value into the stack to avoid double reading from the storage
uint256 head = _queue.head;

priorityOperation = _queue.data[head];
delete _queue.data[head];
_queue.head = head + 1;
}
}
8 changes: 5 additions & 3 deletions test/interfaces/IZkBobPoolAdmin.sol
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ interface IZkBobPoolAdmin {

function pool_index() external view returns (uint256);

function pendingCommitment() external view returns (uint256);

function initialize(uint256 _root) external;

function setTokenSeller(address _tokenSeller) external;
Expand All @@ -27,6 +29,8 @@ interface IZkBobPoolAdmin {

function transact() external;

function proveTreeUpdate(uint256, uint256[8] memory, uint256) external;

function committedForcedExits(uint256 _nullifier) external view returns (bytes32);

function commitForcedExit(
Expand All @@ -52,11 +56,9 @@ interface IZkBobPoolAdmin {
external;

function appendDirectDeposits(
uint256 _root_after,
uint256[] calldata _indices,
uint256 _out_commit,
uint256[8] memory _batch_deposit_proof,
uint256[8] memory _tree_proof
uint256[8] memory _batch_deposit_proof
)
external;

Expand Down
Loading

0 comments on commit 95d8750

Please sign in to comment.