[xc-admin] Refactor wormhole messages as classes (#478)

guibescos · web-flow · commit 8ef49e61283d · 2023-01-12T11:47:56.000-06:00
* Refactor wormhole messages as classes

* Type error

* Export ExecutePostedVaa

* Rename to UnknownGovernanceAction

* Improve interface
diff --git a/xc-admin/packages/xc-admin-common/src/__tests__/GovernancePayload.test.ts b/xc-admin/packages/xc-admin-common/src/__tests__/GovernancePayload.test.ts
@@ -1,17 +1,6 @@
 import { ChainName } from "@certusone/wormhole-sdk";
-import {
-  PACKET_DATA_SIZE,
-  PublicKey,
-  SystemProgram,
-  TransactionInstruction,
-} from "@solana/web3.js";
-import {
-  ActionName,
-  decodeExecutePostedVaa,
-  decodeHeader,
-  encodeHeader,
-} from "..";
-import { encodeExecutePostedVaa } from "../governance_payload/ExecutePostedVaa";
+import { PACKET_DATA_SIZE, PublicKey, SystemProgram } from "@solana/web3.js";
+import { ActionName, decodeHeader, encodeHeader, ExecutePostedVaa } from "..";
 
 test("GovernancePayload ser/de", (done) => {
   jest.setTimeout(60000);
@@ -75,32 +64,25 @@ test("GovernancePayload ser/de", (done) => {
   ).toThrow("Invalid header, action doesn't match module");
 
   // Decode executePostVaa with empty instructions
-  let expectedExecuteVaaArgs = {
-    targetChainId: "pythnet" as ChainName,
-    instructions: [] as TransactionInstruction[],
-  };
-  buffer = encodeExecutePostedVaa(expectedExecuteVaaArgs);
+  let expectedExecutePostedVaa = new ExecutePostedVaa("pythnet", []);
+  buffer = expectedExecutePostedVaa.encode();
   expect(
     buffer.equals(Buffer.from([80, 84, 71, 77, 0, 0, 0, 26, 0, 0, 0, 0]))
   ).toBeTruthy();
-  let executePostedVaaArgs = decodeExecutePostedVaa(buffer);
+  let executePostedVaaArgs = ExecutePostedVaa.decode(buffer);
   expect(executePostedVaaArgs?.targetChainId).toBe("pythnet");
   expect(executePostedVaaArgs?.instructions.length).toBe(0);
 
   // Decode executePostVaa with one system instruction
-  expectedExecuteVaaArgs = {
-    targetChainId: "pythnet" as ChainName,
-    instructions: [
-      SystemProgram.transfer({
-        fromPubkey: new PublicKey(
-          "AWQ18oKzd187aM2oMB4YirBcdgX1FgWfukmqEX91BRES"
-        ),
-        toPubkey: new PublicKey("J25GT2knN8V2Wvg9jNrYBuj9SZdsLnU6bK7WCGrL7daj"),
-        lamports: 890880,
-      }),
-    ] as TransactionInstruction[],
-  };
-  buffer = encodeExecutePostedVaa(expectedExecuteVaaArgs);
+  expectedExecutePostedVaa = new ExecutePostedVaa("pythnet", [
+    SystemProgram.transfer({
+      fromPubkey: new PublicKey("AWQ18oKzd187aM2oMB4YirBcdgX1FgWfukmqEX91BRES"),
+      toPubkey: new PublicKey("J25GT2knN8V2Wvg9jNrYBuj9SZdsLnU6bK7WCGrL7daj"),
+      lamports: 890880,
+    }),
+  ]);
+
+  buffer = expectedExecutePostedVaa.encode();
   expect(
     buffer.equals(
       Buffer.from([
@@ -115,7 +97,7 @@ test("GovernancePayload ser/de", (done) => {
       ])
     )
   ).toBeTruthy();
-  executePostedVaaArgs = decodeExecutePostedVaa(buffer);
+  executePostedVaaArgs = ExecutePostedVaa.decode(buffer);
   expect(executePostedVaaArgs?.targetChainId).toBe("pythnet");
   expect(executePostedVaaArgs?.instructions.length).toBe(1);
   expect(
diff --git a/xc-admin/packages/xc-admin-common/src/__tests__/WormholeMultisigInstruction.test.ts b/xc-admin/packages/xc-admin-common/src/__tests__/WormholeMultisigInstruction.test.ts
@@ -1,4 +1,3 @@
-import { ChainName } from "@certusone/wormhole-sdk";
 import { createWormholeProgramInterface } from "@certusone/wormhole-sdk/lib/cjs/solana/wormhole";
 import { AnchorProvider, Wallet } from "@project-serum/anchor";
 import {
@@ -16,11 +15,8 @@ import {
   MultisigInstructionProgram,
   MultisigParser,
   WORMHOLE_ADDRESS,
+  ExecutePostedVaa,
 } from "..";
-import {
-  encodeExecutePostedVaa,
-  ExecutePostedVaaArgs,
-} from "../governance_payload/ExecutePostedVaa";
 import { WormholeMultisigInstruction } from "../multisig_transaction/WormholeMultisigInstruction";
 
 test("Wormhole multisig instruction parse: send message without governance payload", (done) => {
@@ -184,19 +180,16 @@ test("Wormhole multisig instruction parse: send message with governance payload"
   );
   const parser = MultisigParser.fromCluster(cluster);
 
-  const executePostedVaaArgs: ExecutePostedVaaArgs = {
-    targetChainId: "pythnet" as ChainName,
-    instructions: [
-      SystemProgram.transfer({
-        fromPubkey: PublicKey.unique(),
-        toPubkey: PublicKey.unique(),
-        lamports: 890880,
-      }),
-    ],
-  };
+  const executePostedVaa: ExecutePostedVaa = new ExecutePostedVaa("pythnet", [
+    SystemProgram.transfer({
+      fromPubkey: PublicKey.unique(),
+      toPubkey: PublicKey.unique(),
+      lamports: 890880,
+    }),
+  ]);
 
   wormholeProgram.methods
-    .postMessage(0, encodeExecutePostedVaa(executePostedVaaArgs), 0)
+    .postMessage(0, executePostedVaa.encode(), 0)
     .accounts({
       bridge: PublicKey.unique(),
       message: PublicKey.unique(),
@@ -319,45 +312,47 @@ test("Wormhole multisig instruction parse: send message with governance payload"
 
         expect(parsedInstruction.args.nonce).toBe(0);
         expect(
-          parsedInstruction.args.payload.equals(
-            encodeExecutePostedVaa(executePostedVaaArgs)
-          )
+          parsedInstruction.args.payload.equals(executePostedVaa.encode())
         );
         expect(parsedInstruction.args.consistencyLevel).toBe(0);
 
-        expect(parsedInstruction.args.governanceName).toBe("ExecutePostedVaa");
-
-        expect(parsedInstruction.args.governanceArgs.targetChainId).toBe(
-          "pythnet"
-        );
-
-        (
-          parsedInstruction.args.governanceArgs
-            .instructions as TransactionInstruction[]
-        ).forEach((instruction, i) => {
-          expect(
-            instruction.programId.equals(
-              executePostedVaaArgs.instructions[i].programId
-            )
+        if (
+          parsedInstruction.args.governanceAction instanceof ExecutePostedVaa
+        ) {
+          expect(parsedInstruction.args.governanceAction.targetChainId).toBe(
+            "pythnet"
           );
-          expect(
-            instruction.data.equals(executePostedVaaArgs.instructions[i].data)
-          );
-          instruction.keys.forEach((account, j) => {
+
+          (
+            parsedInstruction.args.governanceAction
+              .instructions as TransactionInstruction[]
+          ).forEach((instruction, i) => {
             expect(
-              account.pubkey.equals(
-                executePostedVaaArgs.instructions[i].keys[j].pubkey
+              instruction.programId.equals(
+                executePostedVaa.instructions[i].programId
               )
-            ).toBeTruthy();
-            expect(account.isSigner).toBe(
-              executePostedVaaArgs.instructions[i].keys[j].isSigner
             );
-            expect(account.isWritable).toBe(
-              executePostedVaaArgs.instructions[i].keys[j].isWritable
+            expect(
+              instruction.data.equals(executePostedVaa.instructions[i].data)
             );
+            instruction.keys.forEach((account, j) => {
+              expect(
+                account.pubkey.equals(
+                  executePostedVaa.instructions[i].keys[j].pubkey
+                )
+              ).toBeTruthy();
+              expect(account.isSigner).toBe(
+                executePostedVaa.instructions[i].keys[j].isSigner
+              );
+              expect(account.isWritable).toBe(
+                executePostedVaa.instructions[i].keys[j].isWritable
+              );
+            });
           });
-        });
-        done();
+          done();
+        } else {
+          done("Not instance of ExecutePostedVaa");
+        }
       } else {
         done("Not instance of WormholeInstruction");
       }
diff --git a/xc-admin/packages/xc-admin-common/src/governance_payload/ExecutePostedVaa.ts b/xc-admin/packages/xc-admin-common/src/governance_payload/ExecutePostedVaa.ts
@@ -3,7 +3,7 @@ import * as BufferLayout from "@solana/buffer-layout";
 import {
   encodeHeader,
   governanceHeaderLayout,
-  PythGovernanceHeader,
+  PythGovernanceAction,
   verifyHeader,
 } from ".";
 import { Layout } from "@solana/buffer-layout";
@@ -77,62 +77,69 @@ export const executePostedVaaLayout: BufferLayout.Structure<
   new Vector<InstructionData>(instructionDataLayout, "instructions"),
 ]);
 
-export type ExecutePostedVaaArgs = {
-  targetChainId: ChainName;
-  instructions: TransactionInstruction[];
-};
+export class ExecutePostedVaa implements PythGovernanceAction {
+  readonly targetChainId: ChainName;
+  readonly instructions: TransactionInstruction[];
+
+  constructor(
+    targetChainId: ChainName,
+    instructions: TransactionInstruction[]
+  ) {
+    this.targetChainId = targetChainId;
+    this.instructions = instructions;
+  }
 
-/** Decode ExecutePostedVaaArgs and return undefined if it failed */
-export function decodeExecutePostedVaa(data: Buffer): ExecutePostedVaaArgs {
-  let deserialized = executePostedVaaLayout.decode(data);
+  /** Decode ExecutePostedVaaArgs */
+  static decode(data: Buffer): ExecutePostedVaa {
+    let deserialized = executePostedVaaLayout.decode(data);
 
-  let header = verifyHeader(deserialized.header);
+    let header = verifyHeader(deserialized.header);
+
+    let instructions: TransactionInstruction[] = deserialized.instructions.map(
+      (ix) => {
+        let programId: PublicKey = new PublicKey(ix.programId);
+        let keys: AccountMeta[] = ix.accounts.map((acc) => {
+          return {
+            pubkey: new PublicKey(acc.pubkey),
+            isSigner: Boolean(acc.isSigner),
+            isWritable: Boolean(acc.isWritable),
+          };
+        });
+        let data: Buffer = Buffer.from(ix.data);
+        return { programId, keys, data };
+      }
+    );
+    return new ExecutePostedVaa(header.targetChainId, instructions);
+  }
 
-  let instructions: TransactionInstruction[] = deserialized.instructions.map(
-    (ix) => {
-      let programId: PublicKey = new PublicKey(ix.programId);
-      let keys: AccountMeta[] = ix.accounts.map((acc) => {
+  /** 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 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) => {
         return {
-          pubkey: new PublicKey(acc.pubkey),
-          isSigner: Boolean(acc.isSigner),
-          isWritable: Boolean(acc.isWritable),
+          pubkey: acc.pubkey.toBytes(),
+          isSigner: acc.isSigner ? 1 : 0,
+          isWritable: acc.isWritable ? 1 : 0,
         };
       });
-      let data: Buffer = Buffer.from(ix.data);
-      return { programId, keys, data };
-    }
-  );
-
-  return { targetChainId: header.targetChainId, instructions };
-}
-
-/** Encode ExecutePostedVaaArgs */
-export function encodeExecutePostedVaa(src: ExecutePostedVaaArgs): Buffer {
-  // PACKET_DATA_SIZE is the maximum transactin size of Solana, so our serialized payload will never be bigger than that
-  const buffer = Buffer.alloc(PACKET_DATA_SIZE);
-  const offset = encodeHeader(
-    { action: "ExecutePostedVaa", targetChainId: src.targetChainId },
-    buffer
-  );
-  let instructions: InstructionData[] = src.instructions.map((ix) => {
-    let programId = ix.programId.toBytes();
-    let accounts: AccountMetadata[] = ix.keys.map((acc) => {
-      return {
-        pubkey: acc.pubkey.toBytes(),
-        isSigner: acc.isSigner ? 1 : 0,
-        isWritable: acc.isWritable ? 1 : 0,
-      };
+      let data = [...ix.data];
+      return { programId, accounts, data };
     });
-    let data = [...ix.data];
-    return { programId, accounts, data };
-  });
 
-  const span =
-    offset +
-    new Vector<InstructionData>(instructionDataLayout, "instructions").encode(
-      instructions,
-      buffer,
-      offset
-    );
-  return buffer.subarray(0, span);
+    const span =
+      offset +
+      new Vector<InstructionData>(instructionDataLayout, "instructions").encode(
+        instructions,
+        buffer,
+        offset
+      );
+    return buffer.subarray(0, span);
+  }
 }
diff --git a/xc-admin/packages/xc-admin-common/src/governance_payload/index.ts b/xc-admin/packages/xc-admin-common/src/governance_payload/index.ts
@@ -5,10 +5,12 @@ import {
   toChainName,
 } from "@certusone/wormhole-sdk";
 import * as BufferLayout from "@solana/buffer-layout";
-import {
-  decodeExecutePostedVaa,
-  ExecutePostedVaaArgs,
-} from "./ExecutePostedVaa";
+import { ExecutePostedVaa } from "./ExecutePostedVaa";
+
+export interface PythGovernanceAction {
+  readonly targetChainId: ChainName;
+  encode(): Buffer;
+}
 
 export const ExecutorAction = {
   ExecutePostedVaa: 0,
@@ -134,17 +136,14 @@ export function verifyHeader(
   return governanceHeader;
 }
 
-export function decodeGovernancePayload(data: Buffer): {
-  name: string;
-  args: ExecutePostedVaaArgs;
-} {
+export function decodeGovernancePayload(data: Buffer): PythGovernanceAction {
   const header = decodeHeader(data);
   switch (header.action) {
     case "ExecutePostedVaa":
-      return { name: "ExecutePostedVaa", args: decodeExecutePostedVaa(data) };
+      return ExecutePostedVaa.decode(data);
     default:
       throw "Not supported";
   }
 }
 
-export { decodeExecutePostedVaa } from "./ExecutePostedVaa";
+export { ExecutePostedVaa } from "./ExecutePostedVaa";
diff --git a/xc-admin/packages/xc-admin-common/src/multisig_transaction/WormholeMultisigInstruction.ts b/xc-admin/packages/xc-admin-common/src/multisig_transaction/WormholeMultisigInstruction.ts
@@ -3,7 +3,10 @@ import { WormholeInstructionCoder } from "@certusone/wormhole-sdk/lib/cjs/solana
 import { getPythClusterApiUrl } from "@pythnetwork/client/lib/cluster";
 import { Connection, TransactionInstruction } from "@solana/web3.js";
 import { MultisigInstruction, MultisigInstructionProgram } from ".";
-import { decodeGovernancePayload } from "../governance_payload";
+import {
+  decodeGovernancePayload,
+  PythGovernanceAction,
+} from "../governance_payload";
 import { AnchorAccounts, resolveAccountNames } from "./anchor";
 
 export class WormholeMultisigInstruction implements MultisigInstruction {
@@ -47,12 +50,12 @@ export class WormholeMultisigInstruction implements MultisigInstruction {
 
       if (result.name === "postMessage") {
         try {
-          const decoded = decodeGovernancePayload(result.args.payload);
-          result.args.governanceName = decoded.name;
-          result.args.governanceArgs = decoded.args;
+          const decoded: PythGovernanceAction = decodeGovernancePayload(
+            result.args.payload
+          );
+          result.args.governanceAction = decoded;
         } catch {
-          result.args.governanceName = "Unrecognized governance message";
-          result.args.governanceArgs = {};
+          result.args.governanceAction = {};
         }
       }
       return result;
diff --git a/xc-admin/packages/xc-admin-common/src/propose.ts b/xc-admin/packages/xc-admin-common/src/propose.ts
