[Orca] Create token account for referral address

packages/swapper/src/orca/fee.ts

@@ -1,10 +1,7 @@
 import { AssetId, Quote } from "@gemwallet/types";
-import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
 import { PublicKey } from "@solana/web3.js";
 import BN from "bn.js";
 
-import { parsePublicKey } from "../chain/solana/account";
-
 export const BASIS_POINTS_DENOMINATOR = 10_000;
 export const MAX_SAFE_NUMBER_BN = new BN(Number.MAX_SAFE_INTEGER.toString());
 
@@ -35,12 +32,7 @@
     return value.toNumber();
 }
 
-export async function applyReferralFee(
-    asset: AssetId,
-    amountIn: bigint,
-    referralBps: bigint,
-    resolveProgram: (mint: PublicKey) => Promise<PublicKey>,
-): Promise<bigint> {
+export async function applyReferralFee(asset: AssetId, amountIn: bigint, referralBps: bigint): Promise<bigint> {
     if (referralBps <= BigInt(0)) {
         return amountIn;
     }
@@ -50,13 +42,6 @@
         return amountIn - referralFee;
     }
 
-    const tokenId = asset.getTokenId();
-    const mintKey = parsePublicKey(tokenId);
-    const programId = await resolveProgram(mintKey);
-    if (!programId.equals(TOKEN_PROGRAM_ID)) {
-        return amountIn;
-    }
-
     const referralFee = (amountIn * referralBps) / BigInt(10_000);
     return amountIn - referralFee;
 }

packages/swapper/src/orca/provider.test.ts

@@ -1,6 +1,6 @@
 import { Quote } from "@gemwallet/types";
-import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import { ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_PROGRAM_ID, TOKEN_2022_PROGRAM_ID } from "@solana/spl-token";
 import { TransactionInstruction, PublicKey } from "@solana/web3.js";
 
 import { SOL_ASSET, buildOrcaQuoteFixture, createOrcaQuoteRequest } from "../testkit/mock";
 import { OrcaWhirlpoolProvider } from "./provider";
@@ -17,6 +17,7 @@
 
 const TEST_LEGACY_MINT = "So11111111111111111111111111111111111111112";
 const TEST_TOKEN2022_MINT = "2b1kV6DkPAnxd5ixfnxCpjxmKwqjjaYmCZfHsFu24GXo";
+const TOKEN2022_PROGRAM_ADDRESS = "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb";
 
 function createQuote(mint: string, referralBps = 100): Quote {
     return buildOrcaQuoteFixture(
@@ -39,7 +40,7 @@
     );
 }
 
-describe("OrcaWhirlpoolProvider.buildReferralInstruction", () => {
+describe("OrcaWhirlpoolProvider.buildReferralInstructions", () => {
     let provider: OrcaWhirlpoolProvider;
     const userKey = new PublicKey("9iqKg7nZFkC6xhnoWvyvCSdrgSX1uxPxL4X4fb97aotW");
 
@@ -49,43 +50,47 @@
         mockFetchAllMint.fetchAllMint?.mockReset?.();
     });
 
-    it("returns null when referral amount is zero", async () => {
+    it("returns empty array when referral amount is zero", async () => {
         const quote = createQuote(TEST_LEGACY_MINT, 0);
 
         // @ts-expect-error accessing private method for test purposes
-        const instruction = await provider.buildReferralInstruction(quote, userKey);
+        const instructions = await provider.buildReferralInstructions(quote, userKey);
 
-        expect(instruction).toBeNull();
+        expect(instructions).toEqual([]);
     });
 
-    it("skips referral when mint uses token-2022 program", async () => {
+    it("builds create-ata and transfer instructions for legacy token program", async () => {
         mockFetchAllMint.fetchAllMint.mockResolvedValue([
             {
                 exists: true,
-                programAddress: "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb",
+                programAddress: TOKEN_PROGRAM_ID.toBase58(),
             },
         ]);
-        const quote = createQuote(TEST_TOKEN2022_MINT);
+        const quote = createQuote(TEST_LEGACY_MINT);
 
         // @ts-expect-error accessing private method for test purposes
-        const instruction = await provider.buildReferralInstruction(quote, userKey);
+        const instructions = await provider.buildReferralInstructions(quote, userKey);
 
-        expect(instruction).toBeNull();
+        expect(instructions).toHaveLength(2);
+        expect(instructions[0].programId).toEqual(ASSOCIATED_TOKEN_PROGRAM_ID);
+        expect(instructions[1].programId).toEqual(TOKEN_PROGRAM_ID);
     });
 
-    it("builds transfer instruction for legacy token program", async () => {
+    it("builds create-ata and transfer instructions for token-2022 program", async () => {
         mockFetchAllMint.fetchAllMint.mockResolvedValue([
             {
                 exists: true,
-                programAddress: TOKEN_PROGRAM_ID.toBase58(),
+                programAddress: TOKEN2022_PROGRAM_ADDRESS,
             },
         ]);
-        const quote = createQuote(TEST_LEGACY_MINT);
+        const quote = createQuote(TEST_TOKEN2022_MINT);
 
         // @ts-expect-error accessing private method for test purposes
-        const instruction = await provider.buildReferralInstruction(quote, userKey);
+        const instructions = await provider.buildReferralInstructions(quote, userKey);
 
-        expect(instruction).toBeInstanceOf(TransactionInstruction);
+        expect(instructions).toHaveLength(2);
+        expect(instructions[0].programId).toEqual(ASSOCIATED_TOKEN_PROGRAM_ID);
+        expect(instructions[1].programId).toEqual(TOKEN_2022_PROGRAM_ID);
     });
 });
 
@@ -118,7 +123,6 @@
 
     it("reduces swap amount for legacy SPL tokens when referral applies", async () => {
         jest.spyOn(provider as any, "findBestPool").mockResolvedValue({ account: { address: "PoolAddress" } });
-        const mintProgramSpy = jest.spyOn(provider as any, "getTokenProgram").mockResolvedValueOnce(TOKEN_PROGRAM_ID);
 
         const quote = await provider.get_quote(
             createOrcaQuoteRequest({
@@ -134,14 +138,10 @@
 
         expect(capturedAmount).toBe(BigInt(990000));
         expect(quote.route_data).toMatchObject({ amount: "990000" });
-        expect(mintProgramSpy).toHaveBeenCalledWith(new PublicKey(TEST_LEGACY_MINT));
     });
 
-    it("uses full input amount for Token-2022 tokens when referral cannot be collected", async () => {
+    it("reduces swap amount for Token-2022 tokens when referral applies", async () => {
         jest.spyOn(provider as any, "findBestPool").mockResolvedValue({ account: { address: "PoolAddress" } });
-        const getProgramSpy = jest
-            .spyOn(provider as any, "getTokenProgram")
-            .mockResolvedValueOnce(new PublicKey("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"));
 
         const quote = await provider.get_quote(
             createOrcaQuoteRequest({
@@ -155,8 +155,7 @@
             }),
         );
 
-        expect(capturedAmount).toBe(BigInt(1000000));
-        expect(quote.route_data).toMatchObject({ amount: "1000000" });
-        expect(getProgramSpy).toHaveBeenCalledWith(new PublicKey(TEST_TOKEN2022_MINT));
+        expect(capturedAmount).toBe(BigInt(990000));
+        expect(quote.route_data).toMatchObject({ amount: "990000" });
     });
 });

packages/swapper/src/orca/provider.ts

@@ -25,7 +25,11 @@ import {
 } from "@orca-so/whirlpools-core";
 import { AccountRole, type AccountLookupMeta, type AccountMeta, type Instruction } from "@solana/instructions";
 import { createSolanaRpc, address as toAddress, type Account, type Address, type TransactionSigner } from "@solana/kit";
-import { createTransferInstruction, getAssociatedTokenAddressSync, TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import {
+    createAssociatedTokenAccountIdempotentInstruction,
+    createTransferCheckedInstruction,
+    getAssociatedTokenAddressSync,
+} from "@solana/spl-token";
 import { Connection, PublicKey, SystemProgram, Transaction, TransactionInstruction } from "@solana/web3.js";
 
 import { BigIntMath } from "../bigint_math";
@@ -110,9 +114,7 @@ export class OrcaWhirlpoolProvider implements Protocol {
 
         const amountIn = BigIntMath.parseString(quoteRequest.from_value);
         const referralBps = BigInt(quoteRequest.referral_bps ?? 0);
-        const swapAmount = await applyReferralFee(fromAsset, amountIn, referralBps, (mint) =>
-            this.getTokenProgram(mint),
-        );
+        const swapAmount = await applyReferralFee(fromAsset, amountIn, referralBps);
 
         if (swapAmount <= BigInt(0)) {
             throw new Error("Swap amount must be greater than zero");
@@ -153,7 +155,7 @@ export class OrcaWhirlpoolProvider implements Protocol {
 
         const signer = this.createPassthroughSigner(userPublicKey);
 
-        const [{ instructions }, priorityFee, { blockhash, lastValidBlockHeight }, referralInstruction] =
+        const [{ instructions }, priorityFee, { blockhash, lastValidBlockHeight }, referralInstructions] =
             await Promise.all([
                 swapInstructions(
                     this.rpc,
@@ -164,18 +166,14 @@ export class OrcaWhirlpoolProvider implements Protocol {
                 ),
                 this.getPriorityFee(),
                 getRecentBlockhash(this.connection, DEFAULT_COMMITMENT),
-                this.buildReferralInstruction(quote, userPublicKey),
+                this.buildReferralInstructions(quote, userPublicKey),
             ]);
 
         const legacyInstructions = instructions.map((instruction) => this.toLegacyInstruction(instruction));
         const computeBudgetInstructions = addComputeBudgetInstructions([], DEFAULT_COMPUTE_UNIT_LIMIT, priorityFee);
 
         const transaction = new Transaction();
-        transaction.add(...computeBudgetInstructions, ...legacyInstructions);
-
-        if (referralInstruction) {
-            transaction.add(referralInstruction);
-        }
+        transaction.add(...computeBudgetInstructions, ...legacyInstructions, ...referralInstructions);
 
         transaction.feePayer = userPublicKey;
         setTransactionBlockhash(transaction, blockhash, lastValidBlockHeight);
@@ -453,13 +451,13 @@ export class OrcaWhirlpoolProvider implements Protocol {
         };
     }
 
-    private async buildReferralInstruction(
+    private async buildReferralInstructions(
         quote: Quote,
         userPublicKey: PublicKey,
-    ): Promise<TransactionInstruction | null> {
+    ): Promise<TransactionInstruction[]> {
         const referralAmount = calculateReferralFeeAmount(quote);
         if (referralAmount.isZero()) {
-            return null;
+            return [];
         }
 
         const referrer = getReferrerAddresses().solana;
@@ -469,34 +467,40 @@ export class OrcaWhirlpoolProvider implements Protocol {
 
         const fromAsset = AssetId.fromString(quote.quote.from_asset.id);
         if (fromAsset.isNative()) {
-            return SystemProgram.transfer({
-                fromPubkey: userPublicKey,
-                toPubkey: new PublicKey(referrer),
-                lamports: bnToNumberSafe(referralAmount),
-            });
+            return [
+                SystemProgram.transfer({
+                    fromPubkey: userPublicKey,
+                    toPubkey: new PublicKey(referrer),
+                    lamports: bnToNumberSafe(referralAmount),
+                }),
+            ];
         }
 
         const tokenId = fromAsset.getTokenId();
         const fromMintKey = parsePublicKey(tokenId);
         const programId = await this.getTokenProgram(fromMintKey);
-        if (!programId.equals(TOKEN_PROGRAM_ID)) {
-            return null;
-        }
+        const referrerPublicKey = new PublicKey(referrer);
         const userTokenAccount = getAssociatedTokenAddressSync(fromMintKey, userPublicKey, false, programId);
-        const referrerTokenAccount = getAssociatedTokenAddressSync(
-            fromMintKey,
-            new PublicKey(referrer),
-            false,
-            programId,
-        );
-
-        return createTransferInstruction(
-            userTokenAccount,
-            referrerTokenAccount,
-            userPublicKey,
-            bnToNumberSafe(referralAmount),
-            [],
-            programId,
-        );
+        const referrerTokenAccount = getAssociatedTokenAddressSync(fromMintKey, referrerPublicKey, false, programId);
+
+        return [
+            createAssociatedTokenAccountIdempotentInstruction(
+                userPublicKey,
+                referrerTokenAccount,
+                referrerPublicKey,
+                fromMintKey,
+                programId,
+            ),
+            createTransferCheckedInstruction(
+                userTokenAccount,
+                fromMintKey,
+                referrerTokenAccount,
+                userPublicKey,
+                bnToNumberSafe(referralAmount),
+                quote.quote.from_asset.decimals,
+                [],
+                programId,
+            ),
+        ];
     }
 }