Add modeIdempotent option to resolveOrCreateATA(s)

packages/common-sdk/README.md

@@ -1,2 +1,13 @@
 # Orca Common SDK
 This package contains a set of utility functions used by other Typescript components in Orca.
+
+## Run Typescript tests via local validator
+To startup local validator, run:
+```
+solana-test-validator
+```
+
+In the common-sdk folder, run:
+```
+npm run test
+```

packages/common-sdk/jest.config.js

@@ -1,7 +1,7 @@
 module.exports = {
   "roots": [
     "<rootDir>/src",
-    "<rootDir>/tests/sdk"
+    "<rootDir>/tests"
   ],
   "testMatch": [
     "**/__tests__/**/*.+(ts|tsx|js)",

packages/common-sdk/src/web3/ata-util.ts

@@ -5,7 +5,12 @@ import {
   TOKEN_PROGRAM_ID,
   u64,
 } from "@solana/spl-token";
-import { Connection, PublicKey } from "@solana/web3.js";
+import {
+  Connection,
+  PublicKey,
+  TransactionInstruction,
+  SystemProgram,
+} from "@solana/web3.js";
 import {
   createWSOLAccountInstructions,
   ResolvedTokenAddressInstruction,
@@ -23,6 +28,7 @@ import { EMPTY_INSTRUCTION } from "./transactions/types";
  * @param tokenMint Token mint address
  * @param wrappedSolAmountIn Optional. Only use for input/source token that could be SOL
  * @param payer Payer that would pay the rent for the creation of the ATAs
+ * @param modeIdempotent Optional. Use CreateIdempotent instruction instead of Create instruction
  * @returns
  */
 export async function resolveOrCreateATA(
@@ -31,14 +37,16 @@ export async function resolveOrCreateATA(
   tokenMint: PublicKey,
   getAccountRentExempt: () => Promise<number>,
   wrappedSolAmountIn = new u64(0),
-  payer = ownerAddress
+  payer = ownerAddress,
+  modeIdempotent: boolean = false,
 ): Promise<ResolvedTokenAddressInstruction> {
   const instructions = await resolveOrCreateATAs(
     connection,
     ownerAddress,
     [{ tokenMint, wrappedSolAmountIn }],
     getAccountRentExempt,
-    payer
+    payer,
+    modeIdempotent,
   );
   return instructions[0]!;
 }
@@ -58,14 +66,16 @@ type ResolvedTokenAddressRequest = {
  * @param tokenMint Token mint address
  * @param wrappedSolAmountIn Optional. Only use for input/source token that could be SOL
  * @param payer Payer that would pay the rent for the creation of the ATAs
+ * @param modeIdempotent Optional. Use CreateIdempotent instruction instead of Create instruction
  * @returns
  */
 export async function resolveOrCreateATAs(
   connection: Connection,
   ownerAddress: PublicKey,
   requests: ResolvedTokenAddressRequest[],
   getAccountRentExempt: () => Promise<number>,
-  payer = ownerAddress
+  payer = ownerAddress,
+  modeIdempotent: boolean = false,
 ): Promise<ResolvedTokenAddressInstruction[]> {
   const nonNativeMints = requests.filter(({ tokenMint }) => !tokenMint.equals(NATIVE_MINT));
   const nativeMints = requests.filter(({ tokenMint }) => tokenMint.equals(NATIVE_MINT));
@@ -89,13 +99,14 @@ export async function resolveOrCreateATAs(
       if (tokenAccount) {
         resolvedInstruction = { address: ataAddress, ...EMPTY_INSTRUCTION };
       } else {
-        const createAtaInstruction = Token.createAssociatedTokenAccountInstruction(
+        const createAtaInstruction = createAssociatedTokenAccountInstruction(
           ASSOCIATED_TOKEN_PROGRAM_ID,
           TOKEN_PROGRAM_ID,
           nonNativeMints[index]!.tokenMint,
           ataAddress,
           ownerAddress,
-          payer
+          payer,
+          modeIdempotent,
         );
 
         resolvedInstruction = {
@@ -131,3 +142,43 @@ export async function deriveATA(ownerAddress: PublicKey, tokenMint: PublicKey):
     ownerAddress
   );
 }
+
+function createAssociatedTokenAccountInstruction(
+  associatedTokenProgramId: PublicKey,
+  tokenProgramId: PublicKey,
+  mint: PublicKey,
+  associatedAccount: PublicKey,
+  owner: PublicKey,
+  payer: PublicKey,
+  modeIdempotent: boolean,
+): TransactionInstruction {
+  if (!modeIdempotent) {
+    return Token.createAssociatedTokenAccountInstruction(
+      associatedTokenProgramId,
+      tokenProgramId,
+      mint,
+      associatedAccount,
+      owner,
+      payer,
+    );
+  }
+
+  // create CreateIdempotent instruction
+  // spl-token v0.1.8 doesn't have a method for CreateIdempotent.
+  // https://github.com/solana-labs/solana-program-library/blob/master/associated-token-account/program/src/instruction.rs#L26
+  const keys = [
+    { pubkey: payer, isSigner: true, isWritable: true },
+    { pubkey: associatedAccount, isSigner: false, isWritable: true },
+    { pubkey: owner, isSigner: false, isWritable: false },
+    { pubkey: mint, isSigner: false, isWritable: false },
+    { pubkey: SystemProgram.programId, isSigner: false, isWritable: false },
+    { pubkey: tokenProgramId, isSigner: false, isWritable: false },
+  ];
+  const instructionData = Buffer.from([1]);
+
+  return new TransactionInstruction({
+    keys,
+    programId: associatedTokenProgramId,
+    data: instructionData,
+  });
+}

packages/common-sdk/tests/ata-util.test.ts

@@ -0,0 +1,245 @@
+import { PublicKey, Connection, Keypair, LAMPORTS_PER_SOL, SystemProgram } from "@solana/web3.js";
+import {
+  Token,
+  TOKEN_PROGRAM_ID,
+  AccountLayout,
+  u64,
+  ASSOCIATED_TOKEN_PROGRAM_ID,
+  NATIVE_MINT,
+} from "@solana/spl-token";
+import { Wallet } from "@project-serum/anchor";
+import { resolveOrCreateATA, resolveOrCreateATAs } from "../src/web3/ata-util";
+import { TransactionBuilder } from "../src/web3/transactions";
+
+jest.setTimeout(100 * 1000 /* ms */);
+
+describe("ata-util", () => {
+  const DEFAULT_RPC_ENDPOINT_URL = "http://localhost:8899";
+
+  const connection = new Connection(DEFAULT_RPC_ENDPOINT_URL, "confirmed");
+  const wallet = new Wallet(Keypair.generate());
+
+  const tokenGetATA = (owner: PublicKey, mint: PublicKey) =>
+    Token.getAssociatedTokenAddress(ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_PROGRAM_ID, mint, owner);
+
+  const createNewMint = () =>
+    Token.createMint(connection, wallet.payer, wallet.publicKey, null, 6, TOKEN_PROGRAM_ID);
+
+  beforeAll(async () => {
+    // airdrop to the test wallet
+    const signature = await connection.requestAirdrop(wallet.publicKey, 1000 * LAMPORTS_PER_SOL);
+    await connection.confirmTransaction(signature);
+  });
+
+  it("resolveOrCreateATA, wrapped sol", async () => {
+    // verify address & instruction
+    const notExpected = await tokenGetATA(wallet.publicKey, NATIVE_MINT);
+    const resolved = await resolveOrCreateATA(
+      connection,
+      wallet.publicKey,
+      NATIVE_MINT,
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      new u64(LAMPORTS_PER_SOL),
+      wallet.publicKey,
+      false
+    );
+    expect(resolved.address.equals(notExpected)).toBeFalsy(); // non-ATA address
+    expect(resolved.instructions.length).toEqual(2);
+    expect(resolved.instructions[0].programId.equals(SystemProgram.programId)).toBeTruthy();
+    expect(resolved.instructions[1].programId.equals(TOKEN_PROGRAM_ID)).toBeTruthy();
+    expect(resolved.cleanupInstructions.length).toEqual(1);
+    expect(resolved.cleanupInstructions[0].programId.equals(TOKEN_PROGRAM_ID)).toBeTruthy();
+  });
+
+  it("resolveOrCreateATA, not exist, modeIdempotent = false", async () => {
+    const mint = await createNewMint();
+
+    // verify address & instruction
+    const expected = await tokenGetATA(wallet.publicKey, mint.publicKey);
+    const resolved = await resolveOrCreateATA(
+      connection,
+      wallet.publicKey,
+      mint.publicKey,
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      new u64(0),
+      wallet.publicKey,
+      false
+    );
+    expect(resolved.address.equals(expected)).toBeTruthy();
+    expect(resolved.instructions.length).toEqual(1);
+    expect(resolved.instructions[0].data.length).toEqual(0); // no instruction data
+
+    // verify transaction
+    const preAccountData = await connection.getAccountInfo(resolved.address);
+    expect(preAccountData).toBeNull();
+
+    const builder = new TransactionBuilder(connection, wallet);
+    builder.addInstruction(resolved);
+    await builder.buildAndExecute();
+
+    const postAccountData = await connection.getAccountInfo(resolved.address);
+    expect(postAccountData?.owner.equals(TOKEN_PROGRAM_ID)).toBeTruthy();
+  });
+
+  it("resolveOrCreateATA, exist, modeIdempotent = false", async () => {
+    const mint = await createNewMint();
+
+    const expected = await mint.createAssociatedTokenAccount(wallet.publicKey);
+    const preAccountData = await connection.getAccountInfo(expected);
+    expect(preAccountData).not.toBeNull();
+
+    // verify address & instruction
+    const resolved = await resolveOrCreateATA(
+      connection,
+      wallet.publicKey,
+      mint.publicKey,
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      new u64(0),
+      wallet.publicKey,
+      false
+    );
+    expect(resolved.address.equals(expected)).toBeTruthy();
+    expect(resolved.instructions.length).toEqual(0);
+  });
+
+  it("resolveOrCreateATA, created before execution, modeIdempotent = false", async () => {
+    const mint = await createNewMint();
+
+    const expected = await tokenGetATA(wallet.publicKey, mint.publicKey);
+    const resolved = await resolveOrCreateATA(
+      connection,
+      wallet.publicKey,
+      mint.publicKey,
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      new u64(0),
+      wallet.publicKey,
+      false
+    );
+    expect(resolved.address.equals(expected)).toBeTruthy();
+    expect(resolved.instructions.length).toEqual(1);
+    expect(resolved.instructions[0].data.length).toEqual(0); // no instruction data
+
+    // created before execution
+    await mint.createAssociatedTokenAccount(wallet.publicKey);
+    const accountData = await connection.getAccountInfo(expected);
+    expect(accountData).not.toBeNull();
+
+    // Tx should be fail
+    const builder = new TransactionBuilder(connection, wallet);
+    builder.addInstruction(resolved);
+    await expect(builder.buildAndExecute()).rejects.toThrow();
+  });
+
+  it("resolveOrCreateATA, created before execution, modeIdempotent = true", async () => {
+    const mint = await createNewMint();
+
+    const expected = await tokenGetATA(wallet.publicKey, mint.publicKey);
+    const resolved = await resolveOrCreateATA(
+      connection,
+      wallet.publicKey,
+      mint.publicKey,
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      new u64(0),
+      wallet.publicKey,
+      true
+    );
+    expect(resolved.address.equals(expected)).toBeTruthy();
+    expect(resolved.instructions.length).toEqual(1);
+    expect(resolved.instructions[0].data[0]).toEqual(1); // 1 byte data
+
+    // created before execution
+    await mint.createAssociatedTokenAccount(wallet.publicKey);
+    const accountData = await connection.getAccountInfo(expected);
+    expect(accountData).not.toBeNull();
+
+    // Tx should be success even if ATA has been created
+    const builder = new TransactionBuilder(connection, wallet);
+    builder.addInstruction(resolved);
+    await expect(builder.buildAndExecute()).resolves.toBeTruthy();
+  });
+
+  it("resolveOrCreateATAs, created before execution, modeIdempotent = false", async () => {
+    const mints = await Promise.all([createNewMint(), createNewMint(), createNewMint()]);
+
+    // create first ATA
+    await mints[0].createAssociatedTokenAccount(wallet.publicKey);
+
+    const expected = await Promise.all(
+      mints.map((mint) => tokenGetATA(wallet.publicKey, mint.publicKey))
+    );
+    const resolved = await resolveOrCreateATAs(
+      connection,
+      wallet.publicKey,
+      mints.map((mint) => ({ tokenMint: mint.publicKey, wrappedSolAmountIn: new u64(0) })),
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      wallet.publicKey,
+      false
+    );
+    expect(resolved[0].address.equals(expected[0])).toBeTruthy();
+    expect(resolved[1].address.equals(expected[1])).toBeTruthy();
+    expect(resolved[2].address.equals(expected[2])).toBeTruthy();
+    expect(resolved[0].instructions.length).toEqual(0); // already exists
+    expect(resolved[1].instructions.length).toEqual(1);
+    expect(resolved[2].instructions.length).toEqual(1);
+    expect(resolved[1].instructions[0].data.length).toEqual(0); // no instruction data
+    expect(resolved[2].instructions[0].data.length).toEqual(0); // no instruction data
+
+    // create second ATA before execution
+    await mints[1].createAssociatedTokenAccount(wallet.publicKey);
+
+    const preAccountData = await connection.getMultipleAccountsInfo(expected);
+    expect(preAccountData[0]).not.toBeNull();
+    expect(preAccountData[1]).not.toBeNull();
+    expect(preAccountData[2]).toBeNull();
+
+    // Tx should be fail
+    const builder = new TransactionBuilder(connection, wallet);
+    builder.addInstructions(resolved);
+    await expect(builder.buildAndExecute()).rejects.toThrow();
+  });
+
+  it("resolveOrCreateATAs, created before execution, modeIdempotent = true", async () => {
+    const mints = await Promise.all([createNewMint(), createNewMint(), createNewMint()]);
+
+    // create first ATA
+    await mints[0].createAssociatedTokenAccount(wallet.publicKey);
+
+    const expected = await Promise.all(
+      mints.map((mint) => tokenGetATA(wallet.publicKey, mint.publicKey))
+    );
+    const resolved = await resolveOrCreateATAs(
+      connection,
+      wallet.publicKey,
+      mints.map((mint) => ({ tokenMint: mint.publicKey, wrappedSolAmountIn: new u64(0) })),
+      () => connection.getMinimumBalanceForRentExemption(AccountLayout.span),
+      wallet.publicKey,
+      true
+    );
+    expect(resolved[0].address.equals(expected[0])).toBeTruthy();
+    expect(resolved[1].address.equals(expected[1])).toBeTruthy();
+    expect(resolved[2].address.equals(expected[2])).toBeTruthy();
+    expect(resolved[0].instructions.length).toEqual(0); // already exists
+    expect(resolved[1].instructions.length).toEqual(1);
+    expect(resolved[2].instructions.length).toEqual(1);
+    expect(resolved[1].instructions[0].data[0]).toEqual(1); // 1 byte data
+    expect(resolved[2].instructions[0].data[0]).toEqual(1); // 1 byte data
+
+    // create second ATA before execution
+    await mints[1].createAssociatedTokenAccount(wallet.publicKey);
+
+    const preAccountData = await connection.getMultipleAccountsInfo(expected);
+    expect(preAccountData[0]).not.toBeNull();
+    expect(preAccountData[1]).not.toBeNull();
+    expect(preAccountData[2]).toBeNull();
+
+    // Tx should be success even if second ATA has been created
+    const builder = new TransactionBuilder(connection, wallet);
+    builder.addInstructions(resolved);
+    await expect(builder.buildAndExecute()).resolves.toBeTruthy();
+
+    const postAccountData = await connection.getMultipleAccountsInfo(expected);
+    expect(postAccountData[0]).not.toBeNull();
+    expect(postAccountData[1]).not.toBeNull();
+    expect(postAccountData[2]).not.toBeNull();
+  });
+});