change ChainId to Word16, reorder TokenMessage fields

Author: chessai

src/Crypto/Hash/HyperlaneNatives.hs

@@ -8,6 +8,7 @@
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
 
 -- | Implementation of Hyperlane natives.
 module Crypto.Hash.HyperlaneNatives
@@ -48,7 +49,7 @@ import Data.Text qualified as Text
 import Data.Text.Encoding qualified as Text
 import Data.Text.Read qualified as Text
 import Data.WideWord.Word256 (Word256(..))
-import Data.Word (Word8, Word32)
+import Data.Word (Word8, Word16, Word32)
 import Ethereum.Misc (keccak256, _getKeccak256Hash, _getBytesN)
 import Pact.JSON.Decode qualified as J
 import Pact.Types.Exp (Literal(..))
@@ -158,9 +159,9 @@ data HyperlaneMessage = HyperlaneMessage
   deriving stock (Eq, Show)
 
 data TokenMessageERC20 = TokenMessageERC20
-  { tmRecipient :: ByteString -- variable
-  , tmAmount :: Word256 -- uint256
-  , tmChainId :: Word256 -- uint256
+  { tmAmount :: Word256 -- uint256
+  , tmChainId :: Word16 -- uint16
+  , tmRecipient :: ByteString -- variable
   }
   deriving stock (Eq, Show)
 
@@ -170,7 +171,7 @@ data TokenMessageERC20 = TokenMessageERC20
 
 packHyperlaneMessage :: HyperlaneMessage -> Builder
 packHyperlaneMessage (HyperlaneMessage{..}) =
-     BB.word8 hmVersion
+  BB.word8 hmVersion
   <> BB.word32BE hmNonce
   <> BB.word32BE hmOriginDomain
   <> BB.byteString hmSender
@@ -217,31 +218,15 @@ packHyperlaneMessage (HyperlaneMessage{..}) =
 -- 5D7D
 packTokenMessageERC20 :: TokenMessageERC20 -> Builder
 packTokenMessageERC20 t =
-  word256BE 96
-  <> word256BE (tmAmount t)
-  <> word256BE (tmChainId t)
-  <> word256BE recipientSize
-  <> BB.byteString recipient
-  where
-    (recipient, recipientSize) = padRight (tmRecipient t)
+  word256BE (tmAmount t)
+  <> BB.word16BE (tmChainId t)
+  <> BB.byteString (tmRecipient t)
 
 unpackTokenMessageERC20 :: Get TokenMessageERC20
 unpackTokenMessageERC20 = do
-  firstOffset <- getWord256BE
-  unless (firstOffset == 96) $ do
-    fail $ "TokenMessage firstOffset expected 96, found " ++ show firstOffset
-
   tmAmount <- getWord256BE
-  tmChainId <- getWord256BE
-
-  recipientSize <- getWord256BE
-  tmRecipient <- do
-    let size = fromIntegral @Word256 @Int recipientSize
-    recipient <- BS.take size
-      <$> Bin.getByteString (fromIntegral @Word256 @Int (recipientSize + restSize recipientSize))
-    if BS.length recipient < size
-    then fail "TokenMessage recipient was smaller than expected"
-    else pure recipient
+  tmChainId <- Bin.getWord16be
+  tmRecipient <- BL.toStrict <$> Bin.getRemainingLazyByteString
 
   pure $ TokenMessageERC20 {..}
 
@@ -274,11 +259,13 @@ decodeBase64AndValidate key expected s = do
 
   return decoded
 
-parseChainId :: Text -> Either HyperlaneError Word256
+parseChainId :: forall a. (a ~ Word16) => Text -> Either HyperlaneError a
 parseChainId s = do
-  cid <- first (HyperlaneErrorInvalidChainId . Text.pack) $ Text.decimal s
-  unless (fst cid >= 0) $ throwError $ HyperlaneErrorInvalidChainId "can't be negative"
-  return $ fst cid
+  (cid, _) <- first (HyperlaneErrorInvalidChainId . Text.pack) $ Text.decimal s
+  unless (cid >= minBound && cid <= maxBound) $ do
+    throwError $ HyperlaneErrorInvalidChainId $ Text.pack $
+      "ChainId must be in [" <> show @a minBound <> ", " <> show @a maxBound <> "]"
+  pure cid
 
 ------------------------------------------------------
 --      Hyperlane Message Pact Object Decoding      --
@@ -370,8 +357,8 @@ getWord256BE = do
 --
 -- > padRight "hello world"
 -- ("hello world\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL\NUL",11)
-padRight :: ByteString -> (ByteString, Word256)
-padRight s =
+_padRight :: ByteString -> (ByteString, Word256)
+_padRight s =
   let
     size = BS.length s
     missingZeroes = restSize size