From 6c740685788b649e8bd74d896c67e3a66cb60229 Mon Sep 17 00:00:00 2001
From: Justin Starry <justin@solana.com>
Date: Wed, 31 Aug 2022 10:43:19 -0400
Subject: [PATCH] feat: implement message v0 decompilation

---
 web3.js/src/message/legacy.ts                 |  36 ++++-
 web3.js/src/message/v0.ts                     |  90 +++++++++++
 web3.js/src/transaction/index.ts              |   1 +
 web3.js/src/transaction/message.ts            | 147 ++++++++++++++++++
 web3.js/test/message-tests/legacy.test.ts     |  91 +++++++++++
 web3.js/test/message-tests/v0.test.ts         | 142 ++++++++++++++++-
 .../test/transaction-tests/message.test.ts    |  89 +++++++++++
 7 files changed, 593 insertions(+), 3 deletions(-)
 create mode 100644 web3.js/src/transaction/message.ts
 create mode 100644 web3.js/test/message-tests/legacy.test.ts
 create mode 100644 web3.js/test/transaction-tests/message.test.ts

diff --git a/web3.js/src/message/legacy.ts b/web3.js/src/message/legacy.ts
index 38faa6320a0dbd..8e5116fb51b637 100644
--- a/web3.js/src/message/legacy.ts
+++ b/web3.js/src/message/legacy.ts
@@ -13,6 +13,9 @@ import {
   MessageAddressTableLookup,
   MessageCompiledInstruction,
 } from './index';
+import {TransactionInstruction} from '../transaction';
+import {CompiledKeys} from './compiled-keys';
+import {MessageAccountKeys} from './account-keys';
 
 /**
  * An instruction to execute by a program
@@ -37,13 +40,19 @@ export type MessageArgs = {
   /** The message header, identifying signed and read-only `accountKeys` */
   header: MessageHeader;
   /** All the account keys used by this transaction */
-  accountKeys: string[];
+  accountKeys: string[] | PublicKey[];
   /** The hash of a recent ledger block */
   recentBlockhash: Blockhash;
   /** Instructions that will be executed in sequence and committed in one atomic transaction if all succeed. */
   instructions: CompiledInstruction[];
 };
 
+export type CompileLegacyArgs = {
+  payerKey: PublicKey;
+  instructions: Array<TransactionInstruction>;
+  recentBlockhash: Blockhash;
+};
+
 /**
  * List of instructions to be processed atomically
  */
@@ -93,6 +102,29 @@ export class Message {
     return [];
   }
 
+  getAccountKeys(): MessageAccountKeys {
+    return new MessageAccountKeys(this.staticAccountKeys);
+  }
+
+  static compile(args: CompileLegacyArgs): Message {
+    const compiledKeys = CompiledKeys.compile(args.instructions, args.payerKey);
+    const [header, staticAccountKeys] = compiledKeys.getMessageComponents();
+    const accountKeys = new MessageAccountKeys(staticAccountKeys);
+    const instructions = accountKeys.compileInstructions(args.instructions).map(
+      (ix: MessageCompiledInstruction): CompiledInstruction => ({
+        programIdIndex: ix.programIdIndex,
+        accounts: ix.accountKeyIndexes,
+        data: bs58.encode(ix.data),
+      }),
+    );
+    return new Message({
+      header,
+      accountKeys: staticAccountKeys,
+      recentBlockhash: args.recentBlockhash,
+      instructions,
+    });
+  }
+
   isAccountSigner(index: number): boolean {
     return index < this.header.numRequiredSignatures;
   }
@@ -250,7 +282,7 @@ export class Message {
     for (let i = 0; i < accountCount; i++) {
       const account = byteArray.slice(0, PUBLIC_KEY_LENGTH);
       byteArray = byteArray.slice(PUBLIC_KEY_LENGTH);
-      accountKeys.push(bs58.encode(Buffer.from(account)));
+      accountKeys.push(new PublicKey(Buffer.from(account)));
     }
 
     const recentBlockhash = byteArray.slice(0, PUBLIC_KEY_LENGTH);
diff --git a/web3.js/src/message/v0.ts b/web3.js/src/message/v0.ts
index 800ad162c163be..0f452602063cff 100644
--- a/web3.js/src/message/v0.ts
+++ b/web3.js/src/message/v0.ts
@@ -40,6 +40,14 @@ export type CompileV0Args = {
   addressLookupTableAccounts?: Array<AddressLookupTableAccount>;
 };
 
+export type GetAccountKeysArgs =
+  | {
+      accountKeysFromLookups: AccountKeysFromLookups;
+    }
+  | {
+      addressLookupTableAccounts: AddressLookupTableAccount[];
+    };
+
 export class MessageV0 {
   header: MessageHeader;
   staticAccountKeys: Array<PublicKey>;
@@ -59,6 +67,88 @@ export class MessageV0 {
     return 0;
   }
 
+  get numAccountKeysFromLookups(): number {
+    let count = 0;
+    for (const lookup of this.addressTableLookups) {
+      count += lookup.readonlyIndexes.length + lookup.writableIndexes.length;
+    }
+    return count;
+  }
+
+  getAccountKeys(args?: GetAccountKeysArgs): MessageAccountKeys {
+    let accountKeysFromLookups: AccountKeysFromLookups | undefined;
+    if (args && 'accountKeysFromLookups' in args) {
+      if (
+        this.numAccountKeysFromLookups !=
+        args.accountKeysFromLookups.writable.length +
+          args.accountKeysFromLookups.readonly.length
+      ) {
+        throw new Error(
+          'Failed to get account keys because of a mismatch in the number of account keys from lookups',
+        );
+      }
+      accountKeysFromLookups = args.accountKeysFromLookups;
+    } else if (args && 'addressLookupTableAccounts' in args) {
+      accountKeysFromLookups = this.resolveAddressTableLookups(
+        args.addressLookupTableAccounts,
+      );
+    } else if (this.addressTableLookups.length > 0) {
+      throw new Error(
+        'Failed to get account keys because address table lookups were not resolved',
+      );
+    }
+    return new MessageAccountKeys(
+      this.staticAccountKeys,
+      accountKeysFromLookups,
+    );
+  }
+
+  resolveAddressTableLookups(
+    addressLookupTableAccounts: AddressLookupTableAccount[],
+  ): AccountKeysFromLookups {
+    const accountKeysFromLookups: AccountKeysFromLookups = {
+      writable: [],
+      readonly: [],
+    };
+
+    for (const tableLookup of this.addressTableLookups) {
+      const tableAccount = addressLookupTableAccounts.find(account =>
+        account.key.equals(tableLookup.accountKey),
+      );
+      if (!tableAccount) {
+        throw new Error(
+          `Failed to find address lookup table account for table key ${tableLookup.accountKey.toBase58()}`,
+        );
+      }
+
+      for (const index of tableLookup.writableIndexes) {
+        if (index < tableAccount.state.addresses.length) {
+          accountKeysFromLookups.writable.push(
+            tableAccount.state.addresses[index],
+          );
+        } else {
+          throw new Error(
+            `Failed to find address for index ${index} in address lookup table ${tableLookup.accountKey.toBase58()}`,
+          );
+        }
+      }
+
+      for (const index of tableLookup.readonlyIndexes) {
+        if (index < tableAccount.state.addresses.length) {
+          accountKeysFromLookups.readonly.push(
+            tableAccount.state.addresses[index],
+          );
+        } else {
+          throw new Error(
+            `Failed to find address for index ${index} in address lookup table ${tableLookup.accountKey.toBase58()}`,
+          );
+        }
+      }
+    }
+
+    return accountKeysFromLookups;
+  }
+
   static compile(args: CompileV0Args): MessageV0 {
     const compiledKeys = CompiledKeys.compile(args.instructions, args.payerKey);
 
diff --git a/web3.js/src/transaction/index.ts b/web3.js/src/transaction/index.ts
index 2f5c19cb2510a5..88d1cb1700ea62 100644
--- a/web3.js/src/transaction/index.ts
+++ b/web3.js/src/transaction/index.ts
@@ -1,4 +1,5 @@
 export * from './constants';
 export * from './expiry-custom-errors';
 export * from './legacy';
+export * from './message';
 export * from './versioned';
diff --git a/web3.js/src/transaction/message.ts b/web3.js/src/transaction/message.ts
new file mode 100644
index 00000000000000..2b0a69572b1b21
--- /dev/null
+++ b/web3.js/src/transaction/message.ts
@@ -0,0 +1,147 @@
+import {
+  AccountKeysFromLookups,
+  MessageAccountKeys,
+} from '../message/account-keys';
+import assert from '../utils/assert';
+import {toBuffer} from '../utils/to-buffer';
+import {Blockhash} from '../blockhash';
+import {Message, MessageV0, VersionedMessage} from '../message';
+import {AddressLookupTableAccount} from '../programs';
+import {AccountMeta, TransactionInstruction} from './legacy';
+
+export type TransactionMessageArgs = {
+  accountKeys: MessageAccountKeys;
+  instructions: Array<TransactionInstruction>;
+  recentBlockhash: Blockhash;
+};
+
+export type DecompileArgs =
+  | {
+      accountKeysFromLookups: AccountKeysFromLookups;
+    }
+  | {
+      addressLookupTableAccounts: AddressLookupTableAccount[];
+    };
+
+export class TransactionMessage {
+  accountKeys: MessageAccountKeys;
+  instructions: Array<TransactionInstruction>;
+  recentBlockhash: Blockhash;
+
+  constructor(args: TransactionMessageArgs) {
+    this.accountKeys = args.accountKeys;
+    this.instructions = args.instructions;
+    this.recentBlockhash = args.recentBlockhash;
+  }
+
+  static decompile(
+    message: VersionedMessage,
+    args?: DecompileArgs,
+  ): TransactionMessage {
+    const {header, compiledInstructions, recentBlockhash} = message;
+
+    const {
+      numRequiredSignatures,
+      numReadonlySignedAccounts,
+      numReadonlyUnsignedAccounts,
+    } = header;
+
+    const numWritableSignedAccounts =
+      numRequiredSignatures - numReadonlySignedAccounts;
+    assert(numWritableSignedAccounts > 0, 'Message header is invalid');
+
+    const numWritableUnsignedAccounts =
+      message.staticAccountKeys.length - numReadonlyUnsignedAccounts;
+    assert(numWritableUnsignedAccounts >= 0, 'Message header is invalid');
+
+    const accountKeys = message.getAccountKeys(args);
+    const instructions: TransactionInstruction[] = [];
+    for (const compiledIx of compiledInstructions) {
+      const keys: AccountMeta[] = [];
+
+      for (const keyIndex of compiledIx.accountKeyIndexes) {
+        const pubkey = accountKeys.get(keyIndex);
+        if (pubkey === undefined) {
+          throw new Error(
+            `Failed to find key for account key index ${keyIndex}`,
+          );
+        }
+
+        const isSigner = keyIndex < numRequiredSignatures;
+
+        let isWritable;
+        if (isSigner) {
+          isWritable = keyIndex < numWritableSignedAccounts;
+        } else if (keyIndex < accountKeys.staticAccountKeys.length) {
+          isWritable =
+            keyIndex - numRequiredSignatures < numWritableUnsignedAccounts;
+        } else {
+          isWritable =
+            keyIndex - accountKeys.staticAccountKeys.length <
+            // accountKeysFromLookups cannot be undefined because we already found a pubkey for this index above
+            accountKeys.accountKeysFromLookups!.writable.length;
+        }
+
+        keys.push({
+          pubkey,
+          isSigner: keyIndex < header.numRequiredSignatures,
+          isWritable,
+        });
+      }
+
+      const programId = accountKeys.get(compiledIx.programIdIndex);
+      if (programId === undefined) {
+        throw new Error(
+          `Failed to find program id for program id index ${compiledIx.programIdIndex}`,
+        );
+      }
+
+      instructions.push(
+        new TransactionInstruction({
+          programId,
+          data: toBuffer(compiledIx.data),
+          keys,
+        }),
+      );
+    }
+
+    return new TransactionMessage({
+      accountKeys,
+      instructions,
+      recentBlockhash,
+    });
+  }
+
+  compileToLegacyMessage(): Message {
+    const payerKey = this.accountKeys.get(0);
+    if (payerKey === undefined) {
+      throw new Error(
+        'Failed to compile message because no account keys were found',
+      );
+    }
+
+    return Message.compile({
+      payerKey,
+      recentBlockhash: this.recentBlockhash,
+      instructions: this.instructions,
+    });
+  }
+
+  compileToV0Message(
+    addressLookupTableAccounts?: AddressLookupTableAccount[],
+  ): MessageV0 {
+    const payerKey = this.accountKeys.get(0);
+    if (payerKey === undefined) {
+      throw new Error(
+        'Failed to compile message because no account keys were found',
+      );
+    }
+
+    return MessageV0.compile({
+      payerKey,
+      recentBlockhash: this.recentBlockhash,
+      instructions: this.instructions,
+      addressLookupTableAccounts,
+    });
+  }
+}
diff --git a/web3.js/test/message-tests/legacy.test.ts b/web3.js/test/message-tests/legacy.test.ts
new file mode 100644
index 00000000000000..a602c189698bf0
--- /dev/null
+++ b/web3.js/test/message-tests/legacy.test.ts
@@ -0,0 +1,91 @@
+import bs58 from 'bs58';
+import {expect} from 'chai';
+import {sha256} from '@noble/hashes/sha256';
+
+import {Message} from '../../src/message';
+import {TransactionInstruction} from '../../src/transaction';
+import {PublicKey} from '../../src/publickey';
+
+function createTestKeys(count: number): Array<PublicKey> {
+  return new Array(count).fill(0).map(() => PublicKey.unique());
+}
+
+describe('Message', () => {
+  it('compile', () => {
+    const keys = createTestKeys(5);
+    const recentBlockhash = bs58.encode(sha256('test'));
+    const payerKey = keys[0];
+    const instructions = [
+      new TransactionInstruction({
+        programId: keys[4],
+        keys: [
+          {pubkey: keys[1], isSigner: true, isWritable: true},
+          {pubkey: keys[2], isSigner: false, isWritable: false},
+          {pubkey: keys[3], isSigner: false, isWritable: false},
+        ],
+        data: Buffer.alloc(1),
+      }),
+      new TransactionInstruction({
+        programId: keys[1],
+        keys: [
+          {pubkey: keys[2], isSigner: true, isWritable: false},
+          {pubkey: keys[3], isSigner: false, isWritable: true},
+        ],
+        data: Buffer.alloc(2),
+      }),
+    ];
+
+    const message = Message.compile({
+      payerKey,
+      recentBlockhash,
+      instructions,
+    });
+
+    expect(message.accountKeys).to.eql([
+      payerKey, // payer is first
+      keys[1], // other writable signer
+      keys[2], // sole readonly signer
+      keys[3], // sole writable non-signer
+      keys[4], // sole readonly non-signer
+    ]);
+    expect(message.header).to.eql({
+      numRequiredSignatures: 3,
+      numReadonlySignedAccounts: 1,
+      numReadonlyUnsignedAccounts: 1,
+    });
+    expect(message.addressTableLookups.length).to.eq(0);
+    expect(message.instructions).to.eql([
+      {
+        programIdIndex: 4,
+        accounts: [1, 2, 3],
+        data: bs58.encode(Buffer.alloc(1)),
+      },
+      {
+        programIdIndex: 1,
+        accounts: [2, 3],
+        data: bs58.encode(Buffer.alloc(2)),
+      },
+    ]);
+    expect(message.recentBlockhash).to.eq(recentBlockhash);
+  });
+
+  it('compile without instructions', () => {
+    const payerKey = PublicKey.unique();
+    const recentBlockhash = bs58.encode(sha256('test'));
+    const message = Message.compile({
+      payerKey,
+      instructions: [],
+      recentBlockhash,
+    });
+
+    expect(message.accountKeys).to.eql([payerKey]);
+    expect(message.header).to.eql({
+      numRequiredSignatures: 1,
+      numReadonlySignedAccounts: 0,
+      numReadonlyUnsignedAccounts: 0,
+    });
+    expect(message.addressTableLookups.length).to.eq(0);
+    expect(message.instructions.length).to.eq(0);
+    expect(message.recentBlockhash).to.eq(recentBlockhash);
+  });
+});
diff --git a/web3.js/test/message-tests/v0.test.ts b/web3.js/test/message-tests/v0.test.ts
index 8538b64ccb5fb5..1f95b1d3b7d682 100644
--- a/web3.js/test/message-tests/v0.test.ts
+++ b/web3.js/test/message-tests/v0.test.ts
@@ -2,7 +2,11 @@ import bs58 from 'bs58';
 import {expect} from 'chai';
 import {sha256} from '@noble/hashes/sha256';
 
-import {MessageV0} from '../../src/message';
+import {
+  MessageAccountKeys,
+  MessageAddressTableLookup,
+  MessageV0,
+} from '../../src/message';
 import {TransactionInstruction} from '../../src/transaction';
 import {PublicKey} from '../../src/publickey';
 import {AddressLookupTableAccount} from '../../src/programs';
@@ -28,6 +32,142 @@ function createTestLookupTable(
 }
 
 describe('MessageV0', () => {
+  it('numAccountKeysFromLookups', () => {
+    const message = MessageV0.compile({
+      payerKey: PublicKey.unique(),
+      recentBlockhash: '',
+      instructions: [],
+    });
+    expect(message.numAccountKeysFromLookups).to.eq(0);
+
+    message.addressTableLookups = [
+      {
+        accountKey: PublicKey.unique(),
+        writableIndexes: [0],
+        readonlyIndexes: [1],
+      },
+      {
+        accountKey: PublicKey.unique(),
+        writableIndexes: [0, 2],
+        readonlyIndexes: [],
+      },
+    ];
+    expect(message.numAccountKeysFromLookups).to.eq(4);
+  });
+
+  it('getAccountKeys', () => {
+    const staticAccountKeys = createTestKeys(3);
+    const lookupTable = createTestLookupTable(createTestKeys(2));
+    const message = new MessageV0({
+      header: {
+        numRequiredSignatures: 1,
+        numReadonlySignedAccounts: 0,
+        numReadonlyUnsignedAccounts: 0,
+      },
+      recentBlockhash: 'test',
+      staticAccountKeys,
+      compiledInstructions: [],
+      addressTableLookups: [
+        {
+          accountKey: lookupTable.key,
+          writableIndexes: [0],
+          readonlyIndexes: [1],
+        },
+      ],
+    });
+
+    expect(() => message.getAccountKeys()).to.throw(
+      'Failed to get account keys because address table lookups were not resolved',
+    );
+    expect(() =>
+      message.getAccountKeys({
+        accountKeysFromLookups: {writable: [PublicKey.unique()], readonly: []},
+      }),
+    ).to.throw(
+      'Failed to get account keys because of a mismatch in the number of account keys from lookups',
+    );
+
+    const accountKeysFromLookups = message.resolveAddressTableLookups([
+      lookupTable,
+    ]);
+    const expectedAccountKeys = new MessageAccountKeys(
+      staticAccountKeys,
+      accountKeysFromLookups,
+    );
+
+    expect(
+      message.getAccountKeys({
+        accountKeysFromLookups,
+      }),
+    ).to.eql(expectedAccountKeys);
+
+    expect(
+      message.getAccountKeys({
+        addressLookupTableAccounts: [lookupTable],
+      }),
+    ).to.eql(expectedAccountKeys);
+  });
+
+  it('resolveAddressTableLookups', () => {
+    const keys = createTestKeys(7);
+    const lookupTable = createTestLookupTable(keys);
+    const createTestMessage = (
+      addressTableLookups: MessageAddressTableLookup[],
+    ): MessageV0 => {
+      return new MessageV0({
+        header: {
+          numRequiredSignatures: 1,
+          numReadonlySignedAccounts: 0,
+          numReadonlyUnsignedAccounts: 0,
+        },
+        recentBlockhash: 'test',
+        staticAccountKeys: [],
+        compiledInstructions: [],
+        addressTableLookups,
+      });
+    };
+
+    expect(
+      createTestMessage([]).resolveAddressTableLookups([lookupTable]),
+    ).to.eql({
+      writable: [],
+      readonly: [],
+    });
+
+    expect(() =>
+      createTestMessage([
+        {
+          accountKey: PublicKey.unique(),
+          writableIndexes: [1, 3, 5],
+          readonlyIndexes: [0, 2, 4],
+        },
+      ]).resolveAddressTableLookups([lookupTable]),
+    ).to.throw('Failed to find address lookup table account for table key');
+
+    expect(() =>
+      createTestMessage([
+        {
+          accountKey: lookupTable.key,
+          writableIndexes: [10],
+          readonlyIndexes: [],
+        },
+      ]).resolveAddressTableLookups([lookupTable]),
+    ).to.throw('Failed to find address for index');
+
+    expect(
+      createTestMessage([
+        {
+          accountKey: lookupTable.key,
+          writableIndexes: [1, 3, 5],
+          readonlyIndexes: [0, 2, 4],
+        },
+      ]).resolveAddressTableLookups([lookupTable]),
+    ).to.eql({
+      writable: [keys[1], keys[3], keys[5]],
+      readonly: [keys[0], keys[2], keys[4]],
+    });
+  });
+
   it('compile', () => {
     const keys = createTestKeys(7);
     const recentBlockhash = bs58.encode(sha256('test'));
diff --git a/web3.js/test/transaction-tests/message.test.ts b/web3.js/test/transaction-tests/message.test.ts
new file mode 100644
index 00000000000000..7aad0e5513e05c
--- /dev/null
+++ b/web3.js/test/transaction-tests/message.test.ts
@@ -0,0 +1,89 @@
+import bs58 from 'bs58';
+import {expect} from 'chai';
+import {sha256} from '@noble/hashes/sha256';
+
+import {
+  TransactionInstruction,
+  TransactionMessage,
+} from '../../src/transaction';
+import {PublicKey} from '../../src/publickey';
+import {AddressLookupTableAccount} from '../../src/programs';
+import {MessageV0} from '../../src/message';
+
+function createTestKeys(count: number): Array<PublicKey> {
+  return new Array(count).fill(0).map(() => PublicKey.unique());
+}
+
+function createTestLookupTable(
+  addresses: Array<PublicKey>,
+): AddressLookupTableAccount {
+  const U64_MAX = 2n ** 64n - 1n;
+  return new AddressLookupTableAccount({
+    key: PublicKey.unique(),
+    state: {
+      lastExtendedSlot: 0,
+      lastExtendedSlotStartIndex: 0,
+      deactivationSlot: U64_MAX,
+      authority: PublicKey.unique(),
+      addresses,
+    },
+  });
+}
+
+describe('TransactionMessage', () => {
+  it('decompile', () => {
+    const keys = createTestKeys(7);
+    const recentBlockhash = bs58.encode(sha256('test'));
+    const payerKey = keys[0];
+    const instructions = [
+      new TransactionInstruction({
+        programId: keys[4],
+        keys: [
+          {pubkey: keys[1], isSigner: true, isWritable: true},
+          {pubkey: keys[2], isSigner: true, isWritable: false},
+          {pubkey: keys[3], isSigner: false, isWritable: true},
+          {pubkey: keys[5], isSigner: false, isWritable: true},
+          {pubkey: keys[6], isSigner: false, isWritable: false},
+        ],
+        data: Buffer.alloc(1),
+      }),
+      new TransactionInstruction({
+        programId: keys[1],
+        keys: [],
+        data: Buffer.alloc(2),
+      }),
+      new TransactionInstruction({
+        programId: keys[3],
+        keys: [],
+        data: Buffer.alloc(3),
+      }),
+    ];
+
+    const addressLookupTableAccounts = [createTestLookupTable(keys)];
+    const message = MessageV0.compile({
+      payerKey,
+      recentBlockhash,
+      instructions,
+      addressLookupTableAccounts,
+    });
+
+    expect(() => TransactionMessage.decompile(message)).to.throw(
+      'Failed to get account keys because address table lookups were not resolved',
+    );
+
+    const accountKeys = message.getAccountKeys({addressLookupTableAccounts});
+    const decompiledMessage = TransactionMessage.decompile(message, {
+      addressLookupTableAccounts,
+    });
+
+    expect(decompiledMessage.accountKeys).to.eql(accountKeys);
+    expect(decompiledMessage.recentBlockhash).to.eq(recentBlockhash);
+    expect(decompiledMessage.instructions).to.eql(instructions);
+
+    expect(decompiledMessage).to.eql(
+      TransactionMessage.decompile(message, {
+        accountKeysFromLookups: accountKeys.accountKeysFromLookups!,
+      }),
+    );
+  });
+});