Merge pull request #24 from KumaCrypto/L04_fix

fix(L04): handle non-native deposit with value & block native deposit with zero value

Author: ccashwell

src/allocators/OnChainAllocator.sol

@@ -87,43 +87,72 @@ contract OnChainAllocator is IOnChainAllocator {
         bytes32 typehash,
         bytes32 witness
     ) public payable returns (bytes32 claimHash, uint256[] memory registeredAmounts, uint256 nonce) {
-        nonce = _getAndUpdateNonce(msg.sender, recipient);
-
+        // Check for empty commitments
         if (commitments.length == 0) {
             revert InvalidCommitments();
         }
 
+        nonce = _getAndUpdateNonce(msg.sender, recipient);
+
+        // Transformed locks will be stored in idsAndAmounts
         uint256[2][] memory idsAndAmounts = new uint256[2][](commitments.length);
 
+        // Init minResetPeriod to the max value
         uint256 minResetPeriod = type(uint256).max;
-        uint256 i = 0;
+
+        // Process native token (zero address) first
+        uint256 i;
         if (commitments[i].token == address(0)) {
+            // The Compact will revert if invalid value is provided for native token
+            // Possible cases:
+            // 1. The callvalue is zero but the first token is native
+            // 2. the callvalue is nonzero but the first token is non-native
+            // 3. the first token is native and the callvalue doesn't equal the first amount
+
+            // Handle first and third points
+            if (commitments[i].amount == 0 || commitments[i].amount != msg.value) {
+                revert InvalidAmount(commitments[i].amount);
+            }
+
             minResetPeriod = _checkInput(commitments[i], recipient, expires, minResetPeriod);
+
             idsAndAmounts[i][0] = AL.toId(commitments[i].lockTag, commitments[i].token);
+            idsAndAmounts[i][1] = msg.value;
 
-            if (commitments[i].amount != 0 && commitments[i].amount != msg.value) {
-                revert InvalidAmount(commitments[i].amount);
+            unchecked {
+                ++i;
+            }
+        } else {
+            // Handle second point
+            if (msg.value != 0) {
+                revert InvalidAmount(msg.value);
             }
-            idsAndAmounts[i][1] = msg.value;
-            i++;
         }
+
+        // Process the rest of the commitments
         for (; i < commitments.length; i++) {
             minResetPeriod = _checkInput(commitments[i], recipient, expires, minResetPeriod);
-            idsAndAmounts[i][0] = AL.toId(commitments[i].lockTag, commitments[i].token);
+
+            address token = commitments[i].token;
+            // Safe to cast - _checkInput validated that the value fits the uint224
             uint224 amount = uint224(commitments[i].amount);
 
-            // If the amount is 0, we use the balance of the contract to deposit.
+            // If the amount is 0, we use the balance of the contract to deposit
             if (amount == 0) {
-                amount = uint224(IERC20(commitments[i].token).balanceOf(address(this)));
+                amount = uint224(IERC20(token).balanceOf(address(this)));
             }
+
+            // Store the lock in idsAndAmounts
+            idsAndAmounts[i][0] = AL.toId(commitments[i].lockTag, token);
             idsAndAmounts[i][1] = amount;
 
             // Approve the compact contract to spend the tokens.
-            if (IERC20(commitments[i].token).allowance(address(this), COMPACT_CONTRACT) < amount) {
-                SafeTransferLib.safeApproveWithRetry(commitments[i].token, COMPACT_CONTRACT, type(uint256).max);
+            if (IERC20(token).allowance(address(this), COMPACT_CONTRACT) < amount) {
+                SafeTransferLib.safeApproveWithRetry(token, COMPACT_CONTRACT, type(uint256).max);
             }
         }
-        // Ensure expiration is not bigger then the smallest reset period
+
+        // Ensure expiration is less then the smallest reset period
         if (expires >= block.timestamp + minResetPeriod) {
             revert InvalidExpiration(expires, block.timestamp + minResetPeriod);
         }

test/OnChainAllocator.t.sol

@@ -1147,6 +1147,27 @@ contract OnChainAllocatorTest is Test, TestHelper {
         );
     }
 
+    function test_allocateAndRegister_revert_InvalidAmount_native_with_zero_deposit() public {
+        Lock[] memory commitments = new Lock[](1);
+        commitments[0] = _makeLock(address(0), 0 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(IOnChainAllocator.InvalidAmount.selector, 0));
+        allocator.allocateAndRegister{value: 0}(
+            recipient, commitments, arbiter, defaultExpiration, BATCH_COMPACT_TYPEHASH, bytes32(0)
+        );
+    }
+
+    function test_allocateAndRegister_revert_InvalidAmount_non_native_with_non_zero_native_call() public {
+        Lock[] memory commitments = new Lock[](1);
+        commitments[0] = _makeLock(address(usdc), 1);
+
+        uint256 amount = 1;
+        vm.expectRevert(abi.encodeWithSelector(IOnChainAllocator.InvalidAmount.selector, amount));
+        allocator.allocateAndRegister{value: amount}(
+            recipient, commitments, arbiter, defaultExpiration, BATCH_COMPACT_TYPEHASH, bytes32(0)
+        );
+    }
+
     function test_allocateAndRegister_revert_InvalidAmount() public {
         Lock[] memory commitments = new Lock[](1);
         commitments[0] = _makeLock(address(usdc), uint256(type(uint224).max) + 1);