From 17090608b23d466cf875e40b4fe81369c7d8779d Mon Sep 17 00:00:00 2001
From: A5 Pickle <a5-pickle@users.noreply.github.com>
Date: Thu, 11 Jan 2024 17:08:47 -0600
Subject: [PATCH] evm: add custom payload ID support

---
 evm/src/libraries/WormholeCctpMessages.sol | 92 +++++++++++++++++++++-
 1 file changed, 88 insertions(+), 4 deletions(-)

diff --git a/evm/src/libraries/WormholeCctpMessages.sol b/evm/src/libraries/WormholeCctpMessages.sol
index ebd5a8a..95980fc 100644
--- a/evm/src/libraries/WormholeCctpMessages.sol
+++ b/evm/src/libraries/WormholeCctpMessages.sol
@@ -30,6 +30,9 @@ library WormholeCctpMessages {
     error InvalidMessage();
     error UnexpectedMessageLength(uint256, uint256);
 
+    /**
+     * @dev NOTE: This method encodes the Wormhole message payload assuming the payload ID == 1.
+     */
     function encodeDeposit(
         address token,
         uint256 amount,
@@ -42,6 +45,7 @@ library WormholeCctpMessages {
     ) internal pure returns (bytes memory encoded) {
         encoded = encodeDeposit(
             token.toUniversalAddress(),
+            DEPOSIT,
             amount,
             sourceCctpDomain,
             targetCctpDomain,
@@ -52,6 +56,9 @@ library WormholeCctpMessages {
         );
     }
 
+    /**
+     * @dev NOTE: This method encodes the Wormhole message payload assuming the payload ID == 1.
+     */
     function encodeDeposit(
         bytes32 universalTokenAddress,
         uint256 amount,
@@ -61,6 +68,54 @@ library WormholeCctpMessages {
         bytes32 burnSource,
         bytes32 mintRecipient,
         bytes memory payload
+    ) internal pure returns (bytes memory encoded) {
+        encoded = encodeDeposit(
+            universalTokenAddress,
+            DEPOSIT,
+            amount,
+            sourceCctpDomain,
+            targetCctpDomain,
+            cctpNonce,
+            burnSource,
+            mintRecipient,
+            payload
+        );
+    }
+
+    function encodeDeposit(
+        address token,
+        uint8 payloadId,
+        uint256 amount,
+        uint32 sourceCctpDomain,
+        uint32 targetCctpDomain,
+        uint64 cctpNonce,
+        bytes32 burnSource,
+        bytes32 mintRecipient,
+        bytes memory payload
+    ) internal pure returns (bytes memory encoded) {
+        encoded = encodeDeposit(
+            token.toUniversalAddress(),
+            payloadId,
+            amount,
+            sourceCctpDomain,
+            targetCctpDomain,
+            cctpNonce,
+            burnSource,
+            mintRecipient,
+            payload
+        );
+    }
+
+    function encodeDeposit(
+        bytes32 universalTokenAddress,
+        uint8 payloadId,
+        uint256 amount,
+        uint32 sourceCctpDomain,
+        uint32 targetCctpDomain,
+        uint64 cctpNonce,
+        bytes32 burnSource,
+        bytes32 mintRecipient,
+        bytes memory payload
     ) internal pure returns (bytes memory encoded) {
         uint256 payloadLen = payload.length;
         if (payloadLen == 0) {
@@ -70,7 +125,7 @@ library WormholeCctpMessages {
         }
 
         encoded = abi.encodePacked(
-            DEPOSIT,
+            payloadId,
             universalTokenAddress,
             amount,
             sourceCctpDomain,
@@ -83,6 +138,9 @@ library WormholeCctpMessages {
         );
     }
 
+    /**
+     * @dev NOTE: This method decodes the VAA payload assuming the payload ID == 1.
+     */
     function decodeDeposit(IWormhole.VM memory vaa)
         internal
         pure
@@ -106,10 +164,36 @@ library WormholeCctpMessages {
             burnSource,
             mintRecipient,
             payload
-        ) = decodeDeposit(vaa, true);
+        ) = decodeDeposit(vaa, DEPOSIT, true);
+    }
+
+    function decodeDeposit(IWormhole.VM memory vaa, uint8 payloadId)
+        internal
+        pure
+        returns (
+            bytes32 token,
+            uint256 amount,
+            uint32 sourceCctpDomain,
+            uint32 targetCctpDomain,
+            uint64 cctpNonce,
+            bytes32 burnSource,
+            bytes32 mintRecipient,
+            bytes memory payload
+        )
+    {
+        (
+            token,
+            amount,
+            sourceCctpDomain,
+            targetCctpDomain,
+            cctpNonce,
+            burnSource,
+            mintRecipient,
+            payload
+        ) = decodeDeposit(vaa, payloadId, true);
     }
 
-    function decodeDeposit(IWormhole.VM memory vaa, bool revertCustomErrors)
+    function decodeDeposit(IWormhole.VM memory vaa, uint8 payloadId, bool revertCustomErrors)
         internal
         pure
         returns (
@@ -124,7 +208,7 @@ library WormholeCctpMessages {
         )
     {
         bytes memory encoded = vaa.payload;
-        uint256 offset = _checkPayloadId(encoded, 0, DEPOSIT, revertCustomErrors);
+        uint256 offset = _checkPayloadId(encoded, 0, payloadId, revertCustomErrors);
 
         (token, offset) = encoded.asBytes32Unchecked(offset);
         (amount, offset) = encoded.asUint256Unchecked(offset);