Balance Change Enforcers Underflow Protection

src/enforcers/ERC1155BalanceChangeEnforcer.sol

@@ -78,12 +78,7 @@ contract ERC1155BalanceChangeEnforcer is CaveatEnforcer {
         override
         onlyDefaultExecutionMode(_mode)
     {
-        (, address token_, address recipient_, uint256 tokenId_,) = getTermsInfo(_terms);
-        bytes32 hashKey_ = _getHashKey(msg.sender, token_, recipient_, tokenId_, _delegationHash);
-        require(!isLocked[hashKey_], "ERC1155BalanceChangeEnforcer:enforcer-is-locked");
-        isLocked[hashKey_] = true;
-        uint256 balance_ = IERC1155(token_).balanceOf(recipient_, tokenId_);
-        balanceCache[hashKey_] = balance_;
+        _validateBeforeHook(_terms, _delegationHash);
     }
 
     /**
@@ -161,4 +156,22 @@ contract ERC1155BalanceChangeEnforcer is CaveatEnforcer {
     {
         return keccak256(abi.encode(_caller, _token, _recipient, _tokenId, _delegationHash));
     }
+
+    /**
+     * @notice Internal function to validate and setup state before a delegation is executed
+     * @dev Caches the initial balance and locks the enforcer to prevent reentrancy
+     * @param _terms The encoded terms of the delegation
+     * @param _delegationHash The hash of the delegation being executed
+     */
+    function _validateBeforeHook(bytes calldata _terms, bytes32 _delegationHash) internal {
+        (bool enforceDecrease_, address token_, address recipient_, uint256 tokenId_, uint256 amount_) = getTermsInfo(_terms);
+        bytes32 hashKey_ = _getHashKey(msg.sender, token_, recipient_, tokenId_, _delegationHash);
+        require(!isLocked[hashKey_], "ERC1155BalanceChangeEnforcer:enforcer-is-locked");
+        isLocked[hashKey_] = true;
+        uint256 balance_ = IERC1155(token_).balanceOf(recipient_, tokenId_);
+        if (enforceDecrease_) {
+            require(balance_ >= amount_, "ERC1155BalanceChangeEnforcer:insufficient-initial-balance");
+        }
+        balanceCache[hashKey_] = balance_;
+    }
 }

src/enforcers/ERC20BalanceChangeEnforcer.sol

@@ -64,11 +64,14 @@ contract ERC20BalanceChangeEnforcer is CaveatEnforcer {
         override
         onlyDefaultExecutionMode(_mode)
     {
-        (, address token_, address recipient_,) = getTermsInfo(_terms);
+        (bool enforceDecrease_, address token_, address recipient_, uint256 amount_) = getTermsInfo(_terms);
         bytes32 hashKey_ = _getHashKey(msg.sender, token_, _delegationHash);
         require(!isLocked[hashKey_], "ERC20BalanceChangeEnforcer:enforcer-is-locked");
         isLocked[hashKey_] = true;
         uint256 balance_ = IERC20(token_).balanceOf(recipient_);
+        if (enforceDecrease_) {
+            require(balance_ >= amount_, "ERC20BalanceChangeEnforcer:insufficient-initial-balance");
+        }
         balanceCache[hashKey_] = balance_;
     }
 

src/enforcers/ERC721BalanceChangeEnforcer.sol

@@ -75,11 +75,14 @@ contract ERC721BalanceChangeEnforcer is CaveatEnforcer {
         override
         onlyDefaultExecutionMode(_mode)
     {
-        (, address token_, address recipient_,) = getTermsInfo(_terms);
+        (bool enforceDecrease_, address token_, address recipient_, uint256 amount_) = getTermsInfo(_terms);
         bytes32 hashKey_ = _getHashKey(msg.sender, token_, recipient_, _delegationHash);
         require(!isLocked[hashKey_], "ERC721BalanceChangeEnforcer:enforcer-is-locked");
         isLocked[hashKey_] = true;
         uint256 balance_ = IERC721(token_).balanceOf(recipient_);
+        if (enforceDecrease_) {
+            require(balance_ >= amount_, "ERC721BalanceChangeEnforcer:insufficient-initial-balance");
+        }
         balanceCache[hashKey_] = balance_;
     }
 

src/enforcers/NativeBalanceChangeEnforcer.sol

@@ -62,10 +62,14 @@ contract NativeBalanceChangeEnforcer is CaveatEnforcer {
         onlyDefaultExecutionMode(_mode)
     {
         bytes32 hashKey_ = _getHashKey(msg.sender, _delegationHash);
-        (, address recipient_,) = getTermsInfo(_terms);
+        (bool enforceDecrease_, address recipient_, uint256 amount_) = getTermsInfo(_terms);
         require(!isLocked[hashKey_], "NativeBalanceChangeEnforcer:enforcer-is-locked");
         isLocked[hashKey_] = true;
-        balanceCache[hashKey_] = recipient_.balance;
+        uint256 balance_ = recipient_.balance;
+        if (enforceDecrease_) {
+            require(balance_ >= amount_, "NativeBalanceChangeEnforcer:insufficient-initial-balance");
+        }
+        balanceCache[hashKey_] = balance_;
     }
 
     /**

test/enforcers/ERC1155BalanceChangeEnforcer.t.sol

@@ -312,6 +312,17 @@ contract ERC1155BalanceChangeEnforcerTest is CaveatEnforcerBaseTest {
         enforcer.beforeHook(hex"", hex"", singleTryMode, hex"", bytes32(0), address(0), address(0));
     }
 
+    // Reverts if the initial balance is insufficient for a decrease operation
+    function test_notAllow_insufficientInitialBalance() public {
+        // Terms: flag=true (decrease expected), token, recipient, tokenId, required amount = 100
+        bytes memory terms_ = abi.encodePacked(true, address(token), address(delegator), uint256(tokenId), uint256(100));
+
+        // Should revert because initial balance (0) is less than required amount (100)
+        vm.prank(dm);
+        vm.expectRevert(bytes("ERC1155BalanceChangeEnforcer:insufficient-initial-balance"));
+        enforcer.beforeHook(terms_, hex"", singleDefaultMode, mintExecutionCallData, bytes32(0), address(0), delegate);
+    }
+
     function _getEnforcer() internal view override returns (ICaveatEnforcer) {
         return ICaveatEnforcer(address(enforcer));
     }

test/enforcers/ERC20BalanceChangeEnforcer.t.sol

@@ -240,6 +240,22 @@ contract ERC20BalanceChangeEnforcerTest is CaveatEnforcerBaseTest {
         enforcer.beforeHook(hex"", hex"", singleTryMode, hex"", bytes32(0), address(0), address(0));
     }
 
+    // Reverts if the initial balance is insufficient for a decrease operation
+    function test_notAllow_insufficientInitialBalance() public {
+        // Set an initial balance for the recipient that's less than the required amount
+        uint256 initialBalance_ = 50;
+        vm.prank(delegator);
+        token.mint(recipient, initialBalance_);
+
+        // Terms: flag=true (decrease expected), token, recipient, required amount = 100
+        bytes memory terms_ = abi.encodePacked(true, address(token), address(recipient), uint256(100));
+
+        // Should revert because initial balance (50) is less than required amount (100)
+        vm.prank(dm);
+        vm.expectRevert(bytes("ERC20BalanceChangeEnforcer:insufficient-initial-balance"));
+        enforcer.beforeHook(terms_, hex"", singleDefaultMode, mintExecutionCallData, bytes32(0), delegator, delegate);
+    }
+
     function _getEnforcer() internal view override returns (ICaveatEnforcer) {
         return ICaveatEnforcer(address(enforcer));
     }

test/enforcers/ERC721BalanceChangeEnforcer.t.sol

@@ -306,6 +306,17 @@ contract ERC721BalanceChangeEnforcerTest is CaveatEnforcerBaseTest {
         enforcer.beforeHook(hex"", hex"", singleTryMode, hex"", bytes32(0), address(0), address(0));
     }
 
+    // Reverts if the initial balance is insufficient for a decrease operation
+    function test_notAllow_insufficientInitialBalance() public {
+        // Terms: flag=true (decrease expected), token, recipient, required amount = 2
+        bytes memory terms_ = abi.encodePacked(true, address(token), address(delegator), uint256(2));
+
+        // Should revert because initial balance (0) is less than required amount (2)
+        vm.prank(dm);
+        vm.expectRevert(bytes("ERC721BalanceChangeEnforcer:insufficient-initial-balance"));
+        enforcer.beforeHook(terms_, hex"", singleDefaultMode, mintExecutionCallData, bytes32(0), address(0), delegate);
+    }
+
     function _getEnforcer() internal view override returns (ICaveatEnforcer) {
         return ICaveatEnforcer(address(enforcer));
     }

test/enforcers/NativeBalanceChangeEnforcer.t.sol

@@ -173,6 +173,17 @@ contract NativeBalanceChangeEnforcerTest is CaveatEnforcerBaseTest {
         enforcer.beforeHook(hex"", hex"", singleTryMode, hex"", bytes32(0), address(0), address(0));
     }
 
+    // Reverts if the initial balance is insufficient for a decrease operation
+    function test_notAllow_insufficientInitialBalance() public {
+        // Terms: flag=true (decrease expected), recipient, required amount = 1000 ether
+        bytes memory terms_ = abi.encodePacked(true, address(delegator), uint256(1000 ether));
+
+        // Should revert because initial balance is less than required amount (1000 ether)
+        vm.prank(dm);
+        vm.expectRevert(bytes("NativeBalanceChangeEnforcer:insufficient-initial-balance"));
+        enforcer.beforeHook(terms_, hex"", singleDefaultMode, executionCallData, bytes32(0), address(0), delegate);
+    }
+
     function _increaseBalance(address _recipient, uint256 _amount) internal {
         vm.deal(_recipient, _recipient.balance + _amount);
     }

test/helpers/DelegationMetaSwapAdapter.t.sol

@@ -92,7 +92,7 @@ abstract contract DelegationMetaSwapAdapterBaseTest is BaseTest {
     /**
      * @dev Generates a valid signature for _apiData with a given _expiration.
      */
-    function _getValidSignature(bytes memory _apiData, uint256 _expiration) internal returns (bytes memory) {
+    function _getValidSignature(bytes memory _apiData, uint256 _expiration) internal view returns (bytes memory) {
         bytes32 messageHash = keccak256(abi.encode(_apiData, _expiration));
         bytes32 ethSignedMessageHash = MessageHashUtils.toEthSignedMessageHash(messageHash);
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(swapSignerPrivateKey, ethSignedMessageHash);
@@ -102,7 +102,7 @@ abstract contract DelegationMetaSwapAdapterBaseTest is BaseTest {
     /**
      * @dev Builds and returns a SignatureData struct from the given apiData.
      */
-    function _buildSigData(bytes memory apiData) internal returns (DelegationMetaSwapAdapter.SignatureData memory) {
+    function _buildSigData(bytes memory apiData) internal view returns (DelegationMetaSwapAdapter.SignatureData memory) {
         uint256 expiration = block.timestamp + 1000;
         bytes memory signature = _getValidSignature(apiData, expiration);
         return DelegationMetaSwapAdapter.SignatureData({ apiData: apiData, expiration: expiration, signature: signature });