diff --git a/web3.js/src/message/compiled-keys.ts b/web3.js/src/message/compiled-keys.ts new file mode 100644 index 00000000000000..a0cf88ea34377a --- /dev/null +++ b/web3.js/src/message/compiled-keys.ts @@ -0,0 +1,165 @@ +import {MessageHeader, MessageAddressTableLookup} from './index'; +import {AccountKeysFromLookups} from './account-keys'; +import {AddressLookupTableAccount} from '../programs'; +import {TransactionInstruction} from '../transaction'; +import assert from '../utils/assert'; +import {PublicKey} from '../publickey'; + +export type CompiledKeyMeta = { + isSigner: boolean; + isWritable: boolean; + isInvoked: boolean; +}; + +type KeyMetaMap = Map; + +export class CompiledKeys { + payer: PublicKey; + keyMetaMap: KeyMetaMap; + + constructor(payer: PublicKey, keyMetaMap: KeyMetaMap) { + this.payer = payer; + this.keyMetaMap = keyMetaMap; + } + + static compile( + instructions: Array, + payer: PublicKey, + ): CompiledKeys { + const keyMetaMap: KeyMetaMap = new Map(); + const getOrInsertDefault = (pubkey: PublicKey): CompiledKeyMeta => { + const address = pubkey.toBase58(); + let keyMeta = keyMetaMap.get(address); + if (keyMeta === undefined) { + keyMeta = { + isSigner: false, + isWritable: false, + isInvoked: false, + }; + keyMetaMap.set(address, keyMeta); + } + return keyMeta; + }; + + const payerKeyMeta = getOrInsertDefault(payer); + payerKeyMeta.isSigner = true; + payerKeyMeta.isWritable = true; + + for (const ix of instructions) { + getOrInsertDefault(ix.programId).isInvoked = true; + for (const accountMeta of ix.keys) { + const keyMeta = getOrInsertDefault(accountMeta.pubkey); + keyMeta.isSigner ||= accountMeta.isSigner; + keyMeta.isWritable ||= accountMeta.isWritable; + } + } + + return new CompiledKeys(payer, keyMetaMap); + } + + getMessageComponents(): [MessageHeader, Array] { + const mapEntries = [...this.keyMetaMap.entries()]; + assert(mapEntries.length <= 256, 'Max static account keys length exceeded'); + + const writableSigners = mapEntries.filter( + ([, meta]) => meta.isSigner && meta.isWritable, + ); + const readonlySigners = mapEntries.filter( + ([, meta]) => meta.isSigner && !meta.isWritable, + ); + const writableNonSigners = mapEntries.filter( + ([, meta]) => !meta.isSigner && meta.isWritable, + ); + const readonlyNonSigners = mapEntries.filter( + ([, meta]) => !meta.isSigner && !meta.isWritable, + ); + + const header: MessageHeader = { + numRequiredSignatures: writableSigners.length + readonlySigners.length, + numReadonlySignedAccounts: readonlySigners.length, + numReadonlyUnsignedAccounts: readonlyNonSigners.length, + }; + + // sanity checks + { + assert( + writableSigners.length > 0, + 'Expected at least one writable signer key', + ); + const [payerAddress] = writableSigners[0]; + assert( + payerAddress === this.payer.toBase58(), + 'Expected first writable signer key to be the fee payer', + ); + } + + const staticAccountKeys = [ + ...writableSigners.map(([address]) => new PublicKey(address)), + ...readonlySigners.map(([address]) => new PublicKey(address)), + ...writableNonSigners.map(([address]) => new PublicKey(address)), + ...readonlyNonSigners.map(([address]) => new PublicKey(address)), + ]; + + return [header, staticAccountKeys]; + } + + extractTableLookup( + lookupTable: AddressLookupTableAccount, + ): [MessageAddressTableLookup, AccountKeysFromLookups] | undefined { + const [writableIndexes, drainedWritableKeys] = + this.drainKeysFoundInLookupTable( + lookupTable.state.addresses, + keyMeta => + !keyMeta.isSigner && !keyMeta.isInvoked && keyMeta.isWritable, + ); + const [readonlyIndexes, drainedReadonlyKeys] = + this.drainKeysFoundInLookupTable( + lookupTable.state.addresses, + keyMeta => + !keyMeta.isSigner && !keyMeta.isInvoked && !keyMeta.isWritable, + ); + + // Don't extract lookup if no keys were found + if (writableIndexes.length === 0 && readonlyIndexes.length === 0) { + return; + } + + return [ + { + accountKey: lookupTable.key, + writableIndexes, + readonlyIndexes, + }, + { + writable: drainedWritableKeys, + readonly: drainedReadonlyKeys, + }, + ]; + } + + /** @internal */ + private drainKeysFoundInLookupTable( + lookupTableEntries: Array, + keyMetaFilter: (keyMeta: CompiledKeyMeta) => boolean, + ): [Array, Array] { + const lookupTableIndexes = new Array(); + const drainedKeys = new Array(); + + for (const [address, keyMeta] of this.keyMetaMap.entries()) { + if (keyMetaFilter(keyMeta)) { + const key = new PublicKey(address); + const lookupTableIndex = lookupTableEntries.findIndex(entry => + entry.equals(key), + ); + if (lookupTableIndex >= 0) { + assert(lookupTableIndex < 256, 'Max lookup table index exceeded'); + lookupTableIndexes.push(lookupTableIndex); + drainedKeys.push(key); + this.keyMetaMap.delete(address); + } + } + } + + return [lookupTableIndexes, drainedKeys]; + } +} diff --git a/web3.js/src/message/index.ts b/web3.js/src/message/index.ts index 5389a89d251194..294a90b17e25ee 100644 --- a/web3.js/src/message/index.ts +++ b/web3.js/src/message/index.ts @@ -1,6 +1,7 @@ import {PublicKey} from '../publickey'; export * from './account-keys'; +// note: compiled-keys is internal and doesn't need to be exported export * from './legacy'; export * from './versioned'; export * from './v0'; diff --git a/web3.js/test/message-tests/compiled-keys.test.ts b/web3.js/test/message-tests/compiled-keys.test.ts new file mode 100644 index 00000000000000..a44f7bdd6cfcb6 --- /dev/null +++ b/web3.js/test/message-tests/compiled-keys.test.ts @@ -0,0 +1,243 @@ +import {expect} from 'chai'; + +import {CompiledKeyMeta, CompiledKeys} from '../../src/message/compiled-keys'; +import {AddressLookupTableAccount} from '../../src/programs'; +import {PublicKey} from '../../src/publickey'; +import {AccountMeta, TransactionInstruction} from '../../src/transaction'; + +function createTestKeys(count: number): Array { + return new Array(count).fill(0).map(() => PublicKey.unique()); +} + +function createTestLookupTable( + addresses: Array, +): 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('CompiledKeys', () => { + it('compile', () => { + const payer = PublicKey.unique(); + const keys = createTestKeys(4); + const programIds = createTestKeys(4); + const compiledKeys = CompiledKeys.compile( + [ + new TransactionInstruction({ + programId: programIds[0], + keys: [ + createAccountMeta(keys[0], false, false), + createAccountMeta(keys[1], true, false), + createAccountMeta(keys[2], false, true), + createAccountMeta(keys[3], true, true), + // duplicate the account metas + createAccountMeta(keys[0], false, false), + createAccountMeta(keys[1], true, false), + createAccountMeta(keys[2], false, true), + createAccountMeta(keys[3], true, true), + // reference program ids + createAccountMeta(programIds[0], false, false), + createAccountMeta(programIds[1], true, false), + createAccountMeta(programIds[2], false, true), + createAccountMeta(programIds[3], true, true), + ], + }), + new TransactionInstruction({programId: programIds[1], keys: []}), + new TransactionInstruction({programId: programIds[2], keys: []}), + new TransactionInstruction({programId: programIds[3], keys: []}), + ], + payer, + ); + + const map = new Map(); + setMapEntry(map, payer, true, true, false); + setMapEntry(map, keys[0], false, false, false); + setMapEntry(map, keys[1], true, false, false); + setMapEntry(map, keys[2], false, true, false); + setMapEntry(map, keys[3], true, true, false); + setMapEntry(map, programIds[0], false, false, true); + setMapEntry(map, programIds[1], true, false, true); + setMapEntry(map, programIds[2], false, true, true); + setMapEntry(map, programIds[3], true, true, true); + expect(compiledKeys.keyMetaMap).to.eql(map); + expect(compiledKeys.payer).to.eq(payer); + }); + + it('compile with dup payer', () => { + const [payer, programId] = createTestKeys(2); + const compiledKeys = CompiledKeys.compile( + [ + new TransactionInstruction({ + programId: programId, + keys: [createAccountMeta(payer, false, false)], + }), + ], + payer, + ); + + const map = new Map(); + setMapEntry(map, payer, true, true, false); + setMapEntry(map, programId, false, false, true); + expect(compiledKeys.keyMetaMap).to.eql(map); + expect(compiledKeys.payer).to.eq(payer); + }); + + it('compile with dup key', () => { + const [payer, key, programId] = createTestKeys(3); + const compiledKeys = CompiledKeys.compile( + [ + new TransactionInstruction({ + programId: programId, + keys: [ + createAccountMeta(key, false, false), + createAccountMeta(key, true, true), + ], + }), + ], + payer, + ); + + const map = new Map(); + setMapEntry(map, payer, true, true, false); + setMapEntry(map, key, true, true, false); + setMapEntry(map, programId, false, false, true); + expect(compiledKeys.keyMetaMap).to.eql(map); + expect(compiledKeys.payer).to.eq(payer); + }); + + it('getMessageComponents', () => { + const keys = createTestKeys(4); + const payer = keys[0]; + const map = new Map(); + setMapEntry(map, payer, true, true, false); + setMapEntry(map, keys[1], true, false, false); + setMapEntry(map, keys[2], false, true, false); + setMapEntry(map, keys[3], false, false, false); + const compiledKeys = new CompiledKeys(payer, map); + const [header, staticAccountKeys] = compiledKeys.getMessageComponents(); + expect(staticAccountKeys).to.eql(keys); + expect(header).to.eql({ + numRequiredSignatures: 2, + numReadonlySignedAccounts: 1, + numReadonlyUnsignedAccounts: 1, + }); + }); + + it('getMessageComponents with overflow', () => { + const keys = createTestKeys(257); + const map = new Map(); + for (const key of keys) { + setMapEntry(map, key, true, true, false); + } + const compiledKeys = new CompiledKeys(keys[0], map); + expect(() => compiledKeys.getMessageComponents()).to.throw( + 'Max static account keys length exceeded', + ); + }); + + it('extractTableLookup', () => { + const keys = createTestKeys(6); + const map = new Map(); + setMapEntry(map, keys[0], true, true, false); + setMapEntry(map, keys[1], true, false, false); + setMapEntry(map, keys[2], false, true, false); + setMapEntry(map, keys[3], false, false, false); + setMapEntry(map, keys[4], true, false, true); + setMapEntry(map, keys[5], false, false, true); + + const lookupTable = createTestLookupTable([...keys, ...keys]); + const compiledKeys = new CompiledKeys(keys[0], map); + const extractResult = compiledKeys.extractTableLookup(lookupTable); + if (extractResult === undefined) { + expect(extractResult).to.not.be.undefined; + return; + } + + const [tableLookup, extractedAddresses] = extractResult; + expect(tableLookup).to.eql({ + accountKey: lookupTable.key, + writableIndexes: [2], + readonlyIndexes: [3], + }); + expect(extractedAddresses).to.eql({ + writable: [keys[2]], + readonly: [keys[3]], + }); + }); + + it('extractTableLookup no extractable keys found', () => { + const keys = createTestKeys(6); + const map = new Map(); + setMapEntry(map, keys[0], true, true, false); + setMapEntry(map, keys[1], true, false, false); + setMapEntry(map, keys[2], true, true, true); + setMapEntry(map, keys[3], true, false, true); + setMapEntry(map, keys[4], false, true, true); + setMapEntry(map, keys[5], false, false, true); + + const lookupTable = createTestLookupTable(keys); + const compiledKeys = new CompiledKeys(keys[0], map); + const extractResult = compiledKeys.extractTableLookup(lookupTable); + expect(extractResult).to.be.undefined; + }); + + it('extractTableLookup with empty lookup table', () => { + const keys = createTestKeys(2); + const map = new Map(); + setMapEntry(map, keys[0], true, true, false); + setMapEntry(map, keys[1], false, false, false); + + const lookupTable = createTestLookupTable([]); + const compiledKeys = new CompiledKeys(keys[0], map); + const extractResult = compiledKeys.extractTableLookup(lookupTable); + expect(extractResult).to.be.undefined; + }); + + it('extractTableLookup with invalid lookup table', () => { + const keys = createTestKeys(257); + const map = new Map(); + setMapEntry(map, keys[0], true, true, false); + setMapEntry(map, keys[256], false, false, false); + + const lookupTable = createTestLookupTable(keys); + const compiledKeys = new CompiledKeys(keys[0], map); + expect(() => compiledKeys.extractTableLookup(lookupTable)).to.throw( + 'Max lookup table index exceeded', + ); + }); +}); + +function setMapEntry( + map: Map, + pubkey: PublicKey, + isSigner: boolean, + isWritable: boolean, + isInvoked: boolean, +) { + map.set(pubkey.toBase58(), { + isSigner, + isWritable, + isInvoked, + }); +} + +function createAccountMeta( + pubkey: PublicKey, + isSigner: boolean, + isWritable: boolean, +): AccountMeta { + return { + pubkey, + isSigner, + isWritable, + }; +}