diff --git a/contracts/evm/src/QuorumCustody.sol b/contracts/evm/src/QuorumCustody.sol index 7df2e4c..14c88a2 100644 --- a/contracts/evm/src/QuorumCustody.sol +++ b/contracts/evm/src/QuorumCustody.sol @@ -13,6 +13,7 @@ import {IDeposit} from "./interfaces/IDeposit.sol"; contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { using SafeERC20 for IERC20; + // Contract-specific errors error InvalidSigner(); error NotSigner(); error AlreadySigner(); @@ -154,12 +155,12 @@ contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { } function deposit(address token, uint256 amount) external payable override nonReentrant { - require(amount != 0, ZeroAmount()); + require(amount != 0, IDeposit.ZeroAmount()); if (token == address(0)) { - require(msg.value == amount, MsgValueMismatch()); + require(msg.value == amount, IDeposit.InvalidMsgValue()); } else { - require(msg.value == 0, NonZeroMsgValueForERC20()); + require(msg.value == 0, IDeposit.InvalidMsgValue()); IERC20(token).safeTransferFrom(msg.sender, address(this), amount); } @@ -174,10 +175,10 @@ contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { returns (bytes32) { require(user != address(0), InvalidUser()); - require(amount != 0, ZeroAmount()); + require(amount != 0, IDeposit.ZeroAmount()); bytes32 withdrawalId = _getWithdrawalId(user, token, amount, nonce); - require(withdrawals[withdrawalId].createdAt == 0, WithdrawalAlreadyExists()); + require(withdrawals[withdrawalId].createdAt == 0, IWithdraw.WithdrawalAlreadyExists()); withdrawals[withdrawalId] = WithdrawalRequest({ user: user, @@ -196,8 +197,8 @@ contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { WithdrawalRequest storage request = withdrawals[withdrawalId]; address signer = msg.sender; - require(request.createdAt != 0, WithdrawalNotFound()); - require(!request.finalized, WithdrawalAlreadyFinalized()); + require(request.createdAt != 0, IWithdraw.WithdrawalNotFound()); + require(!request.finalized, IWithdraw.WithdrawalAlreadyFinalized()); require(block.timestamp <= request.createdAt + OPERATION_EXPIRY, WithdrawalExpired()); require(!withdrawalApprovals[withdrawalId][signer], SignerAlreadyApproved()); @@ -215,8 +216,8 @@ contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { function rejectWithdraw(bytes32 withdrawalId) external override onlySigner nonReentrant { WithdrawalRequest storage request = withdrawals[withdrawalId]; - require(request.createdAt != 0, WithdrawalNotFound()); - require(!request.finalized, WithdrawalAlreadyFinalized()); + require(request.createdAt != 0, IWithdraw.WithdrawalNotFound()); + require(!request.finalized, IWithdraw.WithdrawalAlreadyFinalized()); require(block.timestamp > request.createdAt + OPERATION_EXPIRY, WithdrawalNotExpired()); request.finalized = true; @@ -241,11 +242,11 @@ contract QuorumCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712 { request.finalized = true; if (token == address(0)) { - require(address(this).balance >= amount, InsufficientLiquidity()); + require(address(this).balance >= amount, IWithdraw.InsufficientLiquidity()); (bool success,) = user.call{value: amount}(""); - require(success, ETHTransferFailed()); + require(success, IWithdraw.ETHTransferFailed()); } else { - require(IERC20(token).balanceOf(address(this)) >= amount, InsufficientLiquidity()); + require(IERC20(token).balanceOf(address(this)) >= amount, IWithdraw.InsufficientLiquidity()); IERC20(token).safeTransfer(user, amount); } diff --git a/contracts/evm/src/SimpleCustody.sol b/contracts/evm/src/SimpleCustody.sol index bc5099d..630d57b 100644 --- a/contracts/evm/src/SimpleCustody.sol +++ b/contracts/evm/src/SimpleCustody.sol @@ -32,12 +32,12 @@ contract SimpleCustody is IWithdraw, IDeposit, AccessControl, ReentrancyGuard { } function deposit(address token, uint256 amount) external payable override nonReentrant { - if (amount == 0) revert ZeroAmount(); + if (amount == 0) revert IDeposit.ZeroAmount(); uint256 received = amount; if (token == address(0)) { - if (msg.value != amount) revert MsgValueMismatch(); + if (msg.value != amount) revert IDeposit.InvalidMsgValue(); } else { - if (msg.value != 0) revert NonZeroMsgValueForERC20(); + if (msg.value != 0) revert IDeposit.InvalidMsgValue(); uint256 balanceBefore = IERC20(token).balanceOf(address(this)); IERC20(token).safeTransferFrom(msg.sender, address(this), amount); received = IERC20(token).balanceOf(address(this)) - balanceBefore; @@ -52,10 +52,10 @@ contract SimpleCustody is IWithdraw, IDeposit, AccessControl, ReentrancyGuard { nonReentrant returns (bytes32 withdrawalId) { - if (amount == 0) revert ZeroAmount(); + if (amount == 0) revert IDeposit.ZeroAmount(); withdrawalId = keccak256(abi.encode(block.chainid, address(this), user, token, amount, nonce)); - if (withdrawals[withdrawalId].exists) revert WithdrawalAlreadyExists(); + if (withdrawals[withdrawalId].exists) revert IWithdraw.WithdrawalAlreadyExists(); withdrawals[withdrawalId] = WithdrawalRequest({user: user, token: token, amount: amount, exists: true, finalized: false}); @@ -65,8 +65,8 @@ contract SimpleCustody is IWithdraw, IDeposit, AccessControl, ReentrancyGuard { function finalizeWithdraw(bytes32 withdrawalId) external override onlyRole(NITEWATCH_ROLE) nonReentrant { WithdrawalRequest storage request = withdrawals[withdrawalId]; - if (!request.exists) revert WithdrawalNotFound(); - if (request.finalized) revert WithdrawalAlreadyFinalized(); + if (!request.exists) revert IWithdraw.WithdrawalNotFound(); + if (request.finalized) revert IWithdraw.WithdrawalAlreadyFinalized(); request.finalized = true; address user = request.user; @@ -79,11 +79,11 @@ contract SimpleCustody is IWithdraw, IDeposit, AccessControl, ReentrancyGuard { request.amount = 0; if (token == address(0)) { - if (address(this).balance < amount) revert InsufficientLiquidity(); + if (address(this).balance < amount) revert IWithdraw.InsufficientLiquidity(); (bool success,) = user.call{value: amount}(""); - if (!success) revert ETHTransferFailed(); + if (!success) revert IWithdraw.ETHTransferFailed(); } else { - if (IERC20(token).balanceOf(address(this)) < amount) revert InsufficientLiquidity(); + if (IERC20(token).balanceOf(address(this)) < amount) revert IWithdraw.InsufficientLiquidity(); IERC20(token).safeTransfer(user, amount); } @@ -92,8 +92,8 @@ contract SimpleCustody is IWithdraw, IDeposit, AccessControl, ReentrancyGuard { function rejectWithdraw(bytes32 withdrawalId) external override onlyRole(NITEWATCH_ROLE) nonReentrant { WithdrawalRequest storage request = withdrawals[withdrawalId]; - if (!request.exists) revert WithdrawalNotFound(); - if (request.finalized) revert WithdrawalAlreadyFinalized(); + if (!request.exists) revert IWithdraw.WithdrawalNotFound(); + if (request.finalized) revert IWithdraw.WithdrawalAlreadyFinalized(); request.finalized = true; diff --git a/contracts/evm/src/ThresholdCustody.sol b/contracts/evm/src/ThresholdCustody.sol index 437ff11..67d1257 100644 --- a/contracts/evm/src/ThresholdCustody.sol +++ b/contracts/evm/src/ThresholdCustody.sol @@ -9,19 +9,31 @@ import {MultiSignerERC7913} from "@openzeppelin/contracts/utils/cryptography/sig import {IWithdraw} from "./interfaces/IWithdraw.sol"; import {IDeposit} from "./interfaces/IDeposit.sol"; +import {Utils} from "./Utils.sol"; + +bytes32 constant SET_THRESHOLD_TYPEHASH = + keccak256("SetThreshold(uint256 newThreshold,uint256 nonce,uint256 deadline)"); +bytes32 constant ADD_SIGNERS_TYPEHASH = + keccak256("AddSigners(address[] newSigners,uint256 newThreshold,uint256 nonce,uint256 deadline)"); +bytes32 constant REMOVE_SIGNERS_TYPEHASH = + keccak256("RemoveSigners(address[] signersToRemove,uint256 newThreshold,uint256 nonce,uint256 deadline)"); + +uint256 constant OPERATION_EXPIRY = 1 hours; contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, MultiSignerERC7913 { using SafeERC20 for IERC20; + using {Utils.hashArrayed, Utils.toAddressBytesArray} for address[]; + using {Utils.toBytes} for address; + using {Utils.toAddress} for bytes; + error EmptySignersArray(); + error DeadlineExpired(); + error InvalidSignature(); + error InvalidThreshold(); error NotSigner(); - error InvalidQuorum(); error InvalidUser(); - error WithdrawalExpired(); error SignerAlreadyApproved(); error WithdrawalNotExpired(); - error InvalidSignature(); - error EmptySignersArray(); - error DeadlineExpired(); event WithdrawalApproved(bytes32 indexed withdrawalId, address indexed signer, uint256 currentApprovals); @@ -31,26 +43,18 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi uint256 amount; bool finalized; uint64 createdAt; - uint64 requiredQuorum; + uint64 requiredThreshold; } - bytes32 public constant ADD_SIGNERS_TYPEHASH = - keccak256("AddSigners(address[] newSigners,uint256 newQuorum,uint256 nonce,uint256 deadline)"); - bytes32 public constant REMOVE_SIGNERS_TYPEHASH = - keccak256("RemoveSigners(address[] signersToRemove,uint256 newQuorum,uint256 nonce,uint256 deadline)"); - - uint256 public constant OPERATION_EXPIRY = 1 hours; - mapping(bytes32 withdrawalId => WithdrawalRequest request) public withdrawals; mapping(bytes32 withdrawalId => mapping(address signer => bool hasApproved)) public withdrawalApprovals; uint256 public signerNonce; - constructor(address[] memory initialSigners, uint64 quorum_) + constructor(address[] memory initialSigners, uint64 threshold) EIP712("ThresholdCustody", "1") - MultiSignerERC7913(_toAddressBytesArray(initialSigners), quorum_) + MultiSignerERC7913(initialSigners.toAddressBytesArray(), threshold) { require(initialSigners.length != 0, EmptySignersArray()); - require(quorum_ != 0 && quorum_ <= initialSigners.length, InvalidQuorum()); } modifier onlySigner() { @@ -59,26 +63,37 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi } function isSigner(address signer) public view returns (bool) { - return isSigner(_toBytes(signer)); + return isSigner(signer.toBytes()); + } + + function setThreshold(uint64 newThreshold, uint256 deadline, bytes calldata signatures) external { + require(block.timestamp <= deadline, DeadlineExpired()); + + bytes32 structHash = keccak256(abi.encode(SET_THRESHOLD_TYPEHASH, newThreshold, signerNonce, deadline)); + bytes32 digest = _hashTypedDataV4(structHash); + + require(_rawSignatureValidation(digest, signatures), InvalidSignature()); + + signerNonce++; + + _setThreshold(newThreshold); } function addSigners(address[] calldata newSigners, uint64 newThreshold, uint256 deadline, bytes calldata signatures) external - onlySigner { require(block.timestamp <= deadline, DeadlineExpired()); require(newSigners.length != 0, EmptySignersArray()); - bytes32 structHash = keccak256( - abi.encode(ADD_SIGNERS_TYPEHASH, _hashAddressArray(newSigners), newThreshold, signerNonce, deadline) - ); + bytes32 structHash = + keccak256(abi.encode(ADD_SIGNERS_TYPEHASH, newSigners.hashArrayed(), newThreshold, signerNonce, deadline)); bytes32 digest = _hashTypedDataV4(structHash); require(_rawSignatureValidation(digest, signatures), InvalidSignature()); signerNonce++; - _addSigners(_toAddressBytesArray(newSigners)); + _addSigners(newSigners.toAddressBytesArray()); _setThreshold(newThreshold); } @@ -87,12 +102,12 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi uint64 newThreshold, uint256 deadline, bytes calldata signatures - ) external onlySigner { + ) external { require(block.timestamp <= deadline, DeadlineExpired()); require(signersToRemove.length != 0, EmptySignersArray()); bytes32 structHash = keccak256( - abi.encode(REMOVE_SIGNERS_TYPEHASH, _hashAddressArray(signersToRemove), newThreshold, signerNonce, deadline) + abi.encode(REMOVE_SIGNERS_TYPEHASH, signersToRemove.hashArrayed(), newThreshold, signerNonce, deadline) ); bytes32 digest = _hashTypedDataV4(structHash); @@ -100,17 +115,17 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi signerNonce++; - _removeSigners(_toAddressBytesArray(signersToRemove)); + _removeSigners(signersToRemove.toAddressBytesArray()); _setThreshold(newThreshold); } function deposit(address token, uint256 amount) external payable override nonReentrant { - require(amount != 0, ZeroAmount()); + require(amount != 0, IDeposit.ZeroAmount()); if (token == address(0)) { - require(msg.value == amount, MsgValueMismatch()); + require(msg.value == amount, IDeposit.InvalidMsgValue()); } else { - require(msg.value == 0, NonZeroMsgValueForERC20()); + require(msg.value == 0, IDeposit.InvalidMsgValue()); IERC20(token).safeTransferFrom(msg.sender, address(this), amount); } @@ -125,17 +140,17 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi returns (bytes32) { require(user != address(0), InvalidUser()); - require(amount != 0, ZeroAmount()); + require(amount != 0, IDeposit.ZeroAmount()); - bytes32 withdrawalId = _getWithdrawalId(user, token, amount, nonce); - require(withdrawals[withdrawalId].createdAt == 0, WithdrawalAlreadyExists()); + bytes32 withdrawalId = Utils.getWithdrawalId(address(this), user, token, amount, nonce); + require(withdrawals[withdrawalId].createdAt == 0, IWithdraw.WithdrawalAlreadyExists()); withdrawals[withdrawalId] = WithdrawalRequest({ user: user, token: token, amount: amount, finalized: false, - requiredQuorum: threshold(), + requiredThreshold: threshold(), createdAt: uint64(block.timestamp) }); @@ -147,9 +162,9 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi WithdrawalRequest storage request = withdrawals[withdrawalId]; address signer = msg.sender; - require(request.createdAt != 0, WithdrawalNotFound()); - require(!request.finalized, WithdrawalAlreadyFinalized()); - require(block.timestamp <= request.createdAt + OPERATION_EXPIRY, WithdrawalExpired()); + require(request.createdAt != 0, IWithdraw.WithdrawalNotFound()); + require(!request.finalized, IWithdraw.WithdrawalAlreadyFinalized()); + require(block.timestamp <= request.createdAt + OPERATION_EXPIRY, DeadlineExpired()); require(!withdrawalApprovals[withdrawalId][signer], SignerAlreadyApproved()); withdrawalApprovals[withdrawalId][signer] = true; @@ -157,7 +172,7 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi emit WithdrawalApproved(withdrawalId, signer, validApprovals); - if (validApprovals >= request.requiredQuorum) { + if (validApprovals >= request.requiredThreshold) { _executeWithdrawal(request); emit WithdrawFinalized(withdrawalId, true); } @@ -166,8 +181,8 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi function rejectWithdraw(bytes32 withdrawalId) external override onlySigner nonReentrant { WithdrawalRequest storage request = withdrawals[withdrawalId]; - require(request.createdAt != 0, WithdrawalNotFound()); - require(!request.finalized, WithdrawalAlreadyFinalized()); + require(request.createdAt != 0, IWithdraw.WithdrawalNotFound()); + require(!request.finalized, IWithdraw.WithdrawalAlreadyFinalized()); require(block.timestamp > request.createdAt + OPERATION_EXPIRY, WithdrawalNotExpired()); request.finalized = true; @@ -184,11 +199,11 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi request.finalized = true; if (token == address(0)) { - require(address(this).balance >= amount, InsufficientLiquidity()); + require(address(this).balance >= amount, IWithdraw.InsufficientLiquidity()); (bool success,) = user.call{value: amount}(""); - require(success, ETHTransferFailed()); + require(success, IWithdraw.ETHTransferFailed()); } else { - require(IERC20(token).balanceOf(address(this)) >= amount, InsufficientLiquidity()); + require(IERC20(token).balanceOf(address(this)) >= amount, IWithdraw.InsufficientLiquidity()); IERC20(token).safeTransfer(user, amount); } @@ -201,41 +216,8 @@ contract ThresholdCustody is IWithdraw, IDeposit, ReentrancyGuard, EIP712, Multi bytes[] memory allSigners = getSigners(0, type(uint64).max); for (uint256 i = 0; i < allSigners.length; i++) { - address s = _bytesToAddress(allSigners[i]); + address s = allSigners[i].toAddress(); if (withdrawalApprovals[withdrawalId][s]) count++; } } - - function _hashAddressArray(address[] calldata arr) internal pure returns (bytes32) { - bytes32[] memory encoded = new bytes32[](arr.length); - for (uint256 i = 0; i < arr.length; i++) { - encoded[i] = bytes32(uint256(uint160(arr[i]))); - } - return keccak256(abi.encodePacked(encoded)); - } - - function _getWithdrawalId(address user, address token, uint256 amount, uint256 nonce) - internal - view - returns (bytes32) - { - return keccak256(abi.encode(block.chainid, address(this), user, token, amount, nonce)); - } - - // Helpers for conversion - function _toBytes(address a) internal pure returns (bytes memory) { - return abi.encodePacked(a); - } - - function _toAddressBytesArray(address[] memory addrs) internal pure returns (bytes[] memory) { - bytes[] memory b = new bytes[](addrs.length); - for (uint256 i = 0; i < addrs.length; i++) { - b[i] = _toBytes(addrs[i]); - } - return b; - } - - function _bytesToAddress(bytes memory b) internal pure returns (address) { - return address(bytes20(b)); - } } diff --git a/contracts/evm/src/Utils.sol b/contracts/evm/src/Utils.sol new file mode 100644 index 0000000..452a76d --- /dev/null +++ b/contracts/evm/src/Utils.sol @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.30; + +library Utils { + function getWithdrawalId(address custodyAddress, address user, address token, uint256 amount, uint256 nonce) + internal + view + returns (bytes32) + { + return keccak256(abi.encode(block.chainid, custodyAddress, user, token, amount, nonce)); + } + + // Helpers for conversion + function toBytes(address a) internal pure returns (bytes memory) { + return abi.encodePacked(a); + } + + function toAddressBytesArray(address[] memory addrs) internal pure returns (bytes[] memory) { + bytes[] memory b = new bytes[](addrs.length); + for (uint256 i = 0; i < addrs.length; i++) { + b[i] = toBytes(addrs[i]); + } + return b; + } + + function hashArrayed(address[] memory arr) internal pure returns (bytes32) { + bytes32[] memory encoded = new bytes32[](arr.length); + for (uint256 i = 0; i < arr.length; i++) { + encoded[i] = bytes32(uint256(uint160(arr[i]))); + } + return keccak256(abi.encodePacked(encoded)); + } + + function toAddress(bytes memory b) internal pure returns (address) { + return address(bytes20(b)); + } +} diff --git a/contracts/evm/src/interfaces/IDeposit.sol b/contracts/evm/src/interfaces/IDeposit.sol index ddada0e..1d844bb 100644 --- a/contracts/evm/src/interfaces/IDeposit.sol +++ b/contracts/evm/src/interfaces/IDeposit.sol @@ -2,9 +2,8 @@ pragma solidity ^0.8.30; interface IDeposit { - // ---- errors ---- - error MsgValueMismatch(); - error NonZeroMsgValueForERC20(); + error ZeroAmount(); + error InvalidMsgValue(); /// @notice Emitted when a user deposits funds into custody. event Deposited(address indexed user, address indexed token, uint256 amount); diff --git a/contracts/evm/src/interfaces/IWithdraw.sol b/contracts/evm/src/interfaces/IWithdraw.sol index d3a8cea..f64e384 100644 --- a/contracts/evm/src/interfaces/IWithdraw.sol +++ b/contracts/evm/src/interfaces/IWithdraw.sol @@ -2,8 +2,6 @@ pragma solidity ^0.8.30; interface IWithdraw { - // ---- errors ---- - error ZeroAmount(); error WithdrawalAlreadyExists(); error WithdrawalNotFound(); error WithdrawalAlreadyFinalized(); diff --git a/contracts/evm/test/QuorumCustody.t.sol b/contracts/evm/test/QuorumCustody.t.sol index 6e94f93..0aca51a 100644 --- a/contracts/evm/test/QuorumCustody.t.sol +++ b/contracts/evm/test/QuorumCustody.t.sol @@ -654,14 +654,14 @@ contract QuorumCustodyTest is Test { function test_Fail_DepositZeroAmount() public { vm.prank(user); - vm.expectRevert(IWithdraw.ZeroAmount.selector); + vm.expectRevert(IDeposit.ZeroAmount.selector); custody.deposit(address(0), 0); } function test_Fail_DepositETH_MsgValueMismatch() public { vm.deal(user, 2 ether); vm.prank(user); - vm.expectRevert(IDeposit.MsgValueMismatch.selector); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); custody.deposit{value: 0.5 ether}(address(0), 1 ether); } @@ -670,7 +670,7 @@ contract QuorumCustodyTest is Test { vm.deal(user, 1 ether); vm.startPrank(user); token.approve(address(custody), 100e18); - vm.expectRevert(IDeposit.NonZeroMsgValueForERC20.selector); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); custody.deposit{value: 1 ether}(address(token), 100e18); vm.stopPrank(); } @@ -693,7 +693,7 @@ contract QuorumCustodyTest is Test { function test_Fail_StartWithdraw_ZeroAmount() public { vm.prank(signer1); - vm.expectRevert(IWithdraw.ZeroAmount.selector); + vm.expectRevert(IDeposit.ZeroAmount.selector); custody.startWithdraw(user, address(0), 0, 1); } diff --git a/contracts/evm/test/SimpleCustody.t.sol b/contracts/evm/test/SimpleCustody.t.sol index 003eb73..bbd9b7a 100644 --- a/contracts/evm/test/SimpleCustody.t.sol +++ b/contracts/evm/test/SimpleCustody.t.sol @@ -73,7 +73,7 @@ contract SimpleCustodyTest is Test { function test_depositETH_wrongAmount() public { vm.deal(user, 2 ether); vm.prank(user); - vm.expectRevert(IDeposit.MsgValueMismatch.selector); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); custody.deposit{value: 1 ether}(address(0), 2 ether); } @@ -84,7 +84,7 @@ contract SimpleCustodyTest is Test { vm.startPrank(user); token.approve(address(custody), 100e18); - vm.expectRevert(IDeposit.NonZeroMsgValueForERC20.selector); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); custody.deposit{value: 1 ether}(address(token), 100e18); vm.stopPrank(); } @@ -120,13 +120,13 @@ contract SimpleCustodyTest is Test { function test_deposit_zeroAmount() public { vm.prank(user); - vm.expectRevert(IWithdraw.ZeroAmount.selector); + vm.expectRevert(IDeposit.ZeroAmount.selector); custody.deposit{value: 0}(address(0), 0); } function test_startWithdraw_zeroAmount() public { vm.startPrank(neodax); - vm.expectRevert(IWithdraw.ZeroAmount.selector); + vm.expectRevert(IDeposit.ZeroAmount.selector); custody.startWithdraw(user, address(0), 0, 1); vm.stopPrank(); } diff --git a/contracts/evm/test/TestThresholdCustody.sol b/contracts/evm/test/TestThresholdCustody.sol new file mode 100644 index 0000000..7fd33b4 --- /dev/null +++ b/contracts/evm/test/TestThresholdCustody.sol @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.30; + +import {ThresholdCustody} from "../src/ThresholdCustody.sol"; + +/// @notice Harness contract for testing internal functions of ThresholdCustody +contract TestThresholdCustody is ThresholdCustody { + constructor(address[] memory initialSigners, uint64 threshold) ThresholdCustody(initialSigners, threshold) {} + + /// @notice Exposes the internal _executeWithdrawal function for testing + function exposed_executeWithdrawal(bytes32 withdrawalId) external { + WithdrawalRequest storage request = withdrawals[withdrawalId]; + _executeWithdrawal(request); + } + + /// @notice Exposes the internal _countValidApprovals function for testing + function exposed_countValidApprovals(bytes32 withdrawalId) external view returns (uint256) { + return _countValidApprovals(withdrawalId); + } + + /// @notice Helper to set withdrawal approvals directly for testing + function workaround_setWithdrawalApproval(bytes32 withdrawalId, address signer, bool hasApproved) external { + withdrawalApprovals[withdrawalId][signer] = hasApproved; + } +} diff --git a/contracts/evm/test/ThresholdCustody.sol b/contracts/evm/test/ThresholdCustody.sol deleted file mode 100644 index 5c3d0b2..0000000 --- a/contracts/evm/test/ThresholdCustody.sol +++ /dev/null @@ -1,1271 +0,0 @@ -// SPDX-License-Identifier: MIT -pragma solidity ^0.8.24; - -import {Test} from "forge-std/Test.sol"; -import {ThresholdCustody} from "../src/ThresholdCustody.sol"; -import {IWithdraw} from "../src/interfaces/IWithdraw.sol"; -import {IDeposit} from "../src/interfaces/IDeposit.sol"; -import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; - -contract MockERC20 is ERC20 { - constructor() ERC20("Mock", "MCK") {} - - function mint(address to, uint256 amount) external { - _mint(to, amount); - } -} - -contract ThresholdCustodyTest is Test { - ThresholdCustody public custody; - MockERC20 public token; - - address internal user; - - address internal signer1; - uint256 internal signer1Pk; - address internal signer2; - uint256 internal signer2Pk; - address internal signer3; - uint256 internal signer3Pk; - address internal signer4; - uint256 internal signer4Pk; - address internal signer5; - uint256 internal signer5Pk; - - // EIP-712 domain values (must match contract constructor) - bytes32 constant ADD_SIGNERS_TYPEHASH = - keccak256("AddSigners(address[] newSigners,uint256 newQuorum,uint256 nonce,uint256 deadline)"); - bytes32 constant REMOVE_SIGNERS_TYPEHASH = - keccak256("RemoveSigners(address[] signersToRemove,uint256 newQuorum,uint256 nonce,uint256 deadline)"); - uint256 constant MAX_DEADLINE = type(uint256).max; - - function setUp() public { - user = makeAddr("user"); - (signer1, signer1Pk) = makeAddrAndKey("signer1"); - (signer2, signer2Pk) = makeAddrAndKey("signer2"); - (signer3, signer3Pk) = makeAddrAndKey("signer3"); - (signer4, signer4Pk) = makeAddrAndKey("signer4"); - (signer5, signer5Pk) = makeAddrAndKey("signer5"); - - address[] memory initialSigners = new address[](1); - initialSigners[0] = signer1; - custody = new ThresholdCustody(initialSigners, 1); - token = new MockERC20(); - } - - // ========================================================================= - // EIP-712 signing helpers - // ========================================================================= - - function _domainSeparator() internal view returns (bytes32) { - return keccak256( - abi.encode( - keccak256("EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"), - keccak256("ThresholdCustody"), - keccak256("1"), - block.chainid, - address(custody) - ) - ); - } - - function _hashAddressArray(address[] memory arr) internal pure returns (bytes32) { - bytes32[] memory encoded = new bytes32[](arr.length); - for (uint256 i = 0; i < arr.length; i++) { - encoded[i] = bytes32(uint256(uint160(arr[i]))); - } - return keccak256(abi.encodePacked(encoded)); - } - - function _signAddSigners( - uint256 pk, - address[] memory newSigners, - uint256 newQuorum, - uint256 nonce, - uint256 deadline - ) internal view returns (bytes memory) { - bytes32 structHash = keccak256( - abi.encode(ADD_SIGNERS_TYPEHASH, _hashAddressArray(newSigners), newQuorum, nonce, deadline) - ); - bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparator(), structHash)); - (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); - return abi.encodePacked(r, s, v); - } - - function _signRemoveSigners( - uint256 pk, - address[] memory signersToRemove, - uint256 newQuorum, - uint256 nonce, - uint256 deadline - ) internal view returns (bytes memory) { - bytes32 structHash = keccak256( - abi.encode(REMOVE_SIGNERS_TYPEHASH, _hashAddressArray(signersToRemove), newQuorum, nonce, deadline) - ); - bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparator(), structHash)); - (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); - return abi.encodePacked(r, s, v); - } - - // Helper: deploy 5-signer custody with quorum=3 - function _setup3of5() internal { - address[] memory allSigners = new address[](5); - allSigners[0] = signer1; - allSigners[1] = signer2; - allSigners[2] = signer3; - allSigners[3] = signer4; - allSigners[4] = signer5; - custody = new ThresholdCustody(allSigners, 3); - } - - // Helper: sort two signatures by signer address (ascending) and encode in MultiSignerERC7913 format - function _sortSigs2(address a, bytes memory sigA, address b, bytes memory sigB) - internal - pure - returns (bytes memory) - { - return _encodeMultiSig2(a, sigA, b, sigB); - } - - function _emptySigs() internal pure returns (bytes memory) { - // Return properly encoded empty arrays for MultiSignerERC7913 format - bytes[] memory emptySigners = new bytes[](0); - bytes[] memory emptySignatures = new bytes[](0); - return abi.encode(emptySigners, emptySignatures); - } - - // Helper for self-signature: when a single signer with threshold=1 signs their own operation - function _selfSign( - uint256 pk, - address signer, - address[] memory newSigners, - uint256 newThreshold, - uint256 nonce, - uint256 deadline - ) internal view returns (bytes memory) { - bytes memory sig = _signAddSigners(pk, newSigners, newThreshold, nonce, deadline); - return _encodeMultiSig(signer, sig); - } - - function _selfSignRemove( - uint256 pk, - address signer, - address[] memory signersToRemove, - uint256 newThreshold, - uint256 nonce, - uint256 deadline - ) internal view returns (bytes memory) { - bytes memory sig = _signRemoveSigners(pk, signersToRemove, newThreshold, nonce, deadline); - return _encodeMultiSig(signer, sig); - } - - // Helper to encode a single signature in MultiSignerERC7913 format - function _encodeMultiSig(address signer, bytes memory signature) internal pure returns (bytes memory) { - bytes[] memory signers = new bytes[](1); - signers[0] = abi.encodePacked(signer); - bytes[] memory signatures = new bytes[](1); - signatures[0] = signature; - return abi.encode(signers, signatures); - } - - // Helper to encode two signatures in MultiSignerERC7913 format (sorted by signer) - function _encodeMultiSig2(address signerA, bytes memory sigA, address signerB, bytes memory sigB) - internal - pure - returns (bytes memory) - { - bytes[] memory signers = new bytes[](2); - bytes[] memory signatures = new bytes[](2); - - if (uint160(signerA) < uint160(signerB)) { - signers[0] = abi.encodePacked(signerA); - signers[1] = abi.encodePacked(signerB); - signatures[0] = sigA; - signatures[1] = sigB; - } else { - signers[0] = abi.encodePacked(signerB); - signers[1] = abi.encodePacked(signerA); - signatures[0] = sigB; - signatures[1] = sigA; - } - - return abi.encode(signers, signatures); - } - - // ========================================================================= - // Constructor tests - // ========================================================================= - - function test_InitialState() public view { - assertEq(custody.threshold(), 1); - bytes[] memory signers = custody.getSigners(0, type(uint64).max); - assertEq(signers.length, 1); - assertTrue(custody.isSigner(signer1)); - assertEq(custody.getSignerCount(), 1); - } - - function test_Constructor_MultipleSigners() public { - address[] memory s = new address[](3); - s[0] = signer1; - s[1] = signer2; - s[2] = signer3; - ThresholdCustody c = new ThresholdCustody(s, 2); - - assertEq(c.threshold(), 2); - assertEq(c.getSignerCount(), 3); - assertTrue(c.isSigner(signer1)); - assertTrue(c.isSigner(signer2)); - assertTrue(c.isSigner(signer3)); - } - - function test_Fail_Constructor_EmptySigners() public { - address[] memory s = new address[](0); - // With 0 signers, MultiSignerERC7913 will try to validate threshold against 0 signers - vm.expectRevert(); - new ThresholdCustody(s, 1); - } - - // NOTE: address(0) is actually valid for MultiSignerERC7913 as it's 20 bytes - // This test is removed as it's not a failure case in the new implementation - - function test_Fail_Constructor_DuplicateSigner() public { - address[] memory s = new address[](2); - s[0] = signer1; - s[1] = signer1; - // MultiSignerERC7913 will revert with AlreadyExists error - vm.expectRevert(); - new ThresholdCustody(s, 1); - } - - function test_Fail_Constructor_QuorumZero() public { - address[] memory s = new address[](1); - s[0] = signer1; - // MultiSignerERC7913 will revert with ZeroThreshold error - vm.expectRevert(); - new ThresholdCustody(s, 0); - } - - function test_Fail_Constructor_QuorumTooHigh() public { - address[] memory s = new address[](1); - s[0] = signer1; - // MultiSignerERC7913 will revert with UnreachableThreshold error - vm.expectRevert(); - new ThresholdCustody(s, 2); - } - - // ========================================================================= - // addSigners - // ========================================================================= - - function test_AddSigners_Quorum1_EmptySigs() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - // With threshold=1, signer1 must provide their own signature - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 2, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - custody.addSigners(newSigners, 2, MAX_DEADLINE, sigs); - - assertTrue(custody.isSigner(signer2)); - assertEq(custody.threshold(), 2); - assertEq(custody.getSignerCount(), 2); - } - - function test_AddSigners_WithSignature() public { - // First add signer2 to get threshold=2 - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory sigs1 = _selfSign(signer1Pk, signer1, s1, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 2, MAX_DEADLINE, sigs1); - - // Now add signer3 with threshold=2, need 2 signatures total - address[] memory s2 = new address[](1); - s2[0] = signer3; - uint256 nonce = custody.signerNonce(); - - bytes memory sig1 = _signAddSigners(signer1Pk, s2, 2, nonce, MAX_DEADLINE); - bytes memory sig2 = _signAddSigners(signer2Pk, s2, 2, nonce, MAX_DEADLINE); - bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); - - vm.prank(signer1); - custody.addSigners(s2, 2, MAX_DEADLINE, encodedSigs); - - assertTrue(custody.isSigner(signer3)); - assertEq(custody.getSignerCount(), 3); - } - - function test_AddSigners_BatchMultiple() public { - address[] memory newSigners = new address[](3); - newSigners[0] = signer2; - newSigners[1] = signer3; - newSigners[2] = signer4; - - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 2, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - custody.addSigners(newSigners, 2, MAX_DEADLINE, sigs); - - assertTrue(custody.isSigner(signer2)); - assertTrue(custody.isSigner(signer3)); - assertTrue(custody.isSigner(signer4)); - assertEq(custody.getSignerCount(), 4); - assertEq(custody.threshold(), 2); - } - - function test_AddSigners_ThresholdChanged() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 1, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - custody.addSigners(newSigners, 1, MAX_DEADLINE, sigs); - - assertEq(custody.threshold(), 1); - } - - function test_Fail_AddSigners_NotSigner() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - vm.prank(user); - vm.expectRevert(ThresholdCustody.NotSigner.selector); - custody.addSigners(newSigners, 1, MAX_DEADLINE, _emptySigs()); - } - - function test_Fail_AddSigners_EmptyArray() public { - address[] memory newSigners = new address[](0); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.EmptySignersArray.selector); - custody.addSigners(newSigners, 1, MAX_DEADLINE, _emptySigs()); - } - - // Note: ZeroAddress, Duplicate, and DuplicateInBatch validation is now handled by MultiSignerERC7913 - // These will revert with MultiSignerERC7913 errors instead - - function test_Fail_AddSigners_QuorumZero() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 0, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - // MultiSignerERC7913 will revert with ZeroThreshold error - vm.expectRevert(); - custody.addSigners(newSigners, 0, MAX_DEADLINE, sigs); - } - - function test_Fail_AddSigners_QuorumTooHigh() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 3, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - // MultiSignerERC7913 will revert with UnreachableThreshold error - vm.expectRevert(); - custody.addSigners(newSigners, 3, MAX_DEADLINE, sigs); // max is 2 (1 existing + 1 new) - } - - function test_Fail_AddSigners_InsufficientSignatures() public { - // Setup: 2 signers, threshold=2 - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory sigs1 = _selfSign(signer1Pk, signer1, s1, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 2, MAX_DEADLINE, sigs1); - - // Try to add signer3 with only 1 signature (need 2) - address[] memory s2 = new address[](1); - s2[0] = signer3; - uint256 nonce = custody.signerNonce(); - bytes memory sig1 = _signAddSigners(signer1Pk, s2, 2, nonce, MAX_DEADLINE); - bytes memory onlyOneSig = _encodeMultiSig(signer1, sig1); - - vm.prank(signer1); - // Returns InvalidSignature for insufficient signatures - vm.expectRevert(ThresholdCustody.InvalidSignature.selector); - custody.addSigners(s2, 2, MAX_DEADLINE, onlyOneSig); - } - - function test_Fail_AddSigners_StaleNonce() public { - // Setup: 2 signers, threshold=2 - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory sigs1 = _selfSign(signer1Pk, signer1, s1, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 2, MAX_DEADLINE, sigs1); - - // Pre-sign at current nonce - address[] memory s2 = new address[](1); - s2[0] = signer3; - uint256 staleNonce = custody.signerNonce(); - bytes memory staleSig1 = _signAddSigners(signer1Pk, s2, 2, staleNonce, MAX_DEADLINE); - bytes memory staleSig2 = _signAddSigners(signer2Pk, s2, 2, staleNonce, MAX_DEADLINE); - - // Add signer4 first (advances nonce) - address[] memory s3 = new address[](1); - s3[0] = signer4; - bytes memory sig1 = _signAddSigners(signer1Pk, s3, 2, staleNonce, MAX_DEADLINE); - bytes memory sig2 = _signAddSigners(signer2Pk, s3, 2, staleNonce, MAX_DEADLINE); - bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); - vm.prank(signer1); - custody.addSigners(s3, 2, MAX_DEADLINE, encodedSigs); - - // Now try to use the stale signatures (nonce is now incremented) - bytes memory staleEncodedSigs = _encodeMultiSig2(signer1, staleSig1, signer2, staleSig2); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.InvalidSignature.selector); - custody.addSigners(s2, 2, MAX_DEADLINE, staleEncodedSigs); - } - - function test_AddSigners_IncrementsNonce() public { - uint256 nonceBefore = custody.signerNonce(); - - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 1, nonceBefore, MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(newSigners, 1, MAX_DEADLINE, sigs); - - assertEq(custody.signerNonce(), nonceBefore + 1); - } - - // ========================================================================= - // removeSigners - // ========================================================================= - - function test_RemoveSigners_Quorum1() public { - // Add signer2 first - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s1, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 1, MAX_DEADLINE, addSigs); - - // Remove signer2 - address[] memory toRemove = new address[](1); - toRemove[0] = signer2; - bytes memory removeSigs = _selfSignRemove(signer1Pk, signer1, toRemove, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.removeSigners(toRemove, 1, MAX_DEADLINE, removeSigs); - - assertFalse(custody.isSigner(signer2)); - assertEq(custody.getSignerCount(), 1); - } - - function test_RemoveSigners_WithSignature() public { - // Setup 3 signers, threshold 2 - address[] memory s = new address[](2); - s[0] = signer2; - s[1] = signer3; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, addSigs); - - // Remove signer3, need 2 signatures total - address[] memory toRemove = new address[](1); - toRemove[0] = signer3; - uint256 nonce = custody.signerNonce(); - bytes memory sig1 = _signRemoveSigners(signer1Pk, toRemove, 2, nonce, MAX_DEADLINE); - bytes memory sig2 = _signRemoveSigners(signer2Pk, toRemove, 2, nonce, MAX_DEADLINE); - bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); - - vm.prank(signer1); - custody.removeSigners(toRemove, 2, MAX_DEADLINE, encodedSigs); - - assertFalse(custody.isSigner(signer3)); - assertEq(custody.getSignerCount(), 2); - } - - function test_RemoveSigners_BatchMultiple() public { - _setup3of5(); - - // Remove signer4 and signer5 at once - need 3 signatures for threshold=3 - address[] memory toRemove = new address[](2); - toRemove[0] = signer4; - toRemove[1] = signer5; - uint256 nonce = custody.signerNonce(); - - bytes memory sig1 = _signRemoveSigners(signer1Pk, toRemove, 3, nonce, MAX_DEADLINE); - bytes memory sig2 = _signRemoveSigners(signer2Pk, toRemove, 3, nonce, MAX_DEADLINE); - bytes memory sig3 = _signRemoveSigners(signer3Pk, toRemove, 3, nonce, MAX_DEADLINE); - - // Encode all 3 signatures - bytes[] memory signers = new bytes[](3); - bytes[] memory signatures = new bytes[](3); - address[3] memory addrs = [signer1, signer2, signer3]; - bytes[3] memory sigs = [sig1, sig2, sig3]; - // Sort by address - for (uint256 i = 0; i < 2; i++) { - for (uint256 j = i + 1; j < 3; j++) { - if (uint160(addrs[i]) > uint160(addrs[j])) { - (addrs[i], addrs[j]) = (addrs[j], addrs[i]); - (sigs[i], sigs[j]) = (sigs[j], sigs[i]); - } - } - } - for (uint256 i = 0; i < 3; i++) { - signers[i] = abi.encodePacked(addrs[i]); - signatures[i] = sigs[i]; - } - bytes memory encodedSigs = abi.encode(signers, signatures); - - vm.prank(signer1); - custody.removeSigners(toRemove, 3, MAX_DEADLINE, encodedSigs); - - assertFalse(custody.isSigner(signer4)); - assertFalse(custody.isSigner(signer5)); - assertEq(custody.getSignerCount(), 3); - assertEq(custody.threshold(), 3); - } - - function test_Fail_RemoveSigners_UnreachableThreshold() public { - address[] memory toRemove = new address[](1); - toRemove[0] = signer1; - - vm.prank(signer1); - // MultiSignerERC7913 will revert with UnreachableThreshold error - // when removing would make the threshold impossible to reach - vm.expectRevert(); - custody.removeSigners(toRemove, 1, MAX_DEADLINE, _emptySigs()); - } - - function test_Fail_RemoveSigners_InvalidQuorum() public { - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s1, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 1, MAX_DEADLINE, addSigs); - - address[] memory toRemove = new address[](1); - toRemove[0] = signer2; - bytes memory removeSigs = _selfSignRemove(signer1Pk, signer1, toRemove, 2, custody.signerNonce(), MAX_DEADLINE); - - vm.prank(signer1); - // MultiSignerERC7913 will revert with UnreachableThreshold error - vm.expectRevert(); - custody.removeSigners(toRemove, 2, MAX_DEADLINE, removeSigs); // removing leaves 1, max threshold is 1 - } - - function test_Fail_RemoveSigners_EmptyArray() public { - address[] memory toRemove = new address[](0); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.EmptySignersArray.selector); - custody.removeSigners(toRemove, 1, MAX_DEADLINE, _emptySigs()); - } - - function test_RemovedSignerCannotAct() public { - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s1, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 1, MAX_DEADLINE, addSigs); - - address[] memory toRemove = new address[](1); - toRemove[0] = signer2; - bytes memory removeSigs = _selfSignRemove(signer1Pk, signer1, toRemove, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.removeSigners(toRemove, 1, MAX_DEADLINE, removeSigs); - - vm.prank(signer2); - vm.expectRevert(ThresholdCustody.NotSigner.selector); - custody.startWithdraw(user, address(0), 1 ether, 1); - } - - // ========================================================================= - // Deposit - // ========================================================================= - - function test_DepositETH() public { - vm.deal(user, 1 ether); - vm.prank(user); - custody.deposit{value: 1 ether}(address(0), 1 ether); - assertEq(address(custody).balance, 1 ether); - } - - function test_DepositERC20() public { - token.mint(user, 100e18); - vm.startPrank(user); - token.approve(address(custody), 100e18); - custody.deposit(address(token), 100e18); - vm.stopPrank(); - assertEq(token.balanceOf(address(custody)), 100e18); - } - - function test_DepositETH_EmitsEvent() public { - vm.deal(user, 1 ether); - vm.prank(user); - vm.expectEmit(true, true, false, true); - emit IDeposit.Deposited(user, address(0), 1 ether); - custody.deposit{value: 1 ether}(address(0), 1 ether); - } - - function test_DepositERC20_EmitsEvent() public { - token.mint(user, 50e18); - vm.startPrank(user); - token.approve(address(custody), 50e18); - vm.expectEmit(true, true, false, true); - emit IDeposit.Deposited(user, address(token), 50e18); - custody.deposit(address(token), 50e18); - vm.stopPrank(); - } - - function test_Fail_DepositZeroAmount() public { - vm.prank(user); - vm.expectRevert(IWithdraw.ZeroAmount.selector); - custody.deposit(address(0), 0); - } - - function test_Fail_DepositETH_MsgValueMismatch() public { - vm.deal(user, 2 ether); - vm.prank(user); - vm.expectRevert(IDeposit.MsgValueMismatch.selector); - custody.deposit{value: 0.5 ether}(address(0), 1 ether); - } - - function test_Fail_DepositERC20_NonZeroMsgValue() public { - token.mint(user, 100e18); - vm.deal(user, 1 ether); - vm.startPrank(user); - token.approve(address(custody), 100e18); - vm.expectRevert(IDeposit.NonZeroMsgValueForERC20.selector); - custody.deposit{value: 1 ether}(address(token), 100e18); - vm.stopPrank(); - } - - // ========================================================================= - // startWithdraw - // ========================================================================= - - function test_Fail_StartWithdraw_NotSigner() public { - vm.prank(user); - vm.expectRevert(ThresholdCustody.NotSigner.selector); - custody.startWithdraw(user, address(0), 1 ether, 1); - } - - function test_Fail_StartWithdraw_ZeroUser() public { - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.InvalidUser.selector); - custody.startWithdraw(address(0), address(0), 1 ether, 1); - } - - function test_Fail_StartWithdraw_ZeroAmount() public { - vm.prank(signer1); - vm.expectRevert(IWithdraw.ZeroAmount.selector); - custody.startWithdraw(user, address(0), 0, 1); - } - - function test_Fail_StartWithdraw_DuplicateNonce() public { - vm.startPrank(signer1); - custody.startWithdraw(user, address(0), 1 ether, 1); - vm.expectRevert(IWithdraw.WithdrawalAlreadyExists.selector); - custody.startWithdraw(user, address(0), 1 ether, 1); - vm.stopPrank(); - } - - function test_StartWithdraw_SameParamsDifferentNonce() public { - vm.startPrank(signer1); - bytes32 id1 = custody.startWithdraw(user, address(0), 1 ether, 1); - bytes32 id2 = custody.startWithdraw(user, address(0), 1 ether, 2); - vm.stopPrank(); - assertTrue(id1 != id2); - } - - function test_StartWithdraw_EmitsEvent() public { - vm.prank(signer1); - vm.expectEmit(true, true, true, true); - bytes32 expectedId = - keccak256(abi.encode(block.chainid, address(custody), user, address(0), 1 ether, uint256(1))); - emit IWithdraw.WithdrawStarted(expectedId, user, address(0), 1 ether, 1); - custody.startWithdraw(user, address(0), 1 ether, 1); - } - - function test_StartWithdraw_SnapshotsQuorum() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - (,,,,, uint64 requiredQuorum) = custody.withdrawals(id); - assertEq(requiredQuorum, 1); - } - - // ========================================================================= - // finalizeWithdraw — 1/1 - // ========================================================================= - - function test_FinalizeWithdraw_1_1() public { - vm.deal(address(custody), 1 ether); - - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // finalizeWithdraw — 2/2 progressive - // ========================================================================= - - function test_FinalizeWithdraw_2_2_Progressive() public { - // Setup: 2 signers, threshold=2 - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory sigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, sigs); - - vm.deal(address(custody), 1 ether); - - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - (,,, bool finalized,,) = custody.withdrawals(id); - assertFalse(finalized); - - vm.prank(signer2); - custody.finalizeWithdraw(id); - (,,, finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // finalizeWithdraw — 3/5 - // ========================================================================= - - function test_FinalizeWithdraw_3_5() public { - _setup3of5(); - assertEq(custody.threshold(), 3); - assertEq(custody.getSignerCount(), 5); - - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - (,,, bool finalized,,) = custody.withdrawals(id); - assertFalse(finalized); - - vm.prank(signer2); - custody.finalizeWithdraw(id); - (,,, finalized,,) = custody.withdrawals(id); - assertFalse(finalized); - - vm.prank(signer3); - custody.finalizeWithdraw(id); - (,,, finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // finalizeWithdraw — snapshot quorum - // ========================================================================= - - function test_FinalizeWithdraw_UsesSnapshotQuorum() public { - // Setup: 2 signers, threshold=1 - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory addSigs1 = _selfSign(signer1Pk, signer1, s, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 1, MAX_DEADLINE, addSigs1); - - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - // Raise threshold to 2 AFTER withdrawal was created - address[] memory s2 = new address[](1); - s2[0] = signer3; - uint256 nonce = custody.signerNonce(); - bytes memory sig1 = _signAddSigners(signer1Pk, s2, 2, nonce, MAX_DEADLINE); - bytes memory encodedSigs = _encodeMultiSig(signer1, sig1); - - vm.prank(signer1); - custody.addSigners(s2, 2, MAX_DEADLINE, encodedSigs); - assertEq(custody.threshold(), 2); - - // 1 approval should suffice (snapshot quorum was 1) - vm.prank(signer1); - custody.finalizeWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // finalizeWithdraw — edge cases - // ========================================================================= - - function test_Fail_FinalizeWithdraw_NotSigner() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(user); - vm.expectRevert(ThresholdCustody.NotSigner.selector); - custody.finalizeWithdraw(id); - } - - function test_Fail_FinalizeWithdraw_NonExistent() public { - vm.prank(signer1); - vm.expectRevert(IWithdraw.WithdrawalNotFound.selector); - custody.finalizeWithdraw(bytes32(uint256(999))); - } - - function test_Fail_FinalizeWithdraw_AlreadyFinalized() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - vm.prank(signer1); - vm.expectRevert(IWithdraw.WithdrawalAlreadyFinalized.selector); - custody.finalizeWithdraw(id); - } - - function test_Fail_DuplicateApproval() public { - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory sigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, sigs); - - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.SignerAlreadyApproved.selector); - custody.finalizeWithdraw(id); - } - - function test_Fail_Finalize_Expired() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.WithdrawalExpired.selector); - custody.finalizeWithdraw(id); - } - - function test_FinalizeWithdraw_ExactExpiryBoundary() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - function test_FinalizeWithdraw_ERC20() public { - token.mint(address(custody), 50e18); - - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(token), 50e18, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - assertEq(token.balanceOf(user), 50e18); - assertEq(token.balanceOf(address(custody)), 0); - } - - function test_Fail_FinalizeWithdraw_InsufficientETH() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); - custody.finalizeWithdraw(id); - } - - function test_Fail_FinalizeWithdraw_InsufficientERC20() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(token), 50e18, 1); - - vm.prank(signer1); - vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); - custody.finalizeWithdraw(id); - } - - function test_FinalizeWithdraw_ClearsStorage() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - (address storedUser, address storedToken, uint256 storedAmount, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - assertEq(storedUser, address(0)); - assertEq(storedToken, address(0)); - assertEq(storedAmount, 0); - } - - function test_FinalizeWithdraw_EmitsApprovalEvent() public { - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory sigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, sigs); - - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - vm.expectEmit(true, true, false, true); - emit ThresholdCustody.WithdrawalApproved(id, signer1, 1); - custody.finalizeWithdraw(id); - } - - function test_FinalizeWithdraw_EmitsFinalizedEvent() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - vm.expectEmit(true, false, false, true); - emit IWithdraw.WithdrawFinalized(id, true); - custody.finalizeWithdraw(id); - } - - function test_FinalizeWithdraw_ETH_UserReceivesBalance() public { - vm.deal(address(custody), 5 ether); - uint256 balanceBefore = user.balance; - - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 2 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - assertEq(user.balance, balanceBefore + 2 ether); - assertEq(address(custody).balance, 3 ether); - } - - // ========================================================================= - // rejectWithdraw (expired-only cleanup) - // ========================================================================= - - function test_RejectWithdraw_Expired() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer1); - custody.rejectWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - function test_RejectWithdraw_EmitsEvent() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer1); - vm.expectEmit(true, false, false, true); - emit IWithdraw.WithdrawFinalized(id, false); - custody.rejectWithdraw(id); - } - - function test_Fail_RejectWithdraw_NotExpired() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.WithdrawalNotExpired.selector); - custody.rejectWithdraw(id); - } - - function test_Fail_RejectWithdraw_NonExistent() public { - vm.prank(signer1); - vm.expectRevert(IWithdraw.WithdrawalNotFound.selector); - custody.rejectWithdraw(bytes32(uint256(999))); - } - - function test_Fail_RejectWithdraw_AlreadyFinalized() public { - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer1); - vm.expectRevert(IWithdraw.WithdrawalAlreadyFinalized.selector); - custody.rejectWithdraw(id); - } - - function test_Fail_RejectWithdraw_NotSigner() public { - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(user); - vm.expectRevert(ThresholdCustody.NotSigner.selector); - custody.rejectWithdraw(id); - } - - // ========================================================================= - // Lifecycle: reject expired, then re-create - // ========================================================================= - - function test_RejectExpiredThenRecreate() public { - vm.deal(address(custody), 1 ether); - - vm.prank(signer1); - bytes32 id1 = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer1); - custody.rejectWithdraw(id1); - - vm.prank(signer1); - bytes32 id2 = custody.startWithdraw(user, address(0), 1 ether, 2); - assertTrue(id1 != id2); - - vm.prank(signer1); - custody.finalizeWithdraw(id2); - - assertEq(user.balance, 1 ether); - } - - // ========================================================================= - // Partial approval then expiry - // ========================================================================= - - function test_PartialApprovalThenExpiry() public { - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory sigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, sigs); - - vm.deal(address(custody), 1 ether); - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - - vm.prank(signer1); - custody.finalizeWithdraw(id); - - vm.warp(block.timestamp + 1 hours + 1); - - vm.prank(signer2); - vm.expectRevert(ThresholdCustody.WithdrawalExpired.selector); - custody.finalizeWithdraw(id); - - // Clean up expired - vm.prank(signer1); - custody.rejectWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // Multiple concurrent withdrawals - // ========================================================================= - - function test_MultipleConcurrentWithdrawals() public { - vm.deal(address(custody), 3 ether); - - vm.startPrank(signer1); - bytes32 id1 = custody.startWithdraw(user, address(0), 1 ether, 1); - bytes32 id3 = custody.startWithdraw(user, address(0), 1 ether, 3); - vm.stopPrank(); - - vm.prank(signer1); - custody.finalizeWithdraw(id1); - - // id2 left to expire - vm.prank(signer1); - custody.finalizeWithdraw(id3); - - assertEq(user.balance, 2 ether); - assertEq(address(custody).balance, 1 ether); - } - - // ========================================================================= - // getSignerCount - // ========================================================================= - - function test_GetSignerCount() public { - assertEq(custody.getSignerCount(), 1); - - address[] memory s = new address[](2); - s[0] = signer2; - s[1] = signer3; - bytes memory sigs = _selfSign(signer1Pk, signer1, s, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 1, MAX_DEADLINE, sigs); - assertEq(custody.getSignerCount(), 3); - } - - // ========================================================================= - // Removed signer approvals no longer count (withdrawal) - // ========================================================================= - - function test_FinalizeWithdraw_RemovedSignerApprovalIgnored() public { - // Setup: 3 signers, threshold=2 - address[] memory s = new address[](2); - s[0] = signer2; - s[1] = signer3; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 2, MAX_DEADLINE, addSigs); - - vm.deal(address(custody), 1 ether); - - vm.prank(signer1); - bytes32 id = custody.startWithdraw(user, address(0), 1 ether, 1); - // requiredQuorum snapshotted at 2 - - // signer2 approves - vm.prank(signer2); - custody.finalizeWithdraw(id); - - // Remove signer2 (need 2 sigs since threshold=2) - address[] memory toRemove = new address[](1); - toRemove[0] = signer2; - uint256 nonce = custody.signerNonce(); - bytes memory sig1 = _signRemoveSigners(signer1Pk, toRemove, 2, nonce, MAX_DEADLINE); - bytes memory sig3 = _signRemoveSigners(signer3Pk, toRemove, 2, nonce, MAX_DEADLINE); - bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer3, sig3); - vm.prank(signer1); - custody.removeSigners(toRemove, 2, MAX_DEADLINE, encodedSigs); - - assertFalse(custody.isSigner(signer2)); - - // signer1 approves — only 1 valid approval (signer2's no longer counts) - // snapshotted requiredQuorum is still 2, so not finalized yet - vm.prank(signer1); - custody.finalizeWithdraw(id); - - (,,, bool finalized,,) = custody.withdrawals(id); - assertFalse(finalized); - - // signer3 approves — now 2 valid approvals (signer1 + signer3), meets requiredQuorum=2 - vm.prank(signer3); - custody.finalizeWithdraw(id); - - (,,, finalized,,) = custody.withdrawals(id); - assertTrue(finalized); - } - - // ========================================================================= - // Deadline expiry tests - // ========================================================================= - - function test_Fail_AddSigners_DeadlineExpired() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - uint256 deadline = block.timestamp + 1 hours; - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 1, custody.signerNonce(), deadline); - vm.warp(deadline + 1); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); - custody.addSigners(newSigners, 1, deadline, sigs); - } - - function test_Fail_RemoveSigners_DeadlineExpired() public { - // Add signer2 first - address[] memory s = new address[](1); - s[0] = signer2; - bytes memory addSigs = _selfSign(signer1Pk, signer1, s, 1, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s, 1, MAX_DEADLINE, addSigs); - - address[] memory toRemove = new address[](1); - toRemove[0] = signer2; - - uint256 deadline = block.timestamp + 1 hours; - bytes memory removeSigs = _selfSignRemove(signer1Pk, signer1, toRemove, 1, custody.signerNonce(), deadline); - vm.warp(deadline + 1); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); - custody.removeSigners(toRemove, 1, deadline, removeSigs); - } - - function test_AddSigners_ExactDeadlineBoundary() public { - address[] memory newSigners = new address[](1); - newSigners[0] = signer2; - - uint256 deadline = block.timestamp + 1 hours; - bytes memory sigs = _selfSign(signer1Pk, signer1, newSigners, 1, custody.signerNonce(), deadline); - vm.warp(deadline); - - vm.prank(signer1); - custody.addSigners(newSigners, 1, deadline, sigs); - assertTrue(custody.isSigner(signer2)); - } - - // ========================================================================= - // Exploit: Quorum downgrade via addSigners - // ========================================================================= - - // NOTE: ThresholdCustody does not have the same threshold-reduction protection as QuorumCustody - // The MultiSignerERC7913 implementation allows threshold changes as long as they're reachable - // These exploit tests are removed as they test behavior specific to the old QuorumCustody contract - - function test_Fail_AddSigners_DeadlineExpired_WithSignatures() public { - // Setup: 2 signers, threshold=2 - address[] memory s1 = new address[](1); - s1[0] = signer2; - bytes memory sigs1 = _selfSign(signer1Pk, signer1, s1, 2, custody.signerNonce(), MAX_DEADLINE); - vm.prank(signer1); - custody.addSigners(s1, 2, MAX_DEADLINE, sigs1); - - // Sign with a deadline, then let it expire - address[] memory s2 = new address[](1); - s2[0] = signer3; - uint256 nonce = custody.signerNonce(); - uint256 deadline = block.timestamp + 1 hours; - - bytes memory sig1 = _signAddSigners(signer1Pk, s2, 2, nonce, deadline); - bytes memory sig2 = _signAddSigners(signer2Pk, s2, 2, nonce, deadline); - bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); - - vm.warp(deadline + 1); - - vm.prank(signer1); - vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); - custody.addSigners(s2, 2, deadline, encodedSigs); - } -} diff --git a/contracts/evm/test/ThresholdCustody.t.sol b/contracts/evm/test/ThresholdCustody.t.sol new file mode 100644 index 0000000..bf61479 --- /dev/null +++ b/contracts/evm/test/ThresholdCustody.t.sol @@ -0,0 +1,2143 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import {Test} from "forge-std/Test.sol"; + +import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; +import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import {MultiSignerERC7913} from "@openzeppelin/contracts/utils/cryptography/signers/MultiSignerERC7913.sol"; + +import { + ThresholdCustody, + SET_THRESHOLD_TYPEHASH, + ADD_SIGNERS_TYPEHASH, + REMOVE_SIGNERS_TYPEHASH, + OPERATION_EXPIRY +} from "../src/ThresholdCustody.sol"; +import {IWithdraw} from "../src/interfaces/IWithdraw.sol"; +import {IDeposit} from "../src/interfaces/IDeposit.sol"; +import {Utils} from "../src/Utils.sol"; +import {TestThresholdCustody} from "./TestThresholdCustody.sol"; + +using {Utils.toBytes} for address; + +contract MockERC20 is ERC20 { + constructor() ERC20("Mock", "MCK") {} + + function mint(address to, uint256 amount) external { + _mint(to, amount); + } +} + +contract ThresholdCustodyTest_Base is Test { + using {Utils.hashArrayed} for address[]; + + ThresholdCustody public custody; + MockERC20 public token; + + address internal user; + + address internal signer1; + uint256 internal signer1Pk; + address internal signer2; + uint256 internal signer2Pk; + address internal signer3; + uint256 internal signer3Pk; + address internal signer4; + uint256 internal signer4Pk; + address internal signer5; + uint256 internal signer5Pk; + + address internal notSigner; + uint256 internal notSignerPk; + + uint256 constant MAX_DEADLINE = type(uint256).max; + + address[] oneSigner = new address[](1); + address[] twoSigners = new address[](2); + address[] threeSigners = new address[](3); + address[] fiveSigners = new address[](5); + + function setUp() public virtual { + user = makeAddr("user"); + (signer1, signer1Pk) = makeAddrAndKey("signer1"); + (signer2, signer2Pk) = makeAddrAndKey("signer2"); + (signer3, signer3Pk) = makeAddrAndKey("signer3"); + (signer4, signer4Pk) = makeAddrAndKey("signer4"); + (signer5, signer5Pk) = makeAddrAndKey("signer5"); + (notSigner, notSignerPk) = makeAddrAndKey("notSigner"); + + oneSigner[0] = signer1; + twoSigners[0] = signer1; + twoSigners[1] = signer2; + threeSigners[0] = signer1; + threeSigners[1] = signer2; + threeSigners[2] = signer3; + fiveSigners[0] = signer1; + fiveSigners[1] = signer2; + fiveSigners[2] = signer3; + fiveSigners[3] = signer4; + fiveSigners[4] = signer5; + token = new MockERC20(); + } + + // ========================================================================= + // EIP-712 signing helpers + // ========================================================================= + + function _domainSeparator(address contractAddress) internal view returns (bytes32) { + return keccak256( + abi.encode( + keccak256("EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"), + keccak256("ThresholdCustody"), + keccak256("1"), + block.chainid, + address(contractAddress) + ) + ); + } + + function _domainSeparator() internal view returns (bytes32) { + return _domainSeparator(address(custody)); + } + + function _signSetThreshold(uint256 pk, uint256 newThreshold, uint256 nonce, uint256 deadline) + internal + view + returns (bytes memory) + { + bytes32 structHash = keccak256(abi.encode(SET_THRESHOLD_TYPEHASH, newThreshold, nonce, deadline)); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparator(), structHash)); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); + return abi.encodePacked(r, s, v); + } + + function _signAddSigners( + uint256 pk, + address[] memory newSigners, + uint256 newThreshold, + uint256 nonce, + uint256 deadline + ) internal view returns (bytes memory) { + bytes32 structHash = keccak256( + abi.encode(ADD_SIGNERS_TYPEHASH, newSigners.hashArrayed(), newThreshold, nonce, deadline) + ); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparator(), structHash)); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); + return abi.encodePacked(r, s, v); + } + + function _signRemoveSigners( + uint256 pk, + address[] memory signersToRemove, + uint256 newThreshold, + uint256 nonce, + uint256 deadline + ) internal view returns (bytes memory) { + bytes32 structHash = keccak256( + abi.encode(REMOVE_SIGNERS_TYPEHASH, signersToRemove.hashArrayed(), newThreshold, nonce, deadline) + ); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparator(), structHash)); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); + return abi.encodePacked(r, s, v); + } + + // Helper: deploy 5-signer custody with quorum=3 + function _setup3of5() internal { + address[] memory allSigners = new address[](5); + allSigners[0] = signer1; + allSigners[1] = signer2; + allSigners[2] = signer3; + allSigners[3] = signer4; + allSigners[4] = notSigner; + custody = new ThresholdCustody(allSigners, 3); + } + + function _emptySigs() internal pure returns (bytes memory) { + bytes[] memory emptySigners = new bytes[](0); + bytes[] memory emptySignatures = new bytes[](0); + return abi.encode(emptySigners, emptySignatures); + } + + function _signSingleSetThreshold(uint256 pk, uint256 newThreshold, uint256 nonce, uint256 deadline) + internal + view + returns (bytes memory) + { + address signer = vm.addr(pk); + bytes memory sig = _signSetThreshold(pk, newThreshold, nonce, deadline); + return _encodeMultiSig(signer, sig); + } + + function _signSingleAdd( + uint256 pk, + address[] memory newSigners, + uint256 newThreshold, + uint256 nonce, + uint256 deadline + ) internal view returns (bytes memory) { + address signer = vm.addr(pk); + bytes memory sig = _signAddSigners(pk, newSigners, newThreshold, nonce, deadline); + return _encodeMultiSig(signer, sig); + } + + function _signSingleRemove( + uint256 pk, + address[] memory signersToRemove, + uint256 newThreshold, + uint256 nonce, + uint256 deadline + ) internal view returns (bytes memory) { + address signer = vm.addr(pk); + bytes memory sig = _signRemoveSigners(pk, signersToRemove, newThreshold, nonce, deadline); + return _encodeMultiSig(signer, sig); + } + + // Helper to encode a single signature in MultiSignerERC7913 format + function _encodeMultiSig(address signer, bytes memory signature) internal pure returns (bytes memory) { + bytes[] memory signers = new bytes[](1); + signers[0] = abi.encodePacked(signer); + bytes[] memory signatures = new bytes[](1); + signatures[0] = signature; + return abi.encode(signers, signatures); + } + + // Helper to encode two signatures in MultiSignerERC7913 format (sorted by signer) + function _encodeMultiSig2(address signerA, bytes memory sigA, address signerB, bytes memory sigB) + internal + pure + returns (bytes memory) + { + bytes[] memory signers = new bytes[](2); + bytes[] memory signatures = new bytes[](2); + + if (uint160(signerA) < uint160(signerB)) { + signers[0] = abi.encodePacked(signerA); + signers[1] = abi.encodePacked(signerB); + signatures[0] = sigA; + signatures[1] = sigB; + } else { + signers[0] = abi.encodePacked(signerB); + signers[1] = abi.encodePacked(signerA); + signatures[0] = sigB; + signatures[1] = sigA; + } + + return abi.encode(signers, signatures); + } + + // do not sort as MultiSignerERC7913 accepts not sorted as well + function _encodeMultiSig3( + address signerA, + bytes memory sigA, + address signerB, + bytes memory sigB, + address signerC, + bytes memory sigC + ) internal pure returns (bytes memory) { + bytes[] memory signers = new bytes[](3); + bytes[] memory signatures = new bytes[](3); + + signers[0] = abi.encodePacked(signerA); + signers[1] = abi.encodePacked(signerB); + signers[2] = abi.encodePacked(signerC); + + signatures[0] = sigA; + signatures[1] = sigB; + signatures[2] = sigC; + + return abi.encode(signers, signatures); + } + + function _checkStats(uint256 signersCount, uint64 threshold, uint256 newNonce) internal view { + assertEq(custody.threshold(), threshold); + assertEq(custody.getSignerCount(), signersCount); + assertEq(custody.signerNonce(), newNonce); + } + + function _validateWithdrawalData( + bytes32 id, + address expectedUser, + address expectedToken, + uint256 expectedAmount, + bool expectedFinalized, + uint64 expectedThreshold, + uint64 expectedCreatedAt + ) internal view virtual { + ( + address storedUser, + address storedToken, + uint256 storedAmount, + bool storedFinalized, + uint64 storedCreatedAt, + uint64 storedThreshold + ) = custody.withdrawals(id); + assertEq(storedUser, expectedUser); + assertEq(storedToken, expectedToken); + assertEq(storedAmount, expectedAmount); + assertEq(storedFinalized, expectedFinalized); + assertEq(storedThreshold, expectedThreshold); + assertEq(storedCreatedAt, expectedCreatedAt); + } +} + +// ========================================================================= +// Constructor tests +// ========================================================================= +contract ThresholdCustodyTest_Constructor is ThresholdCustodyTest_Base { + function test_singleSigner() public { + address[] memory s = new address[](1); + s[0] = signer1; + ThresholdCustody c = new ThresholdCustody(s, 1); + + assertEq(c.threshold(), 1); + bytes[] memory signers = c.getSigners(0, type(uint64).max); + assertEq(signers.length, 1); + assertEq(Utils.toAddress(signers[0]), signer1); + assertTrue(c.isSigner(signer1)); + assertEq(c.getSignerCount(), 1); + } + + function test_multipleSigners() public { + address[] memory s = new address[](3); + s[0] = signer1; + s[1] = signer2; + s[2] = signer3; + ThresholdCustody c = new ThresholdCustody(s, 2); + + assertEq(c.threshold(), 2); + assertEq(c.getSignerCount(), 3); + assertTrue(c.isSigner(signer1)); + assertTrue(c.isSigner(signer2)); + assertTrue(c.isSigner(signer3)); + } + + function test_revert_emptySigners() public { + address[] memory s = new address[](0); + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 0, 1) + ); + new ThresholdCustody(s, 1); + } + + function test_revert_duplicateSigners() public { + address[] memory s = new address[](2); + s[0] = signer1; + s[1] = signer1; + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913AlreadyExists.selector, signer1.toBytes()) + ); + new ThresholdCustody(s, 1); + } + + function test_revert_quorumZero() public { + address[] memory s = new address[](1); + s[0] = signer1; + vm.expectRevert(MultiSignerERC7913.MultiSignerERC7913ZeroThreshold.selector); + new ThresholdCustody(s, 0); + } + + function test_revert_quorumNotReachable() public { + address[] memory s = new address[](1); + s[0] = signer1; + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 1, 2) + ); + new ThresholdCustody(s, 2); + } +} + +// ========================================================================= +// setThreshold +// ========================================================================= +contract ThresholdCustodyTest_SetThreshold is ThresholdCustodyTest_Base { + function test_success_1_of_3_signature_increase() public { + custody = new ThresholdCustody(threeSigners, 1); + + uint64 newThreshold = 2; // increase + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + + custody.setThreshold(newThreshold, MAX_DEADLINE, sigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_success_2_of_3_signature_decrease() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 1; // decrease + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_2() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signSetThreshold(signer3Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer3, sig3); + + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_2_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signSetThreshold(signer3Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer3, sig3); + + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_revert_emptySignatures() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 1; + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, _emptySigs()); + } + + function test_revert_signatureFromNotSigner() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + // signature from NOT signer + bytes memory notSignerSig = _signSetThreshold(notSignerPk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, notSigner, notSignerSig); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSingleSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, sig2); + } + + function test_revert_duplicatedSignature_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_duplicatedSignature_thresholdReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig3(signer1, sig1, signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_incorrectSignature() public { + custody = new ThresholdCustody(threeSigners, 1); + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + // corrupt the signature by changing one byte + sig1[10] = bytes1(sig1[10] ^ 0x01); + bytes memory encodedSigs = _encodeMultiSig(signer1, sig1); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_zeroNewThreshold() public { + custody = new ThresholdCustody(oneSigner, 1); + + uint64 newThreshold = 0; + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(MultiSignerERC7913.MultiSignerERC7913ZeroThreshold.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_unreachableNewThreshold() public { + custody = new ThresholdCustody(oneSigner, 1); + + uint64 newThreshold = 3; // cannot be reached with 1 signer + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 1, 3) + ); + custody.setThreshold(newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_deadlinePassed() public { + custody = new ThresholdCustody(oneSigner, 1); + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleSetThreshold(signer1Pk, newThreshold, nonce, block.timestamp - 1); + + vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); + custody.setThreshold(newThreshold, block.timestamp - 1, sigs); + } + + function test_revert_outdatedNonce() public { + custody = new ThresholdCustody(threeSigners, 1); + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signSingleSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + + // First call succeeds + custody.setThreshold(newThreshold, MAX_DEADLINE, sig1); + + // Second call with same nonce should revert + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, sig1); + } + + function test_revert_futureNonce() public { + custody = new ThresholdCustody(oneSigner, 1); + + uint64 newThreshold = 2; // change + uint256 futureNonce = custody.signerNonce() + 42; + bytes memory sigs = _signSingleSetThreshold(signer1Pk, newThreshold, futureNonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.setThreshold(newThreshold, MAX_DEADLINE, sigs); + } +} + +// ========================================================================= +// addSigners +// ========================================================================= +contract ThresholdCustodyTest_AddSigners is ThresholdCustodyTest_Base { + function test_success_1_of_1_signature() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + _checkStats(2, newThreshold, ++nonce); + } + + function test_success_onlyAddSigners() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](2); + newSigners[0] = signer2; + newSigners[1] = signer3; + + uint64 newThreshold = 1; // no change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + _checkStats(3, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_2() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signAddSigners(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signAddSigners(signer2Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + assertTrue(custody.isSigner(signer4)); + _checkStats(4, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signAddSigners(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signAddSigners(signer3Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer3, sig3); + + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + assertTrue(custody.isSigner(signer4)); + _checkStats(4, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_2_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signAddSigners(signer2Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signAddSigners(signer3Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer3, sig3); + + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertTrue(custody.isSigner(signer3)); + assertTrue(custody.isSigner(signer4)); + _checkStats(4, newThreshold, ++nonce); + } + + function test_revert_emptySignatures() public { + custody = new ThresholdCustody(oneSigner, 1); + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, _emptySigs()); + } + + function test_revert_signatureFromNotSigner() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signAddSigners(signer2Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + // signature from NOT signer + bytes memory notSignerSig = _signAddSigners(notSignerPk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, notSigner, notSignerSig); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sig2); + } + + function test_revert_duplicatedSignature_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signAddSigners(signer2Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_duplicatedSignature_thresholdReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer4; + + uint64 newThreshold = 3; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signAddSigners(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signAddSigners(signer2Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig3(signer1, sig1, signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_incorrectSignature() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signAddSigners(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + // corrupt the signature by changing one byte + sig1[10] = bytes1(sig1[10] ^ 0x01); + bytes memory encodedSigs = _encodeMultiSig(signer1, sig1); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_emptyNewArray() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](0); + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.EmptySignersArray.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_newArrayIncludesExistingSigner() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer1; // already existing signer + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913AlreadyExists.selector, signer1.toBytes()) + ); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_newArrayIncludesDuplicatedSigner() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](2); + newSigners[0] = signer2; + newSigners[1] = signer2; // duplicated in the new array + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913AlreadyExists.selector, signer2.toBytes()) + ); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_zeroNewThreshold() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 0; + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(MultiSignerERC7913.MultiSignerERC7913ZeroThreshold.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_unreachableNewThreshold() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 3; // cannot be reached with 2 signers + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 2, 3) + ); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_deadlinePassed() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, block.timestamp - 1); + + vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); + custody.addSigners(newSigners, newThreshold, block.timestamp - 1, sigs); + } + + function test_revert_outdatedNonce() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 2; // change + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, nonce, MAX_DEADLINE); + + // First call succeeds + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + + // Second call with same nonce should revert + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_futureNonce() public { + custody = new ThresholdCustody(oneSigner, 1); + + address[] memory newSigners = new address[](1); + newSigners[0] = signer2; + + uint64 newThreshold = 2; // change + uint256 futureNonce = custody.signerNonce() + 42; + bytes memory sigs = _signSingleAdd(signer1Pk, newSigners, newThreshold, futureNonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.addSigners(newSigners, newThreshold, MAX_DEADLINE, sigs); + } +} + +// ========================================================================= +// removeSigners +// ========================================================================= +contract ThresholdCustodyTest_RemoveSigners is ThresholdCustodyTest_Base { + function test_success_onlyRemoveSigners() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 2; // do NOT change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertFalse(custody.isSigner(signer3)); + _checkStats(2, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_2() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertFalse(custody.isSigner(signer3)); + _checkStats(2, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_1_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signRemoveSigners(signer3Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer3, sig3); + + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertFalse(custody.isSigner(signer3)); + _checkStats(2, newThreshold, ++nonce); + } + + function test_success_2_of_3_signatures_2_and_3() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig3 = _signRemoveSigners(signer3Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer3, sig3); + + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + + assertTrue(custody.isSigner(signer1)); + assertTrue(custody.isSigner(signer2)); + assertFalse(custody.isSigner(signer3)); + _checkStats(2, newThreshold, ++nonce); + } + + function test_revert_emptySignatures() public { + custody = new ThresholdCustody(twoSigners, 1); + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 3; + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, _emptySigs()); + } + + function test_revert_signatureFromNotSigner() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + // signature from NOT signer + bytes memory notSignerSig = _signRemoveSigners(notSignerPk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, notSigner, notSignerSig); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signSingleRemove(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, sig2); + } + + function test_revert_duplicatedSignature_thresholdNotReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig2(signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_duplicatedSignature_thresholdReached() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + // duplicate signature from signer2 + bytes memory encodedSigs = _encodeMultiSig3(signer1, sig1, signer2, sig2, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_incorrectSignature() public { + custody = new ThresholdCustody(twoSigners, 1); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + // corrupt the signature by changing one byte + sig1[10] = bytes1(sig1[10] ^ 0x01); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_emptyToRemoveArray() public { + custody = new ThresholdCustody(twoSigners, 2); + + address[] memory signersToRemove = new address[](0); + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert(ThresholdCustody.EmptySignersArray.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_removeArrayDoesNotIncludeExistingSigner() public { + custody = new ThresholdCustody(twoSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; // not a signer + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913NonexistentSigner.selector, signer3.toBytes()) + ); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_removeArrayIncludesDuplicateSigner() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](2); + signersToRemove[0] = signer3; + signersToRemove[1] = signer3; // duplicate + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913NonexistentSigner.selector, signer3.toBytes()) + ); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_removingAllSigners() public { + custody = new ThresholdCustody(twoSigners, 2); + + address[] memory signersToRemove = new address[](2); + signersToRemove[0] = signer1; + signersToRemove[1] = signer2; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 0, 2) + ); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_zeroNewThreshold() public { + custody = new ThresholdCustody(twoSigners, 1); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 0; + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleRemove(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert(MultiSignerERC7913.MultiSignerERC7913ZeroThreshold.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_unreachableNewThreshold() public { + custody = new ThresholdCustody(twoSigners, 1); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 2; // cannot be reached with 1 signer + uint256 nonce = custody.signerNonce(); + bytes memory sigs = _signSingleRemove(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + + vm.expectRevert( + abi.encodeWithSelector(MultiSignerERC7913.MultiSignerERC7913UnreachableThreshold.selector, 1, 2) + ); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, sigs); + } + + function test_revert_deadlinePassed() public { + custody = new ThresholdCustody(twoSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); + custody.removeSigners(signersToRemove, newThreshold, block.timestamp - 1, encodedSigs); + } + + function test_revert_outdatedNonce() public { + custody = new ThresholdCustody(threeSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer3; + + uint64 newThreshold = 1; // change + uint256 nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + // First call succeeds + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + + // Second call with same nonce should revert + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } + + function test_revert_futureNonce() public { + custody = new ThresholdCustody(twoSigners, 2); + + address[] memory signersToRemove = new address[](1); + signersToRemove[0] = signer2; + + uint64 newThreshold = 1; // change + uint256 futureNonce = custody.signerNonce() + 42; + bytes memory sig1 = _signRemoveSigners(signer1Pk, signersToRemove, newThreshold, futureNonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, signersToRemove, newThreshold, futureNonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.expectRevert(ThresholdCustody.InvalidSignature.selector); + custody.removeSigners(signersToRemove, newThreshold, MAX_DEADLINE, encodedSigs); + } +} + +// ========================================================================= +// Deposit +// ========================================================================= +contract ThresholdCustodyTest_Deposit is ThresholdCustodyTest_Base { + uint256 startEthBalance = 1 ether; + uint256 startErc20Balance = 1e18; + + function setUp() public override { + super.setUp(); + + custody = new ThresholdCustody(oneSigner, 1); + + vm.deal(address(user), startEthBalance); + token.mint(address(user), startErc20Balance); + + vm.prank(user); + token.approve(address(custody), type(uint256).max); + } + + function test_success_eth() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + custody.deposit{value: depositAmount}(address(0), depositAmount); + assertEq(address(custody).balance, depositAmount); + assertEq(user.balance, startEthBalance - depositAmount); + } + + function test_success_erc20() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + custody.deposit(address(token), depositAmount); + assertEq(token.balanceOf(address(custody)), depositAmount); + assertEq(token.balanceOf(user), startErc20Balance - depositAmount); + } + + function test_eth_emitsEvent() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + vm.expectEmit(true, true, false, true); + emit IDeposit.Deposited(user, address(0), depositAmount); + custody.deposit{value: depositAmount}(address(0), depositAmount); + } + + function test_erc20_emitsEvent() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + vm.expectEmit(true, true, false, true); + emit IDeposit.Deposited(user, address(token), depositAmount); + custody.deposit(address(token), depositAmount); + } + + function test_revert_eth_amountZero() public { + vm.prank(user); + vm.expectRevert(IDeposit.ZeroAmount.selector); + custody.deposit(address(0), 0); + } + + function test_revert_eth_msgValueZero() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); + custody.deposit(address(0), depositAmount); + } + + function test_revert_eth_amountMsgValueMismatch() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); + custody.deposit{value: depositAmount / 2}(address(0), depositAmount); + } + + function test_revert_erc20_amountZero() public { + vm.prank(user); + vm.expectRevert(IDeposit.ZeroAmount.selector); + custody.deposit(address(token), 0); + } + + function test_revert_erc20_nonZeroMsgValue() public { + uint256 depositAmount = 1e17; + + vm.prank(user); + vm.expectRevert(IDeposit.InvalidMsgValue.selector); + custody.deposit{value: 1 ether}(address(token), depositAmount); + } +} + +// ========================================================================= +// startWithdraw +// ========================================================================= +contract ThresholdCustodyTest_StartWithdraw is ThresholdCustodyTest_Base { + function setUp() public override { + super.setUp(); + + custody = new ThresholdCustody(oneSigner, 1); + + vm.deal(address(custody), 1 ether); + } + + function test_success() public { + address token = address(0); + uint256 amount = 1 ether; + uint256 nonce = 1; + uint256 timestamp = block.timestamp; + + vm.prank(signer1); + bytes32 id = custody.startWithdraw(user, token, amount, nonce); + + bytes32 expectedId = Utils.getWithdrawalId(address(custody), user, token, amount, nonce); + assertEq(id, expectedId); + _validateWithdrawalData(id, user, token, amount, false, 1, uint64(timestamp)); + } + + function test_success_emitsEvent() public { + address token = address(0); + uint256 amount = 1 ether; + uint256 nonce = 1; + + vm.prank(signer1); + vm.expectEmit(true, true, true, true); + bytes32 expectedId = Utils.getWithdrawalId(address(custody), user, token, amount, nonce); + emit IWithdraw.WithdrawStarted(expectedId, user, token, amount, nonce); + custody.startWithdraw(user, token, amount, nonce); + } + + function test_success_sameParamsDifferentNonce() public { + address token = address(0); + uint256 amount = 1 ether; + uint256 nonce = 1; + uint256 timestamp = block.timestamp; + + vm.startPrank(signer1); + bytes32 id1 = custody.startWithdraw(user, token, amount, nonce); + bytes32 id2 = custody.startWithdraw(user, token, amount, nonce + 1); + vm.stopPrank(); + + assertTrue(id1 != id2); + uint64 expectedThreshold = 1; + _validateWithdrawalData(id1, user, token, amount, false, expectedThreshold, uint64(timestamp)); + _validateWithdrawalData(id2, user, token, amount, false, expectedThreshold, uint64(timestamp)); + } + + function test_revert_callerNotSigner() public { + vm.prank(user); + vm.expectRevert(ThresholdCustody.NotSigner.selector); + custody.startWithdraw(user, address(0), 1 ether, 1); + } + + function test_revert_userAddressZero() public { + vm.prank(signer1); + vm.expectRevert(ThresholdCustody.InvalidUser.selector); + custody.startWithdraw(address(0), address(0), 1 ether, 1); + } + + function test_revert_zeroAmount() public { + vm.prank(signer1); + vm.expectRevert(IDeposit.ZeroAmount.selector); + custody.startWithdraw(user, address(0), 0, 1); + } + + function test_revert_duplicateNonce() public { + vm.startPrank(signer1); + custody.startWithdraw(user, address(0), 1 ether, 1); + vm.expectRevert(IWithdraw.WithdrawalAlreadyExists.selector); + custody.startWithdraw(user, address(0), 1 ether, 1); + vm.stopPrank(); + } +} + +// ========================================================================= +// finalizeWithdraw +// ========================================================================= +contract ThresholdCustodyTest_FinalizeWithdraw is ThresholdCustodyTest_Base { + uint256 custodyNativeBalance = 1 ether; + uint256 custodyErc20Balance = 10e18; + uint256 withdrawalAmount = 1e17; + uint256 nonce = 1; + + function _setUpTest(address[] memory signers, uint64 threshold, address withdrawalToken, uint256 withdrawalAmount_) + public + returns (bytes32) + { + custody = new ThresholdCustody(signers, threshold); + vm.deal(address(custody), custodyNativeBalance); + token.mint(address(custody), custodyErc20Balance); + + vm.prank(signers[0]); + bytes32 id = custody.startWithdraw(user, withdrawalToken, withdrawalAmount_, nonce); + return id; + } + + function _validateInitialBalances(address user, address token) public view { + if (token == address(0)) { + assertEq(address(custody).balance, custodyNativeBalance); + assertEq(user.balance, 0); + } else { + assertEq(ERC20(token).balanceOf(address(custody)), custodyErc20Balance); + assertEq(ERC20(token).balanceOf(user), 0); + } + } + + function _validateBalanceWithdrawn(address user, address token, uint256 withdrawnAmount) public view { + if (token == address(0)) { + assertEq(address(custody).balance, custodyNativeBalance - withdrawnAmount); + assertEq(user.balance, withdrawnAmount); + } else { + assertEq(ERC20(token).balanceOf(address(custody)), custodyErc20Balance - withdrawnAmount); + assertEq(ERC20(token).balanceOf(user), withdrawnAmount); + } + } + + function test_success_1of1_eth() public { + address withdrawalToken = address(0); + uint64 expectedThreshold = 1; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(oneSigner, expectedThreshold, withdrawalToken, withdrawalAmount); + + _validateInitialBalances(user, withdrawalToken); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + // user, token, amount are cleared; finalized=true, threshold and createdAt remain + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, startedAt); + _validateBalanceWithdrawn(user, withdrawalToken, withdrawalAmount); + } + + function test_success_1of1_erc20() public { + address withdrawalToken = address(token); + uint64 expectedThreshold = 1; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(oneSigner, expectedThreshold, withdrawalToken, withdrawalAmount); + + _validateInitialBalances(user, withdrawalToken); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + // user, token, amount are cleared; finalized=true, threshold and createdAt remain + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, startedAt); + _validateBalanceWithdrawn(user, withdrawalToken, withdrawalAmount); + } + + function test_success_2of2_eth() public { + address withdrawalToken = address(0); + uint64 expectedThreshold = 2; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(twoSigners, expectedThreshold, withdrawalToken, withdrawalAmount); + + _validateInitialBalances(user, withdrawalToken); + + vm.warp(block.timestamp + 1 minutes); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, false, expectedThreshold, startedAt); + + vm.prank(signer2); + custody.finalizeWithdraw(id); + + // user, token, amount are cleared; finalized=true, threshold and createdAt remain + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, startedAt); + _validateBalanceWithdrawn(user, withdrawalToken, withdrawalAmount); + } + + function test_success_3of5_eth() public { + address withdrawalToken = address(0); + uint64 expectedThreshold = 3; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(fiveSigners, expectedThreshold, withdrawalToken, withdrawalAmount); + + _validateInitialBalances(user, withdrawalToken); + vm.warp(block.timestamp + 1 minutes); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, false, expectedThreshold, startedAt); + + vm.prank(signer2); + custody.finalizeWithdraw(id); + + _validateInitialBalances(user, withdrawalToken); + vm.warp(block.timestamp + 1 minutes); + + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, false, expectedThreshold, startedAt); + + vm.prank(signer3); + custody.finalizeWithdraw(id); + + // user, token, amount are cleared; finalized=true, threshold and createdAt remain + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, startedAt); + _validateBalanceWithdrawn(user, withdrawalToken, withdrawalAmount); + } + + function test_success_usingSnapshotThreshold_eth() public { + // Setup: 2 signers, threshold=2 + address withdrawalToken = address(0); + uint64 expectedThreshold = 2; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(twoSigners, expectedThreshold, withdrawalToken, withdrawalAmount); + + // Lower threshold to 1 AFTER withdrawal was created + uint64 newThreshold = 1; + nonce = custody.signerNonce(); + bytes memory sig1 = _signSetThreshold(signer1Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory sig2 = _signSetThreshold(signer2Pk, newThreshold, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + vm.prank(signer1); + custody.setThreshold(newThreshold, MAX_DEADLINE, encodedSigs); + assertEq(custody.threshold(), newThreshold); + + vm.warp(block.timestamp + 1 minutes); + + // 1 approval should NOT suffice (snapshot quorum was 2) + vm.prank(signer1); + custody.finalizeWithdraw(id); + + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, false, expectedThreshold, startedAt); + + // 2nd approval should finalize the withdrawal + vm.prank(signer2); + custody.finalizeWithdraw(id); + + // user, token, amount are cleared; finalized=true, threshold and createdAt remain + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, startedAt); + _validateBalanceWithdrawn(user, withdrawalToken, withdrawalAmount); + } + + function test_success_eventsEmitted_2of2_eth() public { + address token = address(0); + bytes32 id = _setUpTest(twoSigners, 2, token, withdrawalAmount); + + vm.expectEmit(true, true, true, true); + emit ThresholdCustody.WithdrawalApproved(id, signer1, 1); + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.expectEmit(true, true, true, true); + emit ThresholdCustody.WithdrawalApproved(id, signer2, 2); + vm.expectEmit(true, true, true, true); + emit IWithdraw.WithdrawFinalized(id, true); + vm.prank(signer2); + custody.finalizeWithdraw(id); + } + + function test_revert_notSigner() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.prank(notSigner); + vm.expectRevert(ThresholdCustody.NotSigner.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_duplicateApproval() public { + bytes32 id = _setUpTest(twoSigners, 2, address(0), withdrawalAmount); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.prank(signer1); + vm.expectRevert(ThresholdCustody.SignerAlreadyApproved.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_nonExistentWithdrawal() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + // corrupt the ID to ensure it doesn't match any real withdrawal + id = bytes32(uint256(id) + 42); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.WithdrawalNotFound.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_alreadyFinalized() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.WithdrawalAlreadyFinalized.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_afterReject() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + // Warp past expiry and reject + vm.warp(block.timestamp + OPERATION_EXPIRY + 1); + vm.prank(signer1); + custody.rejectWithdraw(id); + + // Try to finalize after rejection - should revert + vm.prank(signer1); + vm.expectRevert(IWithdraw.WithdrawalAlreadyFinalized.selector); + custody.finalizeWithdraw(id); + } + + function test_success_exactlyAtExpiryBoundary() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + uint64 createdAt = uint64(block.timestamp); + + // Warp to exactly createdAt + OPERATION_EXPIRY (should succeed) + vm.warp(createdAt + OPERATION_EXPIRY); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + // Verify withdrawal was executed successfully + _validateWithdrawalData(id, address(0), address(0), 0, true, 1, createdAt); + _validateBalanceWithdrawn(user, address(0), withdrawalAmount); + } + + function test_revert_oneSecondAfterExpiry() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + uint64 createdAt = uint64(block.timestamp); + + // Warp to createdAt + OPERATION_EXPIRY + 1 (should revert) + vm.warp(createdAt + OPERATION_EXPIRY + 1); + + vm.prank(signer1); + vm.expectRevert(ThresholdCustody.DeadlineExpired.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_insufficientEth() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), custodyNativeBalance + 1); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); + custody.finalizeWithdraw(id); + } + + function test_revert_insufficientErc20() public { + vm.prank(signer1); + bytes32 id = _setUpTest(oneSigner, 1, address(token), custodyErc20Balance + 1); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); + custody.finalizeWithdraw(id); + } + + function test_multipleConcurrentWithdrawals() public { + custody = new ThresholdCustody(oneSigner, 1); + vm.deal(address(custody), custodyNativeBalance); + + address withdrawalToken = address(0); + uint64 createdAt = uint64(block.timestamp); + + vm.startPrank(signer1); + bytes32 id1 = custody.startWithdraw(user, withdrawalToken, withdrawalAmount, 1); + bytes32 id3 = custody.startWithdraw(user, withdrawalToken, withdrawalAmount, 3); + vm.stopPrank(); + + vm.prank(signer1); + custody.finalizeWithdraw(id1); + + _validateWithdrawalData(id1, address(0), address(0), 0, true, 1, createdAt); + _validateBalanceWithdrawn(user, address(0), withdrawalAmount); + + // id2 left to expire + + vm.prank(signer1); + custody.finalizeWithdraw(id3); + + _validateWithdrawalData(id3, address(0), address(0), 0, true, 1, createdAt); + _validateBalanceWithdrawn(user, address(0), 2 * withdrawalAmount); + } + + function test_removedSignerApprovalIgnored() public { + bytes32 id = _setUpTest(threeSigners, 2, address(0), withdrawalAmount); + + vm.prank(signer2); + custody.finalizeWithdraw(id); + + // Remove signer2 (need 2 sigs since threshold=2) + address[] memory toRemove = new address[](1); + toRemove[0] = signer2; + + nonce = custody.signerNonce(); + bytes memory sig1 = _signRemoveSigners(signer1Pk, toRemove, 2, nonce, MAX_DEADLINE); + bytes memory sig2 = _signRemoveSigners(signer2Pk, toRemove, 2, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig2(signer1, sig1, signer2, sig2); + + custody.removeSigners(toRemove, 2, MAX_DEADLINE, encodedSigs); + + assertFalse(custody.isSigner(signer2)); + + // signer1 approves — only 1 valid approval (signer2's no longer counts) + // snapshotted requiredQuorum is still 2, so not finalized yet + vm.prank(signer1); + custody.finalizeWithdraw(id); + + (,,, bool finalized,,) = custody.withdrawals(id); + assertFalse(finalized); + + // signer3 approves — now 2 valid approvals (signer1 + signer3), meets requiredQuorum=2 + vm.prank(signer3); + custody.finalizeWithdraw(id); + + (,,, finalized,,) = custody.withdrawals(id); + assertTrue(finalized); + } +} + +// ========================================================================= +// rejectWithdraw +// ========================================================================= +contract ThresholdCustodyTest_RejectWithdraw is ThresholdCustodyTest_Base { + uint256 custodyNativeBalance = 1 ether; + uint256 custodyErc20Balance = 10e18; + uint256 withdrawalAmount = 1e17; + uint256 nonce = 1; + + function _setUpTest(address[] memory signers, uint64 threshold, address withdrawalToken, uint256 withdrawalAmount_) + public + returns (bytes32) + { + custody = new ThresholdCustody(signers, threshold); + vm.deal(address(custody), custodyNativeBalance); + token.mint(address(custody), custodyErc20Balance); + + vm.prank(signers[0]); + bytes32 id = custody.startWithdraw(user, withdrawalToken, withdrawalAmount_, nonce); + return id; + } + + function _validateInitialBalances(address user, address token) public view { + if (token == address(0)) { + assertEq(address(custody).balance, custodyNativeBalance); + assertEq(user.balance, 0); + } else { + assertEq(ERC20(token).balanceOf(address(custody)), custodyErc20Balance); + assertEq(ERC20(token).balanceOf(user), 0); + } + } + + function test_success() public { + address withdrawalToken = address(0); + uint64 expectedThreshold = 1; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(oneSigner, expectedThreshold, withdrawalToken, withdrawalAmount); + + vm.warp(startedAt + OPERATION_EXPIRY + 1); + + vm.prank(signer1); + custody.rejectWithdraw(id); + + // rejectWithdraw only sets finalized=true, does NOT clear user/token/amount + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, true, expectedThreshold, startedAt); + } + + function test_success_emitsEvent() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.warp(block.timestamp + OPERATION_EXPIRY + 1); + + vm.expectEmit(true, false, false, true); + emit IWithdraw.WithdrawFinalized(id, false); + vm.prank(signer1); + custody.rejectWithdraw(id); + } + + function test_success_afterPartiallyApproved() public { + address withdrawalToken = address(0); + uint64 expectedThreshold = 2; + uint64 startedAt = uint64(block.timestamp); + bytes32 id = _setUpTest(twoSigners, expectedThreshold, withdrawalToken, withdrawalAmount); + + // Get 1 approval (not enough to finalize) + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.warp(startedAt + OPERATION_EXPIRY + 1); + + (,,, bool finalized,,) = custody.withdrawals(id); + assertFalse(finalized); + + vm.prank(signer1); + custody.rejectWithdraw(id); + + // rejectWithdraw only sets finalized=true, does NOT clear user/token/amount + _validateInitialBalances(user, withdrawalToken); + _validateWithdrawalData(id, user, withdrawalToken, withdrawalAmount, true, expectedThreshold, startedAt); + } + + function test_revert_notExpired() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + // do NOT warp past expiry + + vm.prank(signer1); + vm.expectRevert(ThresholdCustody.WithdrawalNotExpired.selector); + custody.rejectWithdraw(id); + } + + function test_revert_notExpired_lastSecond() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.warp(block.timestamp + OPERATION_EXPIRY); + + vm.prank(signer1); + vm.expectRevert(ThresholdCustody.WithdrawalNotExpired.selector); + custody.rejectWithdraw(id); + } + + function test_revert_nonExistent() public { + custody = new ThresholdCustody(oneSigner, 1); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.WithdrawalNotFound.selector); + custody.rejectWithdraw(bytes32(uint256(999))); + } + + function test_revert_alreadyFinalized() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.warp(block.timestamp + OPERATION_EXPIRY); + + vm.prank(signer1); + vm.expectRevert(IWithdraw.WithdrawalAlreadyFinalized.selector); + custody.rejectWithdraw(id); + } + + function test_revert_notSigner() public { + bytes32 id = _setUpTest(oneSigner, 1, address(0), withdrawalAmount); + + vm.warp(block.timestamp + OPERATION_EXPIRY); + + vm.prank(user); + vm.expectRevert(ThresholdCustody.NotSigner.selector); + custody.rejectWithdraw(id); + } + + function test_success_rejectWithPartialApprovals() public { + bytes32 id = _setUpTest(threeSigners, 3, address(0), withdrawalAmount); + + // Get 2 approvals (not enough to finalize) + vm.prank(signer1); + custody.finalizeWithdraw(id); + + vm.prank(signer2); + custody.finalizeWithdraw(id); + + // Verify not finalized yet + (,,, bool finalized1,,) = custody.withdrawals(id); + assertFalse(finalized1); + + // Warp past expiry + vm.warp(block.timestamp + OPERATION_EXPIRY + 1); + + // Reject the withdrawal + vm.prank(signer1); + custody.rejectWithdraw(id); + + // Verify finalized with success=false + (,,, bool finalized2,,) = custody.withdrawals(id); + assertTrue(finalized2); + } +} + +// ========================================================================= +// _executeWithdrawal (internal function via harness) +// ========================================================================= +contract ThresholdCustodyTest_ExecuteWithdrawal is ThresholdCustodyTest_Base { + TestThresholdCustody public testCustody; + + uint256 custodyNativeBalance = 1 ether; + uint256 custodyErc20Balance = 10e18; + uint256 withdrawalAmount = 1e17; + + function setUp() public override { + super.setUp(); + testCustody = new TestThresholdCustody(oneSigner, 1); + vm.deal(address(testCustody), custodyNativeBalance); + token.mint(address(testCustody), custodyErc20Balance); + } + + function _createWithdrawal(address withdrawalToken, uint256 amount) internal returns (bytes32) { + vm.prank(signer1); + bytes32 id = testCustody.startWithdraw(user, withdrawalToken, amount, 1); + return id; + } + + function _validateWithdrawalData( + bytes32 id, + address expectedUser, + address expectedToken, + uint256 expectedAmount, + bool expectedFinalized, + uint64 expectedThreshold, + uint64 expectedCreatedAt + ) internal view override { + ( + address storedUser, + address storedToken, + uint256 storedAmount, + bool storedFinalized, + uint64 storedCreatedAt, + uint64 storedThreshold + ) = testCustody.withdrawals(id); + assertEq(storedUser, expectedUser); + assertEq(storedToken, expectedToken); + assertEq(storedAmount, expectedAmount); + assertEq(storedFinalized, expectedFinalized); + assertEq(storedThreshold, expectedThreshold); + assertEq(storedCreatedAt, expectedCreatedAt); + } + + function test_success_eth() public { + bytes32 id = _createWithdrawal(address(0), withdrawalAmount); + uint64 expectedThreshold = 1; + uint64 createdAt = uint64(block.timestamp); + + testCustody.exposed_executeWithdrawal(id); + + // Verify finalized is set to true, user/token/amount are cleared + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, createdAt); + + assertEq(address(testCustody).balance, custodyNativeBalance - withdrawalAmount); + assertEq(user.balance, withdrawalAmount); + } + + function test_success_erc20() public { + bytes32 id = _createWithdrawal(address(token), withdrawalAmount); + uint64 expectedThreshold = 1; + uint64 createdAt = uint64(block.timestamp); + + testCustody.exposed_executeWithdrawal(id); + + // Verify finalized is set to true, user/token/amount are cleared + _validateWithdrawalData(id, address(0), address(0), 0, true, expectedThreshold, createdAt); + + assertEq(token.balanceOf(address(testCustody)), custodyErc20Balance - withdrawalAmount); + assertEq(token.balanceOf(user), withdrawalAmount); + } + + function test_revert_eth_insufficientLiquidity() public { + bytes32 id = _createWithdrawal(address(0), custodyNativeBalance + 1); + + vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); + testCustody.exposed_executeWithdrawal(id); + } + + function test_revert_erc20_insufficientLiquidity() public { + bytes32 id = _createWithdrawal(address(token), custodyErc20Balance + 1); + + vm.expectRevert(IWithdraw.InsufficientLiquidity.selector); + testCustody.exposed_executeWithdrawal(id); + } +} + +// ========================================================================= +// _countValidApprovals (internal function via harness) +// ========================================================================= +contract ThresholdCustodyTest_CountValidApprovals is ThresholdCustodyTest_Base { + using {Utils.hashArrayed} for address[]; + + TestThresholdCustody public testCustody; + bytes32 withdrawalId; + + function setUp() public override { + super.setUp(); + } + + function _setupCustody(address[] memory signers, uint64 threshold) internal { + testCustody = new TestThresholdCustody(signers, threshold); + vm.deal(address(testCustody), 1 ether); + + // Create a withdrawal + vm.prank(signers[0]); + withdrawalId = testCustody.startWithdraw(user, address(0), 1e17, 1); + } + + function _testCustodyDomainSeparator() internal view returns (bytes32) { + return _domainSeparator(address(testCustody)); + } + + function _signRemoveSignersForTestCustody( + uint256 pk, + address[] memory signersToRemove, + uint256 newThreshold, + uint256 nonce, + uint256 deadline + ) internal view returns (bytes memory) { + bytes32 structHash = keccak256( + abi.encode(REMOVE_SIGNERS_TYPEHASH, signersToRemove.hashArrayed(), newThreshold, nonce, deadline) + ); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _testCustodyDomainSeparator(), structHash)); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest); + return abi.encodePacked(r, s, v); + } + + function test_zero_forNoApprovals() public { + _setupCustody(threeSigners, 2); + + uint256 count = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(count, 0); + } + + function test_1_for1Approval() public { + _setupCustody(threeSigners, 2); + + testCustody.workaround_setWithdrawalApproval(withdrawalId, signer1, true); + + uint256 count = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(count, 1); + } + + function test_approvalReduces_afterSignerRemoval() public { + _setupCustody(threeSigners, 1); + + testCustody.workaround_setWithdrawalApproval(withdrawalId, signer1, true); + testCustody.workaround_setWithdrawalApproval(withdrawalId, signer2, true); + + uint256 countBefore = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(countBefore, 2, "Should have 2 approvals before removal"); + + // Remove signer2 using threshold=1 + address[] memory toRemove = new address[](1); + toRemove[0] = signer2; + + uint256 nonce = testCustody.signerNonce(); + bytes memory sig = _signRemoveSignersForTestCustody(signer1Pk, toRemove, 1, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig(signer1, sig); + + testCustody.removeSigners(toRemove, 1, MAX_DEADLINE, encodedSigs); + + // Count should now be 1 (only signer1's approval counts) + uint256 countAfter = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(countAfter, 1, "Should have 1 approval after removal"); + } + + function _signerPk(uint256 i) internal pure returns (uint256) { + return uint256(keccak256(abi.encode("fuzz_signer", i))) % (type(uint256).max - 1) + 1; + } + + // restrict to uint8 to avoid memory issues with large arrays in fuzzing + function testFuzz_countValidApprovals(uint8 x, uint8 y, uint8 z) public { + // Constrain: x signers total, y approvals, z signers to remove + vm.assume(x >= 1); + vm.assume(y <= x); + vm.assume(z <= y); + vm.assume(z < x); // Cannot remove all signers + + // Create array of x signers + address[] memory signers = new address[](x); + for (uint256 i = 0; i < x; i++) { + signers[i] = vm.addr(_signerPk(i)); + } + + // Setup custody with threshold=1 for easy operations + testCustody = new TestThresholdCustody(signers, 1); + + withdrawalId = bytes32("deadbeef"); + + // Approve y signers + for (uint256 i = 0; i < y; i++) { + testCustody.workaround_setWithdrawalApproval(withdrawalId, signers[i], true); + } + + uint256 countBefore = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(countBefore, y, "Should have y approvals before removal"); + + // Remove z signers (from the ones that approved) + if (z > 0) { + address[] memory toRemove = new address[](z); + for (uint256 i = 0; i < z; i++) { + toRemove[i] = signers[i]; + } + + uint256 nonce = testCustody.signerNonce(); + uint256 signerPk = _signerPk(0); + bytes memory sig = _signRemoveSignersForTestCustody(signerPk, toRemove, 1, nonce, MAX_DEADLINE); + bytes memory encodedSigs = _encodeMultiSig(signers[0], sig); + + testCustody.removeSigners(toRemove, 1, MAX_DEADLINE, encodedSigs); + + uint256 countAfter = testCustody.exposed_countValidApprovals(withdrawalId); + assertEq(countAfter, y - z, "Should have (y - z) approvals after removal"); + } + } +}