Make function calls atomic with ERC20 transfers in token supply source subnets

contracts/src/lib/SupplySourceHelper.sol

@@ -72,7 +72,7 @@ library SupplySourceHelper {
     /// This function the `safeTransfer` function used before.
     function ierc20Transfer(
         SupplySource memory supplySource,
-        address payable recipient,
+        address recipient,
         uint256 value
     ) internal returns (bool success, bytes memory ret) {
         return
@@ -101,14 +101,38 @@ library SupplySourceHelper {
             // Use the optimized path to send value along with the call.
             (success, ret) = functionCallWithValue({target: target, data: data, value: value});
         } else if (supplySource.kind == SupplyKind.ERC20) {
-            // Transfer the tokens first, _then_ perform the call.
-            (success, ret) = ierc20Transfer(supplySource, target, value);
-            if (success) {
-                // Perform the call only if the ERC20 was successful.
-                (success, ret) = functionCallWithValue(target, data, 0);
-            } else {
-                return (success, ret);
+            (success, ret) = functionCallWithERC20Value({supplySource: supplySource, target: target, data: data, value: value});
+        }
+        return (success, ret);
+    }
+
+    /// @dev Performs the function call with ERC20 value atomically
+    function functionCallWithERC20Value(
+        SupplySource memory supplySource,
+        address target,
+        bytes memory data,
+        uint256 value
+    ) internal returns (bool success, bytes memory ret) {
+        // Transfer the tokens first, _then_ perform the call.
+        (success, ret) = ierc20Transfer(supplySource, target, value);
+
+        if (success) {
+            // Perform the call only if the ERC20 was successful.
+            (success, ret) = functionCallWithValue(target, data, 0);
+        }
+
+        if (!success) {
+            // following the implementation of `openzeppelin-contracts/utils/Address.sol`
+            if (ret.length > 0) {
+                assembly {
+                    let returndata_size := mload(ret)
+                    // see https://ethereum.stackexchange.com/questions/133748/trying-to-understand-solidity-assemblys-revert-function
+                    revert(add(32, ret), returndata_size)
+                }
             }
+            // disable solhint as the failing call does not have return data as well.
+            /* solhint-disable reason-string */
+            revert();
         }
         return (success, ret);
     }

contracts/test/mocks/SupplySourceHelperMock.sol

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.19;
+
+import {SupplySource} from "../../src/structs/Subnet.sol";
+import {SupplySourceHelper} from "../../src/lib/SupplySourceHelper.sol";
+
+/// @notice Helpers to deal with a supply source.
+contract SupplySourceHelperMock {
+    function performCall(
+        SupplySource memory supplySource,
+        address payable target,
+        bytes memory data,
+        uint256 value
+    ) public returns (bool success, bytes memory ret) {
+        return SupplySourceHelper.performCall(supplySource, target, data, value);
+    }
+}

contracts/test/unit/SupplySourceHelper.t.sol

@@ -0,0 +1,116 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.19;
+
+import "forge-std/Test.sol";
+import "openzeppelin-contracts/utils/Strings.sol";
+import "../../src/lib/SubnetIDHelper.sol";
+
+import {SupplySource, SupplyKind} from "../../src/structs/Subnet.sol";
+import {SupplySourceHelper} from "../../src/lib/SupplySourceHelper.sol";
+
+import {SupplySourceHelperMock} from "../mocks/SupplySourceHelperMock.sol";
+
+import {IERC20} from "openzeppelin-contracts/token/ERC20/IERC20.sol";
+
+import {ERC20PresetFixedSupply} from "../helpers/ERC20PresetFixedSupply.sol";
+
+contract FailingContract {
+    error BOOM();
+
+    function failing() external pure {
+        revert BOOM();
+    }
+}
+
+contract SupplySourceHelperTest is Test {
+    /// Call fails but send value works, both should fail
+    function test_revert_atomicity_no_ret() public {
+        uint256 balance = 1_000_000;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        IERC20 token = new ERC20PresetFixedSupply("TestToken", "TEST", balance, address(mock));
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.ERC20, tokenAddress: address(token)});
+
+        bytes memory params = bytes("hello");
+
+        vm.expectRevert();
+        mock.performCall(source, payable(address(this)), params, 100);
+
+        require(token.balanceOf(address(mock)) == balance, "invalid balance");
+    }
+
+    function test_revert_atomicity_with_ret() public {
+        uint256 balance = 1_000_000;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        IERC20 token = new ERC20PresetFixedSupply("TestToken", "TEST", balance, address(mock));
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.ERC20, tokenAddress: address(token)});
+
+        bytes memory params = abi.encodeWithSelector(FailingContract.failing.selector);
+
+        address c = address(new FailingContract());
+        vm.expectRevert(FailingContract.BOOM.selector);
+        mock.performCall(source, payable(c), params, 100);
+
+        require(token.balanceOf(address(mock)) == balance, "invalid balance");
+    }
+
+    function test_call_with_erc20_ok() public {
+        uint256 balance = 1_000_000;
+        uint256 value = 100;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        IERC20 token = new ERC20PresetFixedSupply("TestToken", "TEST", balance, address(mock));
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.ERC20, tokenAddress: address(token)});
+
+        bytes memory params = bytes("hello");
+
+        mock.performCall(source, payable(address(1)), params, value);
+
+        require(token.balanceOf(address(mock)) == balance - value, "invalid balance");
+        require(token.balanceOf(address(1)) == value, "invalid user balance");
+    }
+
+    function test_call_with_native_zero_balance_ok() public {
+        uint256 value = 0;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.Native, tokenAddress: address(0)});
+
+        bytes memory params = bytes("hello");
+
+        mock.performCall(source, payable(address(1)), params, value);
+        require(address(1).balance == 0, "invalid user balance");
+    }
+
+    function test_call_with_native_ok() public {
+        uint256 value = 10;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        vm.deal(address(mock), 1 ether);
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.Native, tokenAddress: address(0)});
+
+        bytes memory params = bytes("hello");
+
+        mock.performCall(source, payable(address(1)), params, value);
+        require(address(1).balance == value, "invalid user balance");
+    }
+
+    function test_call_with_native_reverts() public {
+        uint256 value = 10;
+        SupplySourceHelperMock mock = new SupplySourceHelperMock();
+
+        vm.deal(address(mock), 1 ether);
+
+        SupplySource memory source = SupplySource({kind: SupplyKind.Native, tokenAddress: address(0)});
+
+        bytes memory params = bytes("hello");
+
+        mock.performCall(source, payable(address(this)), params, value);
+        require(address(1).balance == 0, "invalid user balance");
+    }
+}