add IMessageTransformer interface to CCIP onramp and offramp

contracts/scripts/native_solc_compile_all_ccip

@@ -15,16 +15,24 @@ FOUNDRY_PROJECT_SUFFIX="-compile"
 CONTRACTS_DIR="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; cd ../ && pwd -P )"
 export FOUNDRY_PROFILE="$PROJECT"$FOUNDRY_PROJECT_SUFFIX
 
+# first argument is the contract path, second argument is whether to use the --via-ir flag
 compileContract() {
   local contract
   contract=$(basename "$1")
   echo "Compiling" "$contract"
 
+  local viaIR=""
+  if [[ "$2" =~ "true" ]]; then
+    echo "via-ir flag enabled"
+    viaIR="--via-ir"
+  fi
+
   local command
   command="forge build $CONTRACTS_DIR/src/v0.8/$PROJECT/"$1.sol" \
    --root $CONTRACTS_DIR \
    $(getOptimizations "$contract") \
    --extra-output-files bin abi \
+   $viaIR \
    -o $CONTRACTS_DIR/solc/$PROJECT/$contract"
 
   $command
@@ -81,4 +89,8 @@ compileContract test/helpers/CCIPReaderTester
 # Offchain test encoding utils
 compileContract test/helpers/EncodingUtils
 
+# Message Transformer On/OffRamps
+compileContract offRamp/MessageTransformerOffRamp true
+compileContract onRamp/MessageTransformerOnRamp
+
 

contracts/src/v0.8/ccip/interfaces/IMessageTransformer.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {Internal} from "../libraries/Internal.sol";
+
+/// @notice Interface for plug-in message hook contracts that transform OffRamp & OnRamp messages.
+///         The transformer functions are expected to revert on transform failures.
+interface IMessageTransformer {
+  /// @notice Common error that can be thrown on transform failures and used by consumers
+  /// @param errorReason abi encoded revert reason
+  error MessageTransformError(bytes errorReason);
+
+  /// @notice Transforms the given OffRamp message. Reverts on transform failure
+  /// @param message to transform
+  /// @return transformed message
+  function transformInboundMessage(
+    Internal.Any2EVMRampMessage memory message
+  ) external returns (Internal.Any2EVMRampMessage memory);
+
+  /// @notice Transforms the given OnRamp message. Reverts on transform failure
+  /// @param message to transform
+  /// @return transformed message
+  function transformOutboundMessage(
+    Internal.EVM2AnyRampMessage memory message
+  ) external returns (Internal.EVM2AnyRampMessage memory);
+}

contracts/src/v0.8/ccip/offRamp/MessageTransformerOffRamp.sol

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.24;
+
+import {IMessageTransformer} from "../interfaces/IMessageTransformer.sol";
+import {Internal} from "../libraries/Internal.sol";
+import {OffRamp} from "./OffRamp.sol";
+
+/// @notice OffRamp that uses a message transformer to transform messages before execution
+contract MessageTransformerOffRamp is OffRamp {
+  address internal s_messageTransformer;
+
+  constructor(
+    StaticConfig memory staticConfig,
+    DynamicConfig memory dynamicConfig,
+    SourceChainConfigArgs[] memory sourceChainConfigs,
+    address messageTransformerAddr
+  ) OffRamp(staticConfig, dynamicConfig, sourceChainConfigs) {
+    if (address(messageTransformerAddr) == address(0)) {
+      revert ZeroAddressNotAllowed();
+    }
+    s_messageTransformer = messageTransformerAddr;
+  }
+
+  function getMessageTransformerAddress() external view returns (address) {
+    return s_messageTransformer;
+  }
+
+  function _beforeExecuteSingleMessage(
+    Internal.Any2EVMRampMessage memory message
+  ) internal override returns (Internal.Any2EVMRampMessage memory transformedMessage) {
+    try IMessageTransformer(s_messageTransformer).transformInboundMessage(message) returns (
+      Internal.Any2EVMRampMessage memory m
+    ) {
+      transformedMessage = m;
+    } catch (bytes memory err) {
+      revert IMessageTransformer.MessageTransformError(err);
+    }
+    return transformedMessage;
+  }
+}

contracts/src/v0.8/ccip/offRamp/OffRamp.sol

@@ -552,6 +552,15 @@ contract OffRamp is ITypeAndVersion, MultiOCR3Base {
     return (Internal.MessageExecutionState.SUCCESS, "");
   }
 
+  /// @notice hook for applying custom logic to the input message before executeSingleMessage()
+  /// @param message initial message
+  /// @return transformedMessage modified message
+  function _beforeExecuteSingleMessage(
+    Internal.Any2EVMRampMessage memory message
+  ) internal virtual returns (Internal.Any2EVMRampMessage memory transformedMessage) {
+    return message;
+  }
+
   /// @notice Executes a single message.
   /// @param message The message that will be executed.
   /// @param offchainTokenData Token transfer data to be passed to TokenPool.
@@ -565,6 +574,9 @@ contract OffRamp is ITypeAndVersion, MultiOCR3Base {
     uint32[] calldata tokenGasOverrides
   ) external {
     if (msg.sender != address(this)) revert CanOnlySelfCall();
+
+    message = _beforeExecuteSingleMessage(message);
+
     Client.EVMTokenAmount[] memory destTokenAmounts = new Client.EVMTokenAmount[](0);
     if (message.tokenAmounts.length > 0) {
       destTokenAmounts = _releaseOrMintTokens(

contracts/src/v0.8/ccip/onRamp/MessageTransformerOnRamp.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.24;
+
+import {IMessageTransformer} from "../interfaces/IMessageTransformer.sol";
+import {Internal} from "../libraries/Internal.sol";
+import {OnRamp} from "./OnRamp.sol";
+
+/// @notice OnRamp that uses a message transformer to transform messages from router
+contract MessageTransformerOnRamp is OnRamp {
+  address internal s_messageTransformer;
+
+  error ZeroAddressNotAllowed();
+
+  constructor(
+    StaticConfig memory staticConfig,
+    DynamicConfig memory dynamicConfig,
+    DestChainConfigArgs[] memory destChainConfigs,
+    address messageTransformerAddr
+  ) OnRamp(staticConfig, dynamicConfig, destChainConfigs) {
+    if (address(messageTransformerAddr) == address(0)) {
+      revert ZeroAddressNotAllowed();
+    }
+    s_messageTransformer = messageTransformerAddr;
+  }
+
+  function getMessageTransformerAddress() external view returns (address) {
+    return s_messageTransformer;
+  }
+
+  function _postProcessMessage(
+    Internal.EVM2AnyRampMessage memory message
+  ) internal override returns (Internal.EVM2AnyRampMessage memory transformedMessage) {
+    try IMessageTransformer(s_messageTransformer).transformOutboundMessage(message) returns (
+      Internal.EVM2AnyRampMessage memory m
+    ) {
+      transformedMessage = m;
+    } catch (bytes memory err) {
+      revert IMessageTransformer.MessageTransformError(err);
+    }
+    return transformedMessage;
+  }
+}

contracts/src/v0.8/ccip/onRamp/OnRamp.sol

@@ -239,6 +239,8 @@ contract OnRamp is IEVM2AnyOnRampClient, ITypeAndVersion, Ownable2StepMsgSender
       newMessage.tokenAmounts[i].destExecData = destExecDataPerToken[i];
     }
 
+    newMessage = _postProcessMessage(newMessage);
+
     // Hash only after all fields have been set.
     newMessage.header.messageId = Internal._hash(
       newMessage,
@@ -257,6 +259,15 @@ contract OnRamp is IEVM2AnyOnRampClient, ITypeAndVersion, Ownable2StepMsgSender
     return newMessage.header.messageId;
   }
 
+  /// @notice hook for applying custom logic to the message from router
+  /// @param message router message
+  /// @return transformedMessage modified message
+  function _postProcessMessage(
+    Internal.EVM2AnyRampMessage memory message
+  ) internal virtual returns (Internal.EVM2AnyRampMessage memory transformedMessage) {
+    return message;
+  }
+
   /// @notice Uses a pool to lock or burn a token.
   /// @param tokenAndAmount Token address and amount to lock or burn.
   /// @param destChainSelector Target destination chain selector of the message.

contracts/src/v0.8/ccip/test/helpers/MessageTransformerHelper.sol

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity 0.8.26;
+
+import {IMessageTransformer} from "../../interfaces/IMessageTransformer.sol";
+import {Internal} from "../../libraries/Internal.sol";
+
+// This helper is used to test the On/OffRamps
+contract MessageTransformerHelper is IMessageTransformer {
+  error UnknownChain();
+
+  bool public s_shouldRevert;
+
+  function setShouldRevert(
+    bool _shouldRevert
+  ) external {
+    s_shouldRevert = _shouldRevert;
+  }
+
+  /// @inheritdoc IMessageTransformer
+  function transformInboundMessage(
+    Internal.Any2EVMRampMessage memory message
+  ) public view returns (Internal.Any2EVMRampMessage memory) {
+    if (s_shouldRevert) {
+      revert UnknownChain();
+    }
+    return message;
+  }
+
+  /// @inheritdoc IMessageTransformer
+  function transformOutboundMessage(
+    Internal.EVM2AnyRampMessage memory message
+  ) public view returns (Internal.EVM2AnyRampMessage memory) {
+    if (s_shouldRevert) {
+      revert UnknownChain();
+    }
+    return message;
+  }
+}

contracts/src/v0.8/ccip/test/offRamp/OffRamp/MessageTransformerOffRamp.executeSingleMessage.t.sol

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.24;
+
+import {Router} from "../../../Router.sol";
+import {IMessageTransformer} from "../../../interfaces/IMessageTransformer.sol";
+import {Internal} from "../../../libraries/Internal.sol";
+import {MessageTransformerOffRamp} from "../../../offRamp/MessageTransformerOffRamp.sol";
+import {OffRamp} from "../../../offRamp/OffRamp.sol";
+
+import {MessageTransformerHelper} from "../../helpers/MessageTransformerHelper.sol";
+import {OffRampSetup} from "./OffRampSetup.t.sol";
+
+contract MessageTransformerOffRamp_executeSingleMessage is OffRampSetup {
+  MessageTransformerOffRamp internal s_messageTransformerOffRamp;
+
+  function setUp() public virtual override {
+    super.setUp();
+    s_messageTransformerOffRamp = new MessageTransformerOffRamp(
+      s_offRamp.getStaticConfig(),
+      s_offRamp.getDynamicConfig(),
+      new OffRamp.SourceChainConfigArgs[](0),
+      address(s_inboundMessageTransformer)
+    );
+
+    OffRamp.SourceChainConfigArgs[] memory sourceChainConfigs = new OffRamp.SourceChainConfigArgs[](1);
+    sourceChainConfigs[0] = OffRamp.SourceChainConfigArgs({
+      router: s_destRouter,
+      sourceChainSelector: SOURCE_CHAIN_SELECTOR_1,
+      onRamp: ON_RAMP_ADDRESS_1,
+      isEnabled: true
+    });
+    s_messageTransformerOffRamp.applySourceChainConfigUpdates(sourceChainConfigs);
+
+    Router.OnRamp[] memory onRampUpdates = new Router.OnRamp[](0);
+    Router.OffRamp[] memory offRampUpdates = new Router.OffRamp[](2 * sourceChainConfigs.length);
+
+    for (uint256 i = 0; i < sourceChainConfigs.length; ++i) {
+      uint64 sourceChainSelector = sourceChainConfigs[i].sourceChainSelector;
+
+      offRampUpdates[2 * i] =
+        Router.OffRamp({sourceChainSelector: sourceChainSelector, offRamp: address(s_messageTransformerOffRamp)});
+      offRampUpdates[2 * i + 1] = Router.OffRamp({
+        sourceChainSelector: sourceChainSelector,
+        offRamp: s_inboundNonceManager.getPreviousRamps(sourceChainSelector).prevOffRamp
+      });
+    }
+
+    s_destRouter.applyRampUpdates(onRampUpdates, new Router.OffRamp[](0), offRampUpdates);
+  }
+
+  function test_executeSingleMessage_WithMessageTransformer() public {
+    vm.stopPrank();
+    vm.startPrank(address(s_messageTransformerOffRamp));
+    Internal.Any2EVMRampMessage memory message =
+      _generateAny2EVMMessageNoTokens(SOURCE_CHAIN_SELECTOR_1, ON_RAMP_ADDRESS_1, 1);
+    s_messageTransformerOffRamp.executeSingleMessage(message, new bytes[](message.tokenAmounts.length), new uint32[](0));
+  }
+
+  function test_executeSingleMessage_WithMessageTransformer_RevertWhen_UnknownChain() public {
+    vm.stopPrank();
+    vm.startPrank(address(s_messageTransformerOffRamp));
+    Internal.Any2EVMRampMessage memory message =
+      _generateAny2EVMMessageNoTokens(SOURCE_CHAIN_SELECTOR_1, ON_RAMP_ADDRESS_1, 1);
+    // Fail with any error (UnknownChain in this case) to check if OffRamp wraps the error with MessageTransformError during the revert
+    s_inboundMessageTransformer.setShouldRevert(true);
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        IMessageTransformer.MessageTransformError.selector,
+        abi.encodeWithSelector(MessageTransformerHelper.UnknownChain.selector)
+      )
+    );
+    s_messageTransformerOffRamp.executeSingleMessage(message, new bytes[](message.tokenAmounts.length), new uint32[](0));
+  }
+}

contracts/src/v0.8/ccip/test/offRamp/OffRamp/OffRampSetup.t.sol

@@ -14,6 +14,7 @@ import {OffRamp} from "../../../offRamp/OffRamp.sol";
 import {FeeQuoterSetup} from "../../feeQuoter/FeeQuoterSetup.t.sol";
 import {MaybeRevertingBurnMintTokenPool} from "../../helpers/MaybeRevertingBurnMintTokenPool.sol";
 import {MessageInterceptorHelper} from "../../helpers/MessageInterceptorHelper.sol";
+import {MessageTransformerHelper} from "../../helpers/MessageTransformerHelper.sol";
 import {OffRampHelper} from "../../helpers/OffRampHelper.sol";
 import {MaybeRevertMessageReceiver} from "../../helpers/receivers/MaybeRevertMessageReceiver.sol";
 import {MultiOCR3BaseSetup} from "../../ocr/MultiOCR3Base/MultiOCR3BaseSetup.t.sol";
@@ -38,6 +39,7 @@ contract OffRampSetup is FeeQuoterSetup, MultiOCR3BaseSetup {
 
   OffRampHelper internal s_offRamp;
   MessageInterceptorHelper internal s_inboundMessageInterceptor;
+  MessageTransformerHelper internal s_inboundMessageTransformer;
   NonceManager internal s_inboundNonceManager;
 
   bytes32 internal s_configDigestExec;
@@ -53,6 +55,7 @@ contract OffRampSetup is FeeQuoterSetup, MultiOCR3BaseSetup {
     MultiOCR3BaseSetup.setUp();
 
     s_inboundMessageInterceptor = new MessageInterceptorHelper();
+    s_inboundMessageTransformer = new MessageTransformerHelper();
     s_receiver = new MaybeRevertMessageReceiver(false);
     s_secondary_receiver = new MaybeRevertMessageReceiver(false);
     s_reverting_receiver = new MaybeRevertMessageReceiver(true);