Allow retrying transmissions with more gas

contracts/src/v0.8/keystone/KeystoneForwarder.sol

@@ -6,7 +6,6 @@ import {IRouter} from "./interfaces/IRouter.sol";
 import {ITypeAndVersion} from "../shared/interfaces/ITypeAndVersion.sol";
 
 import {OwnerIsCreator} from "../shared/access/OwnerIsCreator.sol";
-import {CallWithExactGas} from "../shared/call/CallWithExactGas.sol";
 
 /// @notice This is an entry point for `write_${chain}` Target capability. It
 /// allows nodes to determine if reports have been processed (successfully or
@@ -91,11 +90,10 @@ contract KeystoneForwarder is OwnerIsCreator, ITypeAndVersion, IRouter {
   uint256 internal constant FORWARDER_METADATA_LENGTH = 45;
   uint256 internal constant SIGNATURE_LENGTH = 65;
 
-  /// @dev The minimum amount of gas to perform the call with exact gas.
-  uint16 internal constant GAS_FOR_CALL_EXACT_CHECK = 5_000;
   /// @dev The gas we require to revert in case of a revert in the call to the
   /// receiver. This is more than enough and does not attempt to be exact.
-  uint256 internal constant GAS_FOR_ROUTER_LOGIC = 40_000;
+  uint256 internal constant REQUIRED_GAS_FOR_ROUTING = 40_000;
+  bytes4 internal constant INSUFFICIENT_GAS_FOR_ROUTING_SIG = 0x0bfecd63;
 
   // ================================================================
   // │                          Router                              │
@@ -110,23 +108,8 @@ contract KeystoneForwarder is OwnerIsCreator, ITypeAndVersion, IRouter {
   }
 
   function removeForwarder(address forwarder) external onlyOwner {
-    s_forwarders[forwarder] = false;
-    emit ForwarderRemoved(forwarder);
-  }
-
-  function reportWithExactGas(
-    uint256 gasLimit,
-    address receiver,
-    bytes calldata metadata,
-    bytes calldata validatedReport
-  ) external returns (bool success) {
-    return
-      CallWithExactGas._callWithExactGas(
-        abi.encodeCall(IReceiver.onReport, (metadata, validatedReport)),
-        receiver,
-        gasLimit,
-        GAS_FOR_CALL_EXACT_CHECK
-      );
+    delete s_forwarders[forwarder];
+    emit IRouter.ForwarderRemoved(forwarder);
   }
 
   function route(
@@ -136,25 +119,49 @@ contract KeystoneForwarder is OwnerIsCreator, ITypeAndVersion, IRouter {
     bytes calldata metadata,
     bytes calldata validatedReport
   ) public returns (bool) {
-    // Calculating the remainder of the gas available after accounting for the
-    // gas needed to handle reverts in the call to the receiver allows us to
-    // avoid passing the gas limit as a parameter to the `route` function.
-    uint256 gasLimit = gasleft() - GAS_FOR_ROUTER_LOGIC;
-
     if (!s_forwarders[msg.sender]) revert UnauthorizedForwarder();
-    if (s_transmissions[transmissionId].transmitter != address(0)) revert AlreadyAttempted(transmissionId);
+
+    TransmissionInfo memory transmission = s_transmissions[transmissionId];
+    uint256 gasLeft = gasleft();
+    if (transmission.transmitter != address(0) && transmission.gasLimit >= gasLeft)
+      revert AlreadyAttempted(transmissionId);
 
     s_transmissions[transmissionId].transmitter = transmitter;
-    s_transmissions[transmissionId].gasLimit = uint88(gasLimit);
-
-    // Making this an external call to be able to catch reverts from the _callWithExactGas function
-    // and avoid having to inline the entire function here.
-    try this.reportWithExactGas(gasLimit, receiver, metadata, validatedReport) returns (bool success) {
-      s_transmissions[transmissionId].success = success;
-      return success;
-    } catch {
-      return false;
+    s_transmissions[transmissionId].gasLimit = uint88(gasLeft);
+
+    bool success;
+    bytes memory payload = abi.encodeCall(IReceiver.onReport, (metadata, validatedReport));
+    assembly {
+      // solidity calls check that a contract actually exists at the destination, so we do the same
+      // Note we do this check prior to measuring gas so REQUIRED_GAS_FOR_ROUTING (our "cushion")
+      // doesn't need to account for it.
+      if iszero(extcodesize(receiver)) {
+        mstore(0x0, 0x00)
+        return(0x0, 0x20)
+      }
+
+      let g := gas()
+      // Compute g -= REQUIRED_GAS_FOR_ROUTING and check for underflow
+      // The gas actually passed to the callee is _min(gasAmount, 63//64*gas available).
+      // We want to ensure that we revert if gasAmount >  63//64*gas available
+      // as we do not want to provide them with less, however that check itself costs
+      // gas. REQUIRED_GAS_FOR_ROUTING ensures we have at least enough gas to be able
+      // to revert if gasAmount >  63//64*gas available.
+      if lt(g, REQUIRED_GAS_FOR_ROUTING) {
+        mstore(0x0, INSUFFICIENT_GAS_FOR_ROUTING_SIG)
+        revert(0x0, 0x4)
+      }
+      g := sub(g, REQUIRED_GAS_FOR_ROUTING)
+
+      // call and return whether we succeeded. ignore return data
+      // call(gas,addr,value,argsOffset,argsLength,retOffset,retLength)
+      success := call(g, receiver, 0, add(payload, 0x20), mload(payload), 0x0, 0x0)
+    }
+
+    if (success) {
+      s_transmissions[transmissionId].success = true;
     }
+    return success;
   }
 
   function getTransmissionId(

contracts/src/v0.8/keystone/interfaces/IRouter.sol

@@ -4,6 +4,9 @@ pragma solidity 0.8.24;
 /// @title IRouter - delivers keystone reports to receiver
 interface IRouter {
   error UnauthorizedForwarder();
+  /// @dev Thrown when the gas limit is insufficient for handling state after
+  /// calling the receiver function.
+  error InsufficientGasForRouting(bytes32 transmissionId);
   error AlreadyAttempted(bytes32 transmissionId);
 
   event ForwarderAdded(address indexed forwarder);

contracts/src/v0.8/keystone/test/KeystoneForwarder_ReportTest.t.sol

@@ -156,7 +156,7 @@ contract KeystoneForwarder_ReportTest is BaseTest {
     vm.expectEmit(address(s_forwarder));
     emit ReportProcessed(address(s_receiver), executionId, reportId, true);
 
-    s_forwarder.report{gas: 200_000}(address(s_receiver), report, reportContext, signatures);
+    s_forwarder.report(address(s_receiver), report, reportContext, signatures);
 
     assertEq(
       s_forwarder.getTransmitter(address(s_receiver), executionId, reportId),
@@ -169,13 +169,35 @@ contract KeystoneForwarder_ReportTest is BaseTest {
       "TransmissionState mismatch"
     );
 
-    assertEq(
+    assertGt(
       s_forwarder.getTransmissionGasLimit(address(s_receiver), executionId, reportId),
-      137_398,
+      100_000,
       "transmission gas limit mismatch"
     );
   }
 
+  // TODO: Add error for insufficient gas
+
+  function test_Report_SuccessfulRetryWithGas() public {
+    s_forwarder.report{gas: 150_000}(address(s_receiver), report, reportContext, signatures);
+
+    // Expect to fail with the receiver running out of gas
+    assertEq(
+      uint8(s_forwarder.getTransmissionState(address(s_receiver), executionId, reportId)),
+      uint8(IRouter.TransmissionState.FAILED)
+    );
+    assertGt(s_forwarder.getTransmissionGasLimit(address(s_receiver), executionId, reportId), 100_000);
+
+    // Should succeed with more gas
+    s_forwarder.report{gas: 300_000}(address(s_receiver), report, reportContext, signatures);
+
+    assertEq(
+      uint8(s_forwarder.getTransmissionState(address(s_receiver), executionId, reportId)),
+      uint8(IRouter.TransmissionState.SUCCEEDED)
+    );
+    assertGt(s_forwarder.getTransmissionGasLimit(address(s_receiver), executionId, reportId), 250_000);
+  }
+
   function test_Report_FailedDeliveryWhenReceiverNotContract() public {
     // Receiver is not a contract
     address receiver = address(404);
@@ -191,6 +213,11 @@ contract KeystoneForwarder_ReportTest is BaseTest {
       uint8(IRouter.TransmissionState.FAILED),
       "TransmissionState mismatch"
     );
+    assertGt(
+      s_forwarder.getTransmissionGasLimit(receiver, executionId, reportId),
+      100_000,
+      "transmission gas limit mismatch"
+    );
   }
 
   function test_Report_FailedDeliveryWhenReceiverInterfaceNotSupported() public {
@@ -208,6 +235,11 @@ contract KeystoneForwarder_ReportTest is BaseTest {
       uint8(IRouter.TransmissionState.FAILED),
       "TransmissionState mismatch"
     );
+    assertGt(
+      s_forwarder.getTransmissionGasLimit(receiver, executionId, reportId),
+      100_000,
+      "transmission gas limit mismatch"
+    );
   }
 
   function test_Report_ConfigVersion() public {
@@ -224,7 +256,7 @@ contract KeystoneForwarder_ReportTest is BaseTest {
     emit ReportProcessed(address(s_receiver), executionId, reportId, true);
 
     vm.prank(TRANSMITTER);
-    s_forwarder.report{gas: 200_000}(address(s_receiver), report, reportContext, signatures);
+    s_forwarder.report(address(s_receiver), report, reportContext, signatures);
 
     // after clear the old version doesn't work anymore
     vm.prank(ADMIN);
@@ -233,7 +265,7 @@ contract KeystoneForwarder_ReportTest is BaseTest {
     uint64 configId = (uint64(DON_ID) << 32) | CONFIG_VERSION;
     vm.expectRevert(abi.encodeWithSelector(KeystoneForwarder.InvalidConfig.selector, configId));
     vm.prank(TRANSMITTER);
-    s_forwarder.report{gas: 200_000}(address(s_receiver), report, reportContext, signatures);
+    s_forwarder.report(address(s_receiver), report, reportContext, signatures);
 
     // but new config does
     bytes32 newExecutionId = hex"6d795f657865637574696f6e5f69640000000000000000000000000000000001";
@@ -257,7 +289,7 @@ contract KeystoneForwarder_ReportTest is BaseTest {
     emit ReportProcessed(address(s_receiver), newExecutionId, reportId, true);
 
     vm.prank(TRANSMITTER);
-    s_forwarder.report{gas: 200_000}(address(s_receiver), newReport, reportContext, newSignatures);
+    s_forwarder.report(address(s_receiver), newReport, reportContext, newSignatures);
 
     // validate transmitter was recorded
     address transmitter = s_forwarder.getTransmitter(address(s_receiver), newExecutionId, reportId);

contracts/src/v0.8/keystone/test/KeystoneRouter_AccessTest.t.sol

@@ -36,6 +36,13 @@ contract KeystoneRouter_SetConfigTest is Test {
     s_router.removeForwarder(FORWARDER);
   }
 
+  function test_RemoveForwarder_Success() public {
+    vm.prank(ADMIN);
+    vm.expectEmit(true, false, false, false);
+    emit IRouter.ForwarderRemoved(FORWARDER);
+    s_router.removeForwarder(FORWARDER);
+  }
+
   function test_Route_RevertWhen_UnauthorizedForwarder() public {
     vm.prank(STRANGER);
     vm.expectRevert(IRouter.UnauthorizedForwarder.selector);

contracts/src/v0.8/keystone/test/mocks/Receiver.sol

@@ -5,10 +5,13 @@ import {IReceiver} from "../../interfaces/IReceiver.sol";
 
 contract Receiver is IReceiver {
   event MessageReceived(bytes metadata, bytes[] mercuryReports);
+  bytes public latestReport;
 
   constructor() {}
 
   function onReport(bytes calldata metadata, bytes calldata rawReport) external {
+    latestReport = rawReport;
+
     // parse actual report
     bytes[] memory mercuryReports = abi.decode(rawReport, (bytes[]));
     emit MessageReceived(metadata, mercuryReports);