Merge pull request #9667 from BitGo/CHALO-1417-close-ata-flow

fix(sdk-coin-sol): support Token-2022 close ATA

Author: MohammedRyaan786

modules/sdk-coin-sol/src/lib/closeAtaBuilder.ts

@@ -14,7 +14,12 @@ const MIX_API_ERROR_MESSAGE =
 
 export class CloseAtaBuilder extends TransactionBuilder {
   // Unified storage for all close entries (single or bulk)
-  protected _closeAtaEntries: { accountAddress: string; destinationAddress: string; authorityAddress: string }[] = [];
+  protected _closeAtaEntries: {
+    accountAddress: string;
+    destinationAddress: string;
+    authorityAddress: string;
+    programId?: string;
+  }[] = [];
 
   // Which API has been used on this builder instance. Locks in on first call so we can
   // reject attempts to mix the legacy single-ATA setters with the bulk addCloseAtaInstruction().
@@ -71,6 +76,16 @@ export class CloseAtaBuilder extends TransactionBuilder {
     return this;
   }
 
+  /** Sets the SPL token program used by the close instruction. */
+  programId(programId: string): this {
+    this._assertSingleAtaApiUsable();
+    validateAddress(programId, 'programId');
+    this._apiMode = 'single';
+    this._ensureSingleEntry();
+    this._closeAtaEntries[0].programId = programId;
+    return this;
+  }
+
   /**
    * Throws if the bulk-ATA API has already been used on this builder.
    */
@@ -93,18 +108,27 @@ export class CloseAtaBuilder extends TransactionBuilder {
    * Add an ATA to close in this transaction (for bulk closure).
    * Cannot be mixed with the single-ATA API (accountAddress/destinationAddress/authorityAddress).
    *
-   * @param {string} accountAddress - the ATA address to close
-   * @param {string} destinationAddress - where rent SOL goes (root wallet address)
-   * @param {string} authorityAddress - ATA owner who must sign
+   * @param accountAddress - the ATA address to close
+   * @param destinationAddress - where rent SOL goes (root wallet address)
+   * @param authorityAddress - ATA owner who must sign
+   * @param programId - SPL token program owning the ATA; omitted for legacy SPL
    */
-  addCloseAtaInstruction(accountAddress: string, destinationAddress: string, authorityAddress: string): this {
+  addCloseAtaInstruction(
+    accountAddress: string,
+    destinationAddress: string,
+    authorityAddress: string,
+    programId?: string
+  ): this {
     if (this._apiMode === 'single') {
       throw new BuildTransactionError(MIX_API_ERROR_MESSAGE);
     }
 
     validateAddress(accountAddress, 'accountAddress');
     validateAddress(destinationAddress, 'destinationAddress');
     validateAddress(authorityAddress, 'authorityAddress');
+    if (programId) {
+      validateAddress(programId, 'programId');
+    }
 
     if (accountAddress === destinationAddress) {
       throw new BuildTransactionError('Account address to close cannot be the same as the destination address');
@@ -115,7 +139,7 @@ export class CloseAtaBuilder extends TransactionBuilder {
     }
 
     this._apiMode = 'bulk';
-    this._closeAtaEntries.push({ accountAddress, destinationAddress, authorityAddress });
+    this._closeAtaEntries.push({ accountAddress, destinationAddress, authorityAddress, programId });
     return this;
   }
 
@@ -129,6 +153,7 @@ export class CloseAtaBuilder extends TransactionBuilder {
           accountAddress: ataCloseInstruction.params.accountAddress,
           destinationAddress: ataCloseInstruction.params.destinationAddress,
           authorityAddress: ataCloseInstruction.params.authorityAddress,
+          programId: ataCloseInstruction.params.programId,
         });
       }
     }
@@ -158,6 +183,7 @@ export class CloseAtaBuilder extends TransactionBuilder {
           accountAddress: entry.accountAddress,
           destinationAddress: entry.destinationAddress,
           authorityAddress: entry.authorityAddress,
+          ...(entry.programId ? { programId: entry.programId } : {}),
         },
       })
     );

modules/sdk-coin-sol/src/lib/iface.ts

@@ -289,7 +289,13 @@ export interface AtaInit {
 
 export interface AtaClose {
   type: InstructionBuilderTypes.CloseAssociatedTokenAccount;
-  params: { accountAddress: string; destinationAddress: string; authorityAddress: string };
+  params: {
+    accountAddress: string;
+    destinationAddress: string;
+    authorityAddress: string;
+    /** SPL token program owning the ATA; omitted for legacy Token Program. */
+    programId?: string;
+  };
 }
 
 export interface AtaRecoverNested {

modules/sdk-coin-sol/src/lib/instructionParamsFactory.ts

@@ -314,6 +314,9 @@ function parseSendInstructions(
             accountAddress,
             destinationAddress,
             authorityAddress,
+            ...(instruction.programId.equals(TOKEN_2022_PROGRAM_ID)
+              ? { programId: instruction.programId.toString() }
+              : {}),
           },
         };
         instructionData.push(ataClose);
@@ -1201,6 +1204,9 @@ function parseAtaCloseInstructions(instructions: TransactionInstruction[]): Arra
             accountAddress: instruction.keys[ataCloseInstructionKeysIndexes.AccountAddress].pubkey.toString(),
             destinationAddress: instruction.keys[ataCloseInstructionKeysIndexes.DestinationAddress].pubkey.toString(),
             authorityAddress: instruction.keys[ataCloseInstructionKeysIndexes.AuthorityAddress].pubkey.toString(),
+            ...(instruction.programId.equals(TOKEN_2022_PROGRAM_ID)
+              ? { programId: instruction.programId.toString() }
+              : {}),
           },
         };
         instructionData.push(ataClose);

modules/sdk-coin-sol/src/lib/solInstructionFactory.ts

@@ -9,6 +9,7 @@ import {
   createTransferCheckedInstruction,
   createTransferCheckedWithFeeInstruction,
   TOKEN_2022_PROGRAM_ID,
+  TOKEN_PROGRAM_ID,
   createApproveInstruction,
 } from '@solana/spl-token';
 import { struct, u8, s8, blob } from '@solana/buffer-layout';
@@ -608,16 +609,19 @@ function createATAInstruction(data: AtaInit): TransactionInstruction[] {
  */
 function closeATAInstruction(data: AtaClose): TransactionInstruction[] {
   const {
-    params: { accountAddress, destinationAddress, authorityAddress },
+    params: { accountAddress, destinationAddress, authorityAddress, programId },
   } = data;
   assert(accountAddress, 'Missing accountAddress param');
   assert(destinationAddress, 'Missing destinationAddress param');
   assert(authorityAddress, 'Missing authorityAddress param');
 
+  const tokenProgramId = programId ? new PublicKey(programId) : TOKEN_PROGRAM_ID;
   const closeAssociatedTokenAccountInstruction = createCloseAccountInstruction(
     new PublicKey(accountAddress),
     new PublicKey(destinationAddress),
-    new PublicKey(authorityAddress)
+    new PublicKey(authorityAddress),
+    [],
+    tokenProgramId
   );
   return [closeAssociatedTokenAccountInstruction];
 }

modules/sdk-coin-sol/test/unit/solInstructionFactory.ts

@@ -39,6 +39,20 @@ describe('Instruction Builder Tests: ', function () {
       ]);
     });
 
+    it('Close ATA uses Token-2022 program when specified', () => {
+      const result = solInstructionFactory({
+        type: InstructionBuilderTypes.CloseAssociatedTokenAccount,
+        params: {
+          accountAddress: testData.authAccount.pub,
+          destinationAddress: testData.authAccount2.pub,
+          authorityAddress: testData.authAccount.pub,
+          programId: TOKEN_2022_PROGRAM_ID.toString(),
+        },
+      });
+
+      result[0].programId.equals(TOKEN_2022_PROGRAM_ID).should.be.true();
+    });
+
     it('Transfer', () => {
       const fromAddress = testData.authAccount.pub;
       const toAddress = testData.nonceAccount.pub;

modules/sdk-coin-sol/test/unit/transactionBuilder/closeAtaBuilder.ts

@@ -220,6 +220,23 @@ describe('Sol Close ATA Builder', () => {
           instruction.params.destinationAddress.should.equal(destinationAddress);
         }
       });
+
+      it('builds mixed legacy SPL and Token-2022 close instructions', async () => {
+        const txBuilder = closeAtaBuilder();
+        txBuilder.addCloseAtaInstruction(ataAddress1, destinationAddress, account.pub);
+        txBuilder.addCloseAtaInstruction(
+          ataAddress2,
+          destinationAddress,
+          account.pub,
+          'TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb'
+        );
+
+        const tx = await txBuilder.build();
+        const instructions = tx.toJson().instructionsData;
+        instructions.length.should.equal(2);
+        should.not.exist(instructions[0].params.programId);
+        instructions[1].params.programId.should.equal('TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb');
+      });
     });
 
     describe('Fail', () => {