STREAM-1417: support token2022 in distributor (#156)

* STREAM-1417: support token2022 in distributor

lerna.json

@@ -2,6 +2,6 @@
   "packages": [
     "packages/*"
   ],
-  "version": "6.0.1",
+  "version": "6.0.2",
   "$schema": "node_modules/lerna/schemas/lerna-schema.json"
 }

packages/common/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@streamflow/common",
-  "version": "6.0.1",
+  "version": "6.0.2",
   "description": "Common utilities and types used by streamflow packages.",
   "homepage": "https://github.com/streamflow-finance/js-sdk/",
   "main": "dist/index.js",

packages/common/solana/types.ts

@@ -21,4 +21,5 @@ export interface CheckAssociatedTokenAccountsData {
 export interface AtaParams {
   mint: PublicKey;
   owner: PublicKey;
+  programId?: PublicKey;
 }

packages/common/solana/utils.ts

@@ -1,4 +1,11 @@
-import { createAssociatedTokenAccountInstruction, getAssociatedTokenAddress } from "@solana/spl-token";
+import {
+  createAssociatedTokenAccountInstruction,
+  getAssociatedTokenAddress,
+  unpackMint,
+  Mint,
+  TOKEN_PROGRAM_ID,
+  TOKEN_2022_PROGRAM_ID,
+} from "@solana/spl-token";
 import { SignerWalletAdapter } from "@solana/wallet-adapter-base";
 import {
   BlockhashWithExpiryBlockHeight,
@@ -146,29 +153,47 @@ export async function signAndExecuteTransaction(
  * Shorthand call signature for getAssociatedTokenAddress, with allowance for address to be offCurve
  * @param {PublicKey} mint - SPL token Mint address.
  * @param {PublicKey} owner - Owner of the Associated Token Address
+ * @param {PublicKey} programId - Program ID of the mint
  * @return {Promise<PublicKey>} - Associated Token Address
  */
-export function ata(mint: PublicKey, owner: PublicKey): Promise<PublicKey> {
-  return getAssociatedTokenAddress(mint, owner, true);
+export function ata(mint: PublicKey, owner: PublicKey, programId?: PublicKey): Promise<PublicKey> {
+  return getAssociatedTokenAddress(mint, owner, true, programId);
 }
 
 /**
  * Function that checks whether ATA exists for each provided owner
  * @param connection - Solana client connection
- * @param paramsBatch - Array of Params for an each ATA account: {mint, owner}
- * @returns Array of boolean where each members corresponds to owners member
+ * @param paramsBatch - Array of Params for each ATA account: {mint, owner}
+ * @returns Array of boolean where each member corresponds to an owner
  */
 export async function ataBatchExist(connection: Connection, paramsBatch: AtaParams[]): Promise<boolean[]> {
   const tokenAccounts = await Promise.all(
-    paramsBatch.map(async ({ mint, owner }) => {
-      const pubkey = await ata(mint, owner);
-      return pubkey;
+    paramsBatch.map(async ({ mint, owner, programId }) => {
+      return ata(mint, owner, programId);
     }),
   );
   const response = await connection.getMultipleAccountsInfo(tokenAccounts);
   return response.map((accInfo) => !!accInfo);
 }
 
+export async function enrichAtaParams(connection: Connection, paramsBatch: AtaParams[]): Promise<AtaParams[]> {
+  const programIdByMint: { [key: string]: PublicKey } = {};
+  return Promise.all(
+    paramsBatch.map(async (params) => {
+      if (params.programId) {
+        return params;
+      }
+      const mintStr = params.mint.toString();
+      if (!(mintStr in programIdByMint)) {
+        const { programId } = await getMintAndProgram(connection, params.mint);
+        programIdByMint[mintStr] = programId;
+      }
+      params.programId = programIdByMint[mintStr];
+      return params;
+    }),
+  );
+}
+
 /**
  * Generates a Transaction to create ATA for an array of owners
  * @param connection - Solana client connection
@@ -184,9 +209,10 @@ export async function generateCreateAtaBatchTx(
   tx: Transaction;
   hash: BlockhashWithExpiryBlockHeight;
 }> {
+  paramsBatch = await enrichAtaParams(connection, paramsBatch);
   const ixs: TransactionInstruction[] = await Promise.all(
-    paramsBatch.map(async ({ mint, owner }) => {
-      return createAssociatedTokenAccountInstruction(payer, await ata(mint, owner), owner, mint);
+    paramsBatch.map(async ({ mint, owner, programId }) => {
+      return createAssociatedTokenAccountInstruction(payer, await ata(mint, owner), owner, mint, programId);
     }),
   );
   const hash = await connection.getLatestBlockhash();
@@ -210,7 +236,11 @@ export async function createAtaBatch(
   invoker: Keypair | SignerWalletAdapter,
   paramsBatch: AtaParams[],
 ): Promise<string> {
-  const { tx, hash } = await generateCreateAtaBatchTx(connection, invoker.publicKey!, paramsBatch);
+  const { tx, hash } = await generateCreateAtaBatchTx(
+    connection,
+    invoker.publicKey!,
+    await enrichAtaParams(connection, paramsBatch),
+  );
   return signAndExecuteTransaction(connection, invoker, tx, hash);
 }
 
@@ -220,24 +250,26 @@ export async function createAtaBatch(
  * @param owners - Array of ATA owners
  * @param mint - Mint for which ATA will be checked
  * @param invoker - Transaction invoker and payer
+ * @param programId - Program ID of the Mint
  * @returns Array of Transaction Instructions that should be added to a transaction
  */
 export async function checkOrCreateAtaBatch(
   connection: Connection,
   owners: PublicKey[],
   mint: PublicKey,
   invoker: SignerWalletAdapter | Keypair,
+  programId?: PublicKey,
 ): Promise<TransactionInstruction[]> {
   const ixs: TransactionInstruction[] = [];
   // TODO: optimize fetching and maps/arrays
   const atas: PublicKey[] = [];
   for (const owner of owners) {
-    atas.push(await ata(mint, owner));
+    atas.push(await ata(mint, owner, programId));
   }
   const response = await connection.getMultipleAccountsInfo(atas);
   for (let i = 0; i < response.length; i++) {
     if (!response[i]) {
-      ixs.push(createAssociatedTokenAccountInstruction(invoker.publicKey!, atas[i], owners[i], mint));
+      ixs.push(createAssociatedTokenAccountInstruction(invoker.publicKey!, atas[i], owners[i], mint, programId));
     }
   }
   return ixs;
@@ -263,3 +295,28 @@ export function prepareBaseInstructions(
 
   return ixs;
 }
+
+/**
+ * Retrieve information about a mint and its program ID, support all Token Programs.
+ *
+ * @param connection Connection to use
+ * @param address    Mint account
+ * @param commitment Desired level of commitment for querying the state
+ *
+ * @return Mint information
+ */
+export async function getMintAndProgram(
+  connection: Connection,
+  address: PublicKey,
+  commitment?: Commitment,
+): Promise<{ mint: Mint; programId: PublicKey }> {
+  const accountInfo = await connection.getAccountInfo(address, commitment);
+  let programId = accountInfo?.owner;
+  if (!programId?.equals(TOKEN_PROGRAM_ID) && !programId?.equals(TOKEN_2022_PROGRAM_ID)) {
+    programId = TOKEN_PROGRAM_ID;
+  }
+  return {
+    mint: unpackMint(address, accountInfo, programId),
+    programId: programId!,
+  };
+}

packages/distributor/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@streamflow/distributor",
-  "version": "6.0.1",
+  "version": "6.0.2",
   "description": "JavaScript SDK to interact with Streamflow Airdrop protocol.",
   "homepage": "https://github.com/streamflow-finance/js-sdk/",
   "main": "dist/index.js",

packages/distributor/solana/client.ts

@@ -1,11 +1,5 @@
 import BN from "bn.js";
-import {
-  ASSOCIATED_TOKEN_PROGRAM_ID,
-  NATIVE_MINT,
-  TOKEN_PROGRAM_ID,
-  createTransferCheckedInstruction,
-  getMint,
-} from "@solana/spl-token";
+import { ASSOCIATED_TOKEN_PROGRAM_ID, NATIVE_MINT, createTransferCheckedInstruction } from "@solana/spl-token";
 import {
   Connection,
   PublicKey,
@@ -21,6 +15,7 @@ import {
   prepareWrappedAccount,
   prepareBaseInstructions,
   prepareTransaction,
+  getMintAndProgram,
 } from "@streamflow/common/solana";
 
 import { DISTRIBUTOR_PROGRAM_ID } from "./constants";
@@ -86,10 +81,10 @@ export default class SolanaDistributorClient {
 
     const ixs: TransactionInstruction[] = prepareBaseInstructions(this.connection, { computePrice, computeLimit });
     const mint = isNative ? NATIVE_MINT : new PublicKey(data.mint);
-    const mintAccount = await getMint(this.connection, mint);
+    const { mint: mintAccount, programId } = await getMintAndProgram(this.connection, mint);
     const distributorPublicKey = getDistributorPda(this.programId, mint, data.version);
-    const tokenVault = await ata(mint, distributorPublicKey);
-    const senderTokens = await ata(mint, invoker.publicKey);
+    const tokenVault = await ata(mint, distributorPublicKey, programId);
+    const senderTokens = await ata(mint, invoker.publicKey, programId);
 
     const args: NewDistributorArgs = {
       version: new BN(data.version, 10),
@@ -110,7 +105,7 @@ export default class SolanaDistributorClient {
       admin: invoker.publicKey,
       systemProgram: SystemProgram.programId,
       associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
-      tokenProgram: TOKEN_PROGRAM_ID,
+      tokenProgram: programId,
     };
 
     if (isNative) {
@@ -133,6 +128,8 @@ export default class SolanaDistributorClient {
         invoker.publicKey,
         BigInt(data.maxTotalClaim.toString()),
         mintAccount.decimals,
+        undefined,
+        programId,
       ),
     );
 
@@ -161,9 +158,12 @@ export default class SolanaDistributorClient {
       throw new Error("Couldn't get account info");
     }
 
+    const { programId } = await getMintAndProgram(this.connection, distributor.mint);
     const ixs: TransactionInstruction[] = prepareBaseInstructions(this.connection, { computePrice, computeLimit });
-    ixs.push(...(await checkOrCreateAtaBatch(this.connection, [invoker.publicKey], distributor.mint, invoker)));
-    const invokerTokens = await ata(distributor.mint, invoker.publicKey);
+    ixs.push(
+      ...(await checkOrCreateAtaBatch(this.connection, [invoker.publicKey], distributor.mint, invoker, programId)),
+    );
+    const invokerTokens = await ata(distributor.mint, invoker.publicKey, programId);
     const claimStatusPublicKey = getClaimantStatusPda(this.programId, distributorPublicKey, invoker.publicKey);
     const claimStatus = await ClaimStatus.fetch(this.connection, claimStatusPublicKey);
 
@@ -174,7 +174,7 @@ export default class SolanaDistributorClient {
       to: invokerTokens,
       claimant: invoker.publicKey,
       mint: distributor.mint,
-      tokenProgram: TOKEN_PROGRAM_ID,
+      tokenProgram: programId,
       systemProgram: SystemProgram.programId,
     };
 
@@ -213,16 +213,19 @@ export default class SolanaDistributorClient {
       throw new Error("Couldn't get account info");
     }
 
+    const { programId } = await getMintAndProgram(this.connection, distributor.mint);
     const ixs: TransactionInstruction[] = prepareBaseInstructions(this.connection, { computePrice, computeLimit });
-    ixs.push(...(await checkOrCreateAtaBatch(this.connection, [invoker.publicKey], distributor.mint, invoker)));
+    ixs.push(
+      ...(await checkOrCreateAtaBatch(this.connection, [invoker.publicKey], distributor.mint, invoker, programId)),
+    );
     const accounts: ClawbackAccounts = {
       distributor: distributorPublicKey,
       from: distributor.tokenVault,
       to: distributor.clawbackReceiver,
       admin: invoker.publicKey,
       mint: distributor.mint,
       systemProgram: SystemProgram.programId,
-      tokenProgram: TOKEN_PROGRAM_ID,
+      tokenProgram: programId,
     };
 
     ixs.push(clawback(accounts, this.programId));

packages/stream/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@streamflow/stream",
-  "version": "6.0.1",
+  "version": "6.0.2",
   "description": "JavaScript SDK to interact with Streamflow protocol.",
   "homepage": "https://github.com/streamflow-finance/js-sdk/",
   "main": "dist/index.js",