enhance BAM payload parsing

wormhole-foundation · Oct 2, 2023 · a537b45 · a537b45
1 parent 224d582
commit a537b45
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 37 deletions.
diff --git a/core/definitions/__tests__/bamVaa.ts b/core/definitions/__tests__/bamVaa.ts
@@ -1,21 +1,20 @@
+import { uint8ArrayToHexByteString } from "@wormhole-foundation/sdk-base/src/utils";
 import "../src/payloads/bam";
 import { deserialize, deserializePayload, serializePayload } from "../src/vaa";
 
-const payloadLiteral = "BAMessage";
-const magicByte = 0xBB;
-const baseGoerliId = 30;
-const goerliBamAppAddress = "44fbfee0af8efa9e580760844f6159a8e2124b53";
+const payloadLiteral = "BamMessage";
+const goerliBamAppAddress = "0x44fbfee0af8efa9e580760844f6159a8e2124b53";
 // these two contracts are deployed at the same address, made 2 variables to avoid confusion
 const baseGoerliBamAppAddress = goerliBamAppAddress;
 
 const cases = [
   {
     vaa: "AQAAAAABAIoYgY9KWhpDVdehi/6jlPnKwiZPMXrE103Pmqm+d0inIdgSnSYcUjExiGB18t2UKDjhbCo+alECY4jXJ1t2iXoAZRML0AAAAAAAAgAAAAAAAAAAAAAAAHmVUoxxQayZCePil5D8T4JCGCRPAAAAAAAAAADKuwAAAAAAAAAAAAAAHgAURPv+4K+O+p5YB2CET2FZqOISS1MAFET7/uCvjvqeWAdghE9hWajiEktTAGAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFaGVsbG8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=",
-    targetChain: baseGoerliId,
+    targetChain: "Base",
     targetAddress: baseGoerliBamAppAddress,
     senderAddress: goerliBamAppAddress,
-    contents: "0000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000568656c6c6f000000000000000000000000000000000000000000000000000000",
-    payload: "bb00000000000000000000001e001444fbfee0af8efa9e580760844f6159a8e2124b53001444fbfee0af8efa9e580760844f6159a8e2124b5300600000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000568656c6c6f000000000000000000000000000000000000000000000000000000",
+    contents: "0x0000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000568656c6c6f000000000000000000000000000000000000000000000000000000",
+    payload: "0xbb00000000000000000000001e001444fbfee0af8efa9e580760844f6159a8e2124b53001444fbfee0af8efa9e580760844f6159a8e2124b5300600000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000568656c6c6f000000000000000000000000000000000000000000000000000000",
   }
 ];
 
@@ -25,16 +24,14 @@ describe("BAM VAA tests", function () {
     for (const testCase of cases) {
       const vaaBytes = Buffer.from(testCase.vaa, "base64");
 
-      const parsed = deserialize("Uint8Array", new Uint8Array(vaaBytes));
-      const x = deserializePayload(payloadLiteral, parsed.payload);
+      const {payload} = deserialize(payloadLiteral, vaaBytes);
 
-      const targetAddress = Buffer.from(x.targetAddress).toString("hex");
-      const senderAddress = Buffer.from(x.senderAddress).toString("hex");
-      const parsedContents = Buffer.from(x.contents).toString("hex");
+      const targetAddress = uint8ArrayToHexByteString(payload.targetAddress);
+      const senderAddress = uint8ArrayToHexByteString(payload.senderAddress);
+      const parsedContents = uint8ArrayToHexByteString(payload.contents);
 
-      expect(x).toBeTruthy();
-      expect(x.magicByte).toEqual(magicByte);
-      expect(x.targetChain).toEqual(testCase.targetChain);
+      expect(payload).toBeTruthy();
+      expect(payload.targetChain).toEqual(testCase.targetChain);
       expect(targetAddress).toEqual(testCase.targetAddress);
       expect(senderAddress).toEqual(testCase.senderAddress);
       expect(parsedContents).toEqual(testCase.contents);
@@ -45,11 +42,10 @@ describe("BAM VAA tests", function () {
     for (const testCase of cases) {
       const vaaBytes = Buffer.from(testCase.vaa, "base64");
 
-      const parsed = deserialize("Uint8Array", new Uint8Array(vaaBytes));
+      const {payload} = deserialize(payloadLiteral, vaaBytes);
 
-      const x = deserializePayload(payloadLiteral, parsed.payload);
-      const serialized = serializePayload(payloadLiteral, x);
-      const serializedHex = Buffer.from(serialized).toString("hex");
+      const serialized = serializePayload(payloadLiteral, payload);
+      const serializedHex = uint8ArrayToHexByteString(serialized);
       expect(serializedHex).toEqual(testCase.payload);
     }
   })

diff --git a/core/definitions/src/payloads/bam.ts b/core/definitions/src/payloads/bam.ts
@@ -1,26 +1,55 @@
-import { Layout, LengthPrefixedBytesLayoutItem, ShallowMapping } from "@wormhole-foundation/sdk-base";
-import { chainItemBase, sequenceItem } from "../layout-items";
+import { Layout, UintLayoutItem, LengthPrefixedBytesLayoutItem, ShallowMapping } from "@wormhole-foundation/sdk-base";
+import { chainItem, amountItem } from "../layout-items";
 import { registerPayloadType } from "../vaa";
 
-const bamCommonLayout: Layout = [
-  { name: "magicByte", binary: 'uint', size: 1 },
-  { name: "version", binary: 'uint', size: 1 },
-  { name: "type", binary: 'uint', size: 1 },
-]
+const bamAddressItem = {
+  binary: "bytes",
+  lengthSize: 2
+} as const satisfies Omit<LengthPrefixedBytesLayoutItem, "name">;
 
-const bamLayout0: Layout = [
-  { name: "index", ...sequenceItem },
-  { name: "targetChain", ...chainItemBase },
-  { name: "targetAddress", binary: 'bytes', lengthSize: 2 } as LengthPrefixedBytesLayoutItem,
-  { name: "senderAddress", binary: 'bytes', lengthSize: 2 } as LengthPrefixedBytesLayoutItem,
-  { name: "contents", binary: 'bytes', lengthSize: 2 } as LengthPrefixedBytesLayoutItem,
-];
+const customOrEmpty = (custom: any) => custom ? { custom } : {};
+
+const messageLayout = <
+  const T extends number,
+  const C extends Pick<LengthPrefixedBytesLayoutItem, "custom">,
+>(type: T, customContents?: C) => [
+  { name: "magicByte", binary: "uint", size: 1, custom: 0xbb, omit: true },
+  { name: "version", binary: "uint", size: 1, custom: 0, omit: true },
+  { name: "type", binary: "uint", size: 1, custom: type, omit: true },
+  { name: "index", binary: "uint", size: 8 },
+  { name: "targetChain", ...chainItem() },
+  { name: "targetAddress", ...bamAddressItem },
+  { name: "senderAddress", ...bamAddressItem },
+  { name: "contents", binary: "bytes", lengthSize: 2, ...customOrEmpty(customContents) },
+] as const satisfies Layout;
+
+const tokenMessageLayout = <
+  const C extends Pick<LengthPrefixedBytesLayoutItem, "custom">,
+  const B extends Pick<UintLayoutItem, "custom">,
+  const A extends Pick<UintLayoutItem, "custom">,
+>(
+  custom?: { contents?: C, bridge?: B, assetIdentifier?: A },
+) => [
+  ...messageLayout(1, custom?.contents),
+  { name: "bridge", binary: "uint", size: 1, ...customOrEmpty(custom?.bridge) },
+  { name: "assetIdentifier", binary: "uint", size: 2, ...customOrEmpty(custom?.assetIdentifier) },
+  { name: "amount", ...amountItem },
+] as const satisfies Layout;
+
+const extendedMessageLayout = <
+  const C extends Pick<LengthPrefixedBytesLayoutItem, "custom">,
+  const R extends Pick<LengthPrefixedBytesLayoutItem, "custom">,
+>(
+  custom?: { contents?: C, relaySignal?: R },
+) => [
+  ...messageLayout(2, custom?.contents),
+  { name: "relaySignal", binary: "bytes", lengthSize: 2, ...customOrEmpty(custom?.relaySignal) },
+] as const satisfies Layout;
 
 export const bamPayloads = [
-  [
-    "BAMessage",
-    [ ...bamCommonLayout, ...bamLayout0 ],
-  ],
+  [ "BamMessage", messageLayout(0) ],
+  [ "BamTokenMessage", tokenMessageLayout() ],
+  [ "BamExtendedMessage", extendedMessageLayout() ],
 ] as const satisfies readonly (readonly [string, Layout])[];
 
 // factory registration: