Add support of new message prefix

zkBob · Dec 11, 2023 · 70cf973 · 70cf973
1 parent 17f2adf
commit 70cf973
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/src/zkbob/ZkBobPool.sol b/src/zkbob/ZkBobPool.sol
@@ -34,6 +34,7 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
 
     uint256 internal constant MAX_POOL_ID = 0xffffff;
     bytes4 internal constant MESSAGE_PREFIX_COMMON_V1 = 0x00000000;
+    bytes4 internal constant MESSAGE_PREFIX_COMMON_V2 = 0x00000100;
     uint256 internal constant FORCED_EXIT_MIN_DELAY = 1 hours;
     uint256 internal constant FORCED_EXIT_MAX_DELAY = 24 hours;
 
@@ -236,7 +237,7 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
             roots[poolIndex] = _tree_root_after();
             bytes memory message = _memo_message();
             // restrict memo message prefix (items count in little endian) to be < 2**16
-            require(bytes4(message) & 0x0000ffff == MESSAGE_PREFIX_COMMON_V1, "ZkBobPool: bad message prefix");
+            require(_isValidPrefix(bytes4(message) & 0x0000ffff), "ZkBobPool: bad message prefix");
             bytes32 message_hash = keccak256(message);
             bytes32 _all_messages_hash = keccak256(abi.encodePacked(all_messages_hash, message_hash));
             all_messages_hash = _all_messages_hash;
@@ -512,4 +513,13 @@ abstract contract ZkBobPool is IZkBobPool, EIP1967Admin, Ownable, Parameters, Ex
     function _isOwner() internal view override returns (bool) {
         return super._isOwner() || _admin() == _msgSender();
     }
+
+    /**
+     * @dev Tells if given message prefix is valid.
+     * @param _prefix prefix to check.
+     * @return true, if prefix is valid.
+     */
+    function _isValidPrefix(bytes4 _prefix) internal pure returns (bool) {
+        return _prefix == MESSAGE_PREFIX_COMMON_V1 || _prefix == MESSAGE_PREFIX_COMMON_V2;
+    }
 }
diff --git a/test/zkbob/ZkBobPool.t.sol b/test/zkbob/ZkBobPool.t.sol
@@ -719,6 +719,23 @@ abstract contract AbstractZkBobPoolTest is AbstractForkTest {
         assertApproxEqAbs(rewardToken.balanceOf(user1), 1785 ether, 200 ether);
     }
 
+    function testTransactAcceptsOnlyValidPrefixes() public {
+        bytes memory data = _encodePermitDeposit(int256(0.5 ether / D), 0.01 ether / D);
+        _transact(data);
+
+        data = _encodeTransferWithPrefix(0.01 ether / D, bytes2(0x0000));
+        _transact(data);
+
+        data = _encodeTransferWithPrefix(0.01 ether / D, bytes2(0x0100));
+        _transact(data);
+
+        data = _encodeTransferWithPrefix(0.01 ether / D, bytes2(0x0001));
+        _transactReverted(data, "ZkBobPool: bad message prefix");
+
+         data = _encodeTransferWithPrefix(0.01 ether / D, bytes2(0x1234));
+        _transactReverted(data, "ZkBobPool: bad message prefix");
+    }
+
     function _encodeDeposit(int256 _amount, uint256 _fee) internal returns (bytes memory) {
         bytes32 nullifier = bytes32(_randFR());
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk1, ECDSA.toEthSignedMessageHash(nullifier));
@@ -770,13 +787,17 @@ abstract contract AbstractZkBobPoolTest is AbstractForkTest {
     }
 
     function _encodeTransfer(uint256 _fee) internal returns (bytes memory) {
+        return _encodeTransferWithPrefix(_fee, bytes2(0x0000));
+    }
+
+    function _encodeTransferWithPrefix(uint256 _fee, bytes2 prefix) internal returns (bytes memory) {
         bytes memory data = abi.encodePacked(
             ZkBobPool.transact.selector, _randFR(), _randFR(), uint48(0), uint112(0), -int64(uint64(_fee / denominator))
         );
         for (uint256 i = 0; i < 17; i++) {
             data = abi.encodePacked(data, _randFR());
         }
-        return abi.encodePacked(data, uint16(1), uint16(44), uint64(_fee / denominator), bytes4(0x01000000), _randFR());
+        return abi.encodePacked(data, uint16(1), uint16(44), uint64(_fee / denominator), bytes2(0x0100), prefix, _randFR());
     }
 
     function _transact(bytes memory _data) internal {