Merge pull request #115 from GenerationSoftware/gen-1778-163-add-prot…

…ection-against-gas-bombs-in-prize-hooks

Add protection against gas bombs in prize hooks

.gitmodules

@@ -37,3 +37,6 @@
 [submodule "lib/yield-daddy"]
 	path = lib/yield-daddy
 	url = https://github.com/timeless-fi/yield-daddy
+[submodule "lib/ExcessivelySafeCall"]
+	path = lib/ExcessivelySafeCall
+	url = https://github.com/nomad-xyz/ExcessivelySafeCall

remappings.txt

@@ -6,6 +6,7 @@ brokentoken/=lib/brokentoken/src/
 openzeppelin/=lib/openzeppelin-contracts/contracts/
 owner-manager-contracts/=lib/owner-manager-contracts/contracts/
 yield-daddy/=lib/yield-daddy/src/
+excessively-safe-call/=lib/ExcessivelySafeCall/src/
 
 pt-v5-liquidator-interfaces/=lib/pt-v5-liquidator-interfaces/src/interfaces/
 pt-v5-twab-controller=lib/pt-v5-twab-controller/src/

src/abstract/Claimable.sol

@@ -3,14 +3,17 @@ pragma solidity ^0.8.24;
 
 import { IClaimable } from "pt-v5-claimable-interface/interfaces/IClaimable.sol";
 import { PrizePool } from "pt-v5-prize-pool/PrizePool.sol";
+import { ExcessivelySafeCall  } from "excessively-safe-call/ExcessivelySafeCall.sol";
 
 import { HookManager } from "./HookManager.sol";
+import { IPrizeHooks } from "../interfaces/IPrizeHooks.sol";
 
 /// @title  PoolTogether V5 Claimable Vault Extension
 /// @author G9 Software Inc.
 /// @notice Provides an interface for Claimer contracts to interact with a vault in PoolTogether
 /// V5 while allowing each account to set and manage prize hooks that are called when they win.
 abstract contract Claimable is HookManager, IClaimable {
+    using ExcessivelySafeCall for address;
 
     ////////////////////////////////////////////////////////////////////////////////
     // Public Constants and Variables
@@ -20,6 +23,13 @@ abstract contract Claimable is HookManager, IClaimable {
     /// @dev This should be enough gas to mint an NFT if needed.
     uint24 public constant HOOK_GAS = 150_000;
 
+    /// @notice The number of bytes to limit hook return / revert data.
+    /// @dev If this limit is exceeded for `beforeClaimPrize` return data, the claim will revert.
+    /// @dev Revert data for both hooks will also be limited to this size.
+    /// @dev 128 bytes is enough for `beforeClaimPrize` to return the `_prizeRecipient` address as well
+    /// as 32 bytes of additional `_hookData` byte string data (32 for offset, 32 for length, 32 for data).
+    uint16 public constant HOOK_RETURN_DATA_LIMIT = 128;
+
     /// @notice Address of the PrizePool that computes prizes.
     PrizePool public immutable prizePool;
 
@@ -44,6 +54,11 @@ abstract contract Claimable is HookManager, IClaimable {
     /// @param claimer The claimer address
     error CallerNotClaimer(address caller, address claimer);
 
+    /// @notice Thrown if relevant hook return data is greater than the `HOOK_RETURN_DATA_LIMIT`.
+    /// @param returnDataSize The actual size of the return data
+    /// @param hookDataLimit The return data size limit for hooks
+    error ReturnDataOverLimit(uint256 returnDataSize, uint256 hookDataLimit);
+
     ////////////////////////////////////////////////////////////////////////////////
     // Modifiers
     ////////////////////////////////////////////////////////////////////////////////
@@ -73,6 +88,7 @@ abstract contract Claimable is HookManager, IClaimable {
 
     /// @inheritdoc IClaimable
     /// @dev Also calls the before and after claim hooks if set by the winner.
+    /// @dev Reverts if the return data size of the `beforeClaimPrize` hook exceeds `HOOK_RETURN_DATA_LIMIT`.
     function claimPrize(
         address _winner,
         uint8 _tier,
@@ -84,13 +100,23 @@ abstract contract Claimable is HookManager, IClaimable {
         bytes memory _hookData;
 
         if (_hooks[_winner].useBeforeClaimPrize) {
-            (_prizeRecipient, _hookData) = _hooks[_winner].implementation.beforeClaimPrize{ gas: HOOK_GAS }(
-                _winner,
-                _tier,
-                _prizeIndex,
-                _reward,
-                _rewardRecipient
+            (bytes memory _returnData, uint256 _actualReturnDataSize) = _safeHookCall(
+                _hooks[_winner].implementation,
+                abi.encodeWithSelector(
+                    IPrizeHooks.beforeClaimPrize.selector,
+                    _winner,
+                    _tier,
+                    _prizeIndex,
+                    _reward,
+                    _rewardRecipient
+                )
             );
+            // If the actual return data is greater than the `HOOK_RETURN_DATA_LIMIT` then we must revert since the
+            // integrity of the data is not guaranteed.
+            if (_actualReturnDataSize > HOOK_RETURN_DATA_LIMIT) {
+                revert ReturnDataOverLimit(_actualReturnDataSize, HOOK_RETURN_DATA_LIMIT);
+            }
+            (_prizeRecipient, _hookData) = abi.decode(_returnData, (address, bytes));
         } else {
             _prizeRecipient = _winner;
         }
@@ -107,13 +133,17 @@ abstract contract Claimable is HookManager, IClaimable {
         );
 
         if (_hooks[_winner].useAfterClaimPrize) {
-            _hooks[_winner].implementation.afterClaimPrize{ gas: HOOK_GAS }(
-                _winner,
-                _tier,
-                _prizeIndex,
-                _prizeTotal - _reward,
-                _prizeRecipient,
-                _hookData
+            _safeHookCall(
+                _hooks[_winner].implementation,
+                abi.encodeWithSelector(
+                    IPrizeHooks.afterClaimPrize.selector,
+                    _winner,
+                    _tier,
+                    _prizeIndex,
+                    _prizeTotal - _reward,
+                    _prizeRecipient,
+                    _hookData
+                )
             );
         }
 
@@ -132,5 +162,36 @@ abstract contract Claimable is HookManager, IClaimable {
         claimer = _claimer;
         emit ClaimerSet(_claimer);
     }
+
+    /// @notice Uses ExcessivelySafeCall to limit the return data size to a safe limit.
+    /// @dev This is used for both hook calls to prevent gas bombs that can be triggered using a large
+    /// amount of return data or a large revert string.
+    /// @dev In the case of an unsuccessful call, the revert reason will be bubbled up if it is within
+    /// the safe data limit. Otherwise, a `ReturnDataOverLimit` reason will be thrown.
+    /// @return _returnData The safe, size limited return data
+    /// @return _actualReturnDataSize The actual return data size of the original result
+    function _safeHookCall(IPrizeHooks _implementation, bytes memory _calldata) internal returns (bytes memory _returnData, uint256 _actualReturnDataSize) {
+        bool _success;
+        (_success, _returnData) = address(_implementation).excessivelySafeCall(
+            HOOK_GAS,
+            0, // value
+            HOOK_RETURN_DATA_LIMIT,
+            _calldata
+        );
+        assembly {
+            _actualReturnDataSize := returndatasize()
+        }
+
+        if (!_success) {
+            // If we can't access the full revert data, we use a generic revert
+            if (_actualReturnDataSize > HOOK_RETURN_DATA_LIMIT) {
+                revert ReturnDataOverLimit(_actualReturnDataSize, HOOK_RETURN_DATA_LIMIT);
+            }
+            // Otherwise, we use a low level revert to bubble up the revert reason
+            assembly {
+                revert(add(32, _returnData), mload(_returnData))
+            }
+        }
+    }
 
 }

test/unit/Claimable.t.sol

@@ -250,6 +250,236 @@ contract ClaimableTest is Test, IPrizeHooks {
         claimable.claimPrize(alice, 1, 2, 1e17, bob);
     }
 
+    function testClaimPrize_beforeClaimPrize_ReturnDataOverLimit_hookData() public {
+        bytes memory tooMuchHookData = new bytes(33); // 33 bytes of data
+
+        // mock to return too much hook data
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.beforeClaimPrize.selector,
+                alice, 1, 2, 1e17, bob
+            ),
+            abi.encode(alice, tooMuchHookData)
+        );
+
+        PrizeHooks memory bothHooks = PrizeHooks(true, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(bothHooks);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        // We expect a revert since the `beforeClaimPrize` function returns too much data
+        vm.expectRevert(abi.encodeWithSelector(Claimable.ReturnDataOverLimit.selector, 160, 128)); // 5 words used, which is 32 bytes over the limit
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_beforeClaimPrize_ReturnDataOverLimit_onRevert() public {
+        bytes memory tooMuchRevertData = new bytes(129); // 129 bytes of revert data
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.beforeClaimPrize.selector,
+                alice, 1, 2, 1e17, bob
+            ),
+            tooMuchRevertData
+        );
+
+        PrizeHooks memory beforeHookOnly = PrizeHooks(true, false, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(beforeHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, prizeRedirectionAddress, 1e17, bob);
+
+        vm.expectRevert(abi.encodeWithSelector(Claimable.ReturnDataOverLimit.selector, 129, 128)); // 1 byte over the limit
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_afterClaimPrize_ReturnDataOverLimit_onRevert() public {
+        bytes memory tooMuchRevertData = new bytes(129); // 129 bytes of revert data
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.afterClaimPrize.selector,
+                alice, 1, 2, 9e17, alice, ""
+            ),
+            tooMuchRevertData
+        );
+
+        PrizeHooks memory afterHookOnly = PrizeHooks(false, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(afterHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        vm.expectRevert(abi.encodeWithSelector(Claimable.ReturnDataOverLimit.selector, 129, 128)); // 1 byte over the limit
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_beforeClaimPrize_bubblesRevertData() public {
+        bytes memory maxRevertData = new bytes(128); // 128 bytes of revert data
+        maxRevertData[0] = bytes1(0x01); // make the data non-zero
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.beforeClaimPrize.selector,
+                alice, 1, 2, 1e17, bob
+            ),
+            maxRevertData
+        );
+
+        PrizeHooks memory beforeHookOnly = PrizeHooks(true, false, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(beforeHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, prizeRedirectionAddress, 1e17, bob);
+
+        vm.expectRevert(maxRevertData); // expected bubbled revert data
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_afterClaimPrize_bubblesRevertData() public {
+        bytes memory maxRevertData = new bytes(128); // 128 bytes of revert data
+        maxRevertData[0] = bytes1(0x01); // make the data non-zero
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.afterClaimPrize.selector,
+                alice, 1, 2, 9e17, alice, ""
+            ),
+            maxRevertData
+        );
+
+        PrizeHooks memory afterHookOnly = PrizeHooks(false, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(afterHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        vm.expectRevert(maxRevertData); // expected bubbled revert data
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_beforeClaimPrize_bubblesEmptyRevertData() public {
+        bytes memory emptyRevertData = new bytes(0); // 0 bytes
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.beforeClaimPrize.selector,
+                alice, 1, 2, 1e17, bob
+            ),
+            emptyRevertData
+        );
+
+        PrizeHooks memory beforeHookOnly = PrizeHooks(true, false, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(beforeHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, prizeRedirectionAddress, 1e17, bob);
+
+        vm.expectRevert(emptyRevertData); // expected bubbled empty revert data
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_afterClaimPrize_bubblesEmptyRevertData() public {
+        bytes memory emptyRevertData = new bytes(0); // 0 bytes
+
+        // mock to revert with too much data
+        vm.mockCallRevert(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.afterClaimPrize.selector,
+                alice, 1, 2, 9e17, alice, ""
+            ),
+            emptyRevertData
+        );
+
+        PrizeHooks memory afterHookOnly = PrizeHooks(false, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(afterHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        vm.expectRevert(emptyRevertData); // expected bubbled empty revert data
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_beforeClaimPrize_noRevertOnMaxHookData() public {
+        bytes memory maxHookData = new bytes(32); // 32 bytes of data
+
+        // mock to return max hook data
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.beforeClaimPrize.selector,
+                alice, 1, 2, 1e17, bob
+            ),
+            abi.encode(alice, maxHookData)
+        );
+
+        PrizeHooks memory bothHooks = PrizeHooks(true, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(bothHooks);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        // expect the after hook to be called with the max hook data with no issues
+        vm.expectEmit();
+        emit AfterClaimPrizeCalled(alice, 1, 2, 9e17, alice, maxHookData);
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
+    function testClaimPrize_afterClaimPrize_ignoresHookData() public {
+        bytes memory tooMuchHookData = new bytes(33); // 33 bytes of data
+
+        // mock to return too much hook data
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSelector(
+                IPrizeHooks.afterClaimPrize.selector,
+                alice, 1, 2, 9e17, alice, ""
+            ),
+            abi.encode(alice, tooMuchHookData)
+        );
+
+        PrizeHooks memory afterHookOnly = PrizeHooks(false, true, hooks);
+
+        vm.startPrank(alice);
+        claimable.setHooks(afterHookOnly);
+        vm.stopPrank();
+
+        mockClaimPrize(alice, 1, 2, alice, 1e17, bob);
+
+        // expect no revert even though the return data has been limited since the afterClaimPrize return data isn't used
+        claimable.claimPrize(alice, 1, 2, 1e17, bob);
+    }
+
     /* ============ IPrizeHooks Implementation ============ */
 
     function beforeClaimPrize(