create ata on claim/clawback

streamflow-finance · Feb 13, 2024 · 7d0faab · 7d0faab
1 parent 1e43767
commit 7d0faab
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 36 deletions.
diff --git a/packages/stream/solana/StreamClient.ts b/packages/stream/solana/StreamClient.ts
@@ -15,7 +15,6 @@ import {
   Commitment,
   ConnectionConfig,
 } from "@solana/web3.js";
-import { createAssociatedTokenAccountInstruction } from "@solana/spl-token";
 import * as borsh from "borsh";
 
 import {
@@ -32,6 +31,7 @@ import {
 } from "./types";
 import {
   ata,
+  checkOrCreateAtaBatch,
   decodeStream,
   extractSolanaErrorCode,
   getProgramAccounts,
@@ -920,35 +920,16 @@ export default class SolanaStreamClient extends BaseStreamClient {
     if (!checkTokenAccounts) {
       return;
     }
-    const checkedKeys: Set<string> = new Set();
-    // TODO: optimize fetching and maps/arrays
-    const accountArrays = [
-      [data.sender, data.senderTokens],
-      [data.recipient, data.recipientTokens],
-      [data.partner, data.partnerTokens],
-      [data.streamflowTreasury, data.streamflowTreasuryTokens],
-    ].filter((value) => {
-      if (checkedKeys.has(value[1].toBase58())) {
-        return false;
-      }
-      checkedKeys.add(value[1].toBase58());
-      return true;
-    });
-    const response = await this.connection.getMultipleAccountsInfo(
-      accountArrays.map((item) => item[1])
+    const owners = Array.from(
+      new Set([
+        data.sender.toBase58(),
+        data.recipient.toBase58(),
+        data.partner.toBase58(),
+        data.streamflowTreasury.toBase58(),
+      ]),
+      (address) => new PublicKey(address)
     );
-    for (let i = 0; i < response.length; i++) {
-      if (!response[i]) {
-        ixs.push(
-          createAssociatedTokenAccountInstruction(
-            invoker.publicKey!,
-            accountArrays[i][1],
-            accountArrays[i][0],
-            data.mint
-          )
-        );
-      }
-    }
+    ixs.push(...(await checkOrCreateAtaBatch(this.connection, owners, data.mint, invoker)));
   }
 
   /**

diff --git a/packages/stream/solana/distributor/client.ts b/packages/stream/solana/distributor/client.ts
@@ -27,7 +27,7 @@ import {
   IGetDistributors,
 } from "./types";
 import { ICreateSolanaExt, IInteractStreamSolanaExt } from "../types";
-import { ata } from "../utils";
+import { ata, checkOrCreateAtaBatch } from "../utils";
 import {
   ClaimLockedAccounts,
   ClawbackAccounts,
@@ -151,14 +151,19 @@ export default class SolanaDistributorClient {
       throw new Error("Invoker's PublicKey is not available, check passed wallet adapter!");
     }
 
-    const ixs: TransactionInstruction[] = [];
     const distributorPublicKey = new PublicKey(data.id);
     const distributor = await MerkleDistributor.fetch(this.connection, distributorPublicKey);
 
     if (!distributor) {
       throw new Error("Couldn't get account info");
     }
 
+    const ixs = await checkOrCreateAtaBatch(
+      this.connection,
+      [invoker.publicKey],
+      distributor.mint,
+      invoker
+    );
     const invokerTokens = await ata(distributor.mint, invoker.publicKey);
     const claimStatusPublicKey = getClaimantStatusPda(
       this.programId,
@@ -210,14 +215,19 @@ export default class SolanaDistributorClient {
       throw new Error("Invoker's PublicKey is not available, check passed wallet adapter!");
     }
 
-    const ixs: TransactionInstruction[] = [];
     const distributorPublicKey = new PublicKey(data.id);
     const distributor = await MerkleDistributor.fetch(this.connection, distributorPublicKey);
 
     if (!distributor) {
       throw new Error("Couldn't get account info");
     }
 
+    const ixs = await checkOrCreateAtaBatch(
+      this.connection,
+      [invoker.publicKey],
+      distributor.mint,
+      invoker
+    );
     const accounts: ClawbackAccounts = {
       distributor: distributorPublicKey,
       from: distributor.tokenVault,

diff --git a/packages/stream/solana/types.ts b/packages/stream/solana/types.ts
@@ -432,13 +432,9 @@ export interface CreateMultipleStreamsValues {
 
 export interface CheckAssociatedTokenAccountsData {
   sender: PublicKey;
-  senderTokens: PublicKey;
   recipient: PublicKey;
-  recipientTokens: PublicKey;
   partner: PublicKey;
-  partnerTokens: PublicKey;
   streamflowTreasury: PublicKey;
-  streamflowTreasuryTokens: PublicKey;
   mint: PublicKey;
 }
 

diff --git a/packages/stream/solana/utils.ts b/packages/stream/solana/utils.ts
@@ -318,3 +318,34 @@ export function extractSolanaErrorCode(errorText: string): string | null {
   const errorCode = Number(match[1]);
   return SOLANA_ERROR_MAP[errorCode] || null;
 }
+
+/**
+ * Utility function that checks whether associated token accounts exist and return instructions to populate them if not
+ * @param connection - Solana client connection
+ * @param owners - Array of ATA owners
+ * @param mint - Mint for which ATA will be checked
+ * @param invoker - Transaction invoker and payer
+ * @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
+): 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));
+  }
+  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)
+      );
+    }
+  }
+  return ixs;
+}