Header becomes class

xc-admin/packages/xc-admin-common/src/__tests__/GovernancePayload.test.ts

@@ -1,66 +1,61 @@
 import { ChainName } from "@certusone/wormhole-sdk";
 import { PACKET_DATA_SIZE, PublicKey, SystemProgram } from "@solana/web3.js";
-import { ActionName, decodeHeader, encodeHeader, ExecutePostedVaa } from "..";
+import { ActionName, PythGovernanceHeader, ExecutePostedVaa } from "..";
 
 test("GovernancePayload ser/de", (done) => {
   jest.setTimeout(60000);
 
   // Valid header 1
-  let expectedGovernanceHeader = {
-    targetChainId: "pythnet" as ChainName,
-    action: "ExecutePostedVaa" as ActionName,
-  };
-  let buffer = Buffer.alloc(PACKET_DATA_SIZE);
-  let span = encodeHeader(expectedGovernanceHeader, buffer);
+  let expectedGovernanceHeader = new PythGovernanceHeader(
+    "pythnet",
+    "ExecutePostedVaa"
+  );
+  let buffer = expectedGovernanceHeader.encode();
   expect(
-    buffer.subarray(0, span).equals(Buffer.from([80, 84, 71, 77, 0, 0, 0, 26]))
+    buffer.equals(Buffer.from([80, 84, 71, 77, 0, 0, 0, 26]))
   ).toBeTruthy();
-
-  let governanceHeader = decodeHeader(buffer.subarray(0, span));
-  expect(governanceHeader?.targetChainId).toBe("pythnet");
-  expect(governanceHeader?.action).toBe("ExecutePostedVaa");
+  let governanceHeader = PythGovernanceHeader.decode(buffer);
+  expect(governanceHeader.targetChainId).toBe("pythnet");
+  expect(governanceHeader.action).toBe("ExecutePostedVaa");
 
   // Valid header 2
-  expectedGovernanceHeader = {
-    targetChainId: "unset" as ChainName,
-    action: "ExecutePostedVaa" as ActionName,
-  };
-  buffer = Buffer.alloc(PACKET_DATA_SIZE);
-  span = encodeHeader(expectedGovernanceHeader, buffer);
-  expect(
-    buffer.subarray(0, span).equals(Buffer.from([80, 84, 71, 77, 0, 0, 0, 0]))
-  ).toBeTruthy();
-  governanceHeader = decodeHeader(buffer.subarray(0, span));
+  expectedGovernanceHeader = new PythGovernanceHeader(
+    "unset",
+    "ExecutePostedVaa"
+  );
+  buffer = expectedGovernanceHeader.encode();
+  expect(buffer.equals(Buffer.from([80, 84, 71, 77, 0, 0, 0, 0]))).toBeTruthy();
+  governanceHeader = PythGovernanceHeader.decode(buffer);
   expect(governanceHeader?.targetChainId).toBe("unset");
   expect(governanceHeader?.action).toBe("ExecutePostedVaa");
 
   // Valid header 3
-  expectedGovernanceHeader = {
-    targetChainId: "solana" as ChainName,
-    action: "SetFee" as ActionName,
-  };
-  buffer = Buffer.alloc(PACKET_DATA_SIZE);
-  span = encodeHeader(expectedGovernanceHeader, buffer);
-  expect(
-    buffer.subarray(0, span).equals(Buffer.from([80, 84, 71, 77, 1, 3, 0, 1]))
-  ).toBeTruthy();
-  governanceHeader = decodeHeader(buffer.subarray(0, span));
+  expectedGovernanceHeader = new PythGovernanceHeader("solana", "SetFee");
+  buffer = expectedGovernanceHeader.encode();
+  expect(buffer.equals(Buffer.from([80, 84, 71, 77, 1, 3, 0, 1]))).toBeTruthy();
+  governanceHeader = PythGovernanceHeader.decode(buffer);
   expect(governanceHeader?.targetChainId).toBe("solana");
   expect(governanceHeader?.action).toBe("SetFee");
 
   // Wrong magic number
   expect(() =>
-    decodeHeader(Buffer.from([0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0]))
+    PythGovernanceHeader.decode(
+      Buffer.from([0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0])
+    )
   ).toThrow("Wrong magic number");
 
   // Wrong chain
   expect(() =>
-    decodeHeader(Buffer.from([80, 84, 71, 77, 0, 0, 255, 255, 0, 0, 0, 0]))
+    PythGovernanceHeader.decode(
+      Buffer.from([80, 84, 71, 77, 0, 0, 255, 255, 0, 0, 0, 0])
+    )
   ).toThrow("Chain Id not found");
 
   // Wrong module/action combination
   expect(() =>
-    decodeHeader(Buffer.from([80, 84, 71, 77, 0, 1, 0, 26, 0, 0, 0, 0]))
+    PythGovernanceHeader.decode(
+      Buffer.from([80, 84, 71, 77, 0, 1, 0, 26, 0, 0, 0, 0])
+    )
   ).toThrow("Invalid header, action doesn't match module");
 
   // Decode executePostVaa with empty instructions

xc-admin/packages/xc-admin-common/src/governance_payload/ExecutePostedVaa.ts

@@ -1,11 +1,6 @@
 import { ChainId, ChainName } from "@certusone/wormhole-sdk";
 import * as BufferLayout from "@solana/buffer-layout";
-import {
-  encodeHeader,
-  governanceHeaderLayout,
-  PythGovernanceAction,
-  verifyHeader,
-} from ".";
+import { PythGovernanceAction, PythGovernanceHeader } from ".";
 import { Layout } from "@solana/buffer-layout";
 import {
   AccountMeta,
@@ -73,7 +68,7 @@ export const executePostedVaaLayout: BufferLayout.Structure<
     instructions: InstructionData[];
   }>
 > = BufferLayout.struct([
-  governanceHeaderLayout(),
+  PythGovernanceHeader.layout,
   new Vector<InstructionData>(instructionDataLayout, "instructions"),
 ]);
 
@@ -92,8 +87,7 @@ export class ExecutePostedVaa implements PythGovernanceAction {
   /** Decode ExecutePostedVaaArgs */
   static decode(data: Buffer): ExecutePostedVaa {
     let deserialized = executePostedVaaLayout.decode(data);
-
-    let header = verifyHeader(deserialized.header);
+    let header = PythGovernanceHeader.verify(deserialized.header);
 
     let instructions: TransactionInstruction[] = deserialized.instructions.map(
       (ix) => {
@@ -114,12 +108,13 @@ export class ExecutePostedVaa implements PythGovernanceAction {
 
   /** Encode ExecutePostedVaaArgs */
   encode(): Buffer {
-    // PACKET_DATA_SIZE is the maximum transaction size of Solana, so our serialized payload will never be bigger than that
+    const headerBuffer = new PythGovernanceHeader(
+      this.targetChainId,
+      "ExecutePostedVaa"
+    ).encode();
+
+    // The code will crash if the payload is actually bigger than PACKET_DATA_SIZE. But PACKET_DATA_SIZE is the maximum transaction size of Solana, so our serialized payload should never be bigger than this anyway
     const buffer = Buffer.alloc(PACKET_DATA_SIZE);
-    const offset = encodeHeader(
-      { action: "ExecutePostedVaa", targetChainId: this.targetChainId },
-      buffer
-    );
     let instructions: InstructionData[] = this.instructions.map((ix) => {
       let programId = ix.programId.toBytes();
       let accounts: AccountMetadata[] = ix.keys.map((acc) => {
@@ -133,13 +128,10 @@ export class ExecutePostedVaa implements PythGovernanceAction {
       return { programId, accounts, data };
     });
 
-    const span =
-      offset +
-      new Vector<InstructionData>(instructionDataLayout, "instructions").encode(
-        instructions,
-        buffer,
-        offset
-      );
-    return buffer.subarray(0, span);
+    const span = new Vector<InstructionData>(
+      instructionDataLayout,
+      "instructions"
+    ).encode(instructions, buffer, PythGovernanceHeader.span);
+    return Buffer.concat([headerBuffer, buffer.subarray(0, span)]);
   }
 }

xc-admin/packages/xc-admin-common/src/governance_payload/index.ts

@@ -5,6 +5,7 @@ import {
   toChainName,
 } from "@certusone/wormhole-sdk";
 import * as BufferLayout from "@solana/buffer-layout";
+import { PACKET_DATA_SIZE } from "@solana/web3.js";
 import { ExecutePostedVaa } from "./ExecutePostedVaa";
 
 export interface PythGovernanceAction {
@@ -52,24 +53,21 @@ export declare type ActionName =
   | keyof typeof ExecutorAction
   | keyof typeof TargetAction;
 
-export type PythGovernanceHeader = {
-  targetChainId: ChainName;
-  action: ActionName;
-};
-
 export const MAGIC_NUMBER = 0x4d475450;
 export const MODULE_EXECUTOR = 0;
 export const MODULE_TARGET = 1;
 
-export function governanceHeaderLayout(): BufferLayout.Structure<
-  Readonly<{
-    magicNumber: number;
-    module: number;
-    action: number;
-    chain: ChainId;
-  }>
-> {
-  return BufferLayout.struct(
+export class PythGovernanceHeader {
+  readonly targetChainId: ChainName;
+  readonly action: ActionName;
+  static layout: BufferLayout.Structure<
+    Readonly<{
+      magicNumber: number;
+      module: number;
+      action: number;
+      chain: ChainId;
+    }>
+  > = BufferLayout.struct(
     [
       BufferLayout.u32("magicNumber"),
       BufferLayout.u8("module"),
@@ -78,66 +76,71 @@ export function governanceHeaderLayout(): BufferLayout.Structure<
     ],
     "header"
   );
-}
+  static span: 8;
 
-/** Decode Pyth Governance Header and return undefined if the header is invalid */
-export function decodeHeader(data: Buffer): PythGovernanceHeader {
-  let deserialized = governanceHeaderLayout().decode(data);
-  return verifyHeader(deserialized);
-}
-
-export function encodeHeader(
-  src: PythGovernanceHeader,
-  buffer: Buffer
-): number {
-  let module: number;
-  let action: number;
-  if (src.action in ExecutorAction) {
-    module = MODULE_EXECUTOR;
-    action = ExecutorAction[src.action as keyof typeof ExecutorAction];
-  } else {
-    module = MODULE_TARGET;
-    action = TargetAction[src.action as keyof typeof TargetAction];
+  constructor(targetChainId: ChainName, action: ActionName) {
+    this.targetChainId = targetChainId;
+    this.action = action;
   }
-  return governanceHeaderLayout().encode(
-    {
-      magicNumber: MAGIC_NUMBER,
-      module,
-      action,
-      chain: toChainId(src.targetChainId),
-    },
-    buffer
-  );
-}
-
-export function verifyHeader(
-  deserialized: Readonly<{
-    magicNumber: number;
-    module: number;
-    action: number;
-    chain: ChainId;
-  }>
-): PythGovernanceHeader {
-  if (deserialized.magicNumber !== MAGIC_NUMBER) {
-    throw new Error("Wrong magic number");
+  /** Decode Pyth Governance Header */
+  static decode(data: Buffer): PythGovernanceHeader {
+    let deserialized = this.layout.decode(data);
+    return this.verify(deserialized);
   }
 
-  if (!toChainName(deserialized.chain)) {
-    throw new Error("Chain Id not found");
+  /** Verify header fields, takes in a raw deserialized header  */
+  static verify(
+    deserialized: Readonly<{
+      magicNumber: number;
+      module: number;
+      action: number;
+      chain: ChainId;
+    }>
+  ): PythGovernanceHeader {
+    if (deserialized.magicNumber !== MAGIC_NUMBER) {
+      throw new Error("Wrong magic number");
+    }
+
+    if (!toChainName(deserialized.chain)) {
+      throw new Error("Chain Id not found");
+    }
+
+    return new PythGovernanceHeader(
+      toChainName(deserialized.chain),
+      toActionName({
+        actionId: deserialized.action,
+        moduleId: deserialized.module,
+      })
+    );
   }
 
-  let governanceHeader: PythGovernanceHeader = {
-    targetChainId: toChainName(deserialized.chain),
-    action: toActionName({
-      actionId: deserialized.action,
-      moduleId: deserialized.module,
-    }),
-  };
-  return governanceHeader;
+  encode(): Buffer {
+    // The code will crash if the payload is actually bigger than PACKET_DATA_SIZE. But PACKET_DATA_SIZE is the maximum transaction size of Solana, so our serialized payload should never be bigger than this anyway
+    const buffer = Buffer.alloc(PACKET_DATA_SIZE);
+    let module: number;
+    let action: number;
+    if (this.action in ExecutorAction) {
+      module = MODULE_EXECUTOR;
+      action = ExecutorAction[this.action as keyof typeof ExecutorAction];
+    } else {
+      module = MODULE_TARGET;
+      action = TargetAction[this.action as keyof typeof TargetAction];
+    }
+    const span = PythGovernanceHeader.layout.encode(
+      {
+        magicNumber: MAGIC_NUMBER,
+        module,
+        action,
+        chain: toChainId(this.targetChainId),
+      },
+      buffer
+    );
+    return buffer.subarray(0, span);
+  }
 }
 
 export function decodeGovernancePayload(data: Buffer): PythGovernanceAction {
-  const header = decodeHeader(data);
+  const header = PythGovernanceHeader.decode(data);
   switch (header.action) {
     case "ExecutePostedVaa":
       return ExecutePostedVaa.decode(data);