BitGo · baltiyal · Sep 2, 2024 · Aug 29, 2024
@@ -1,11 +1,19 @@
 import { TransactionBuilder } from './transactionBuilder';
 import { BuildTransactionError, TransactionType } from '@bitgo/sdk-core';
-import { getAssociatedTokenAccountAddress, getSolTokenFromTokenName, isValidAmount, validateAddress } from './utils';
-import { BaseCoin as CoinConfig, SolCoin } from '@bitgo/statics';
 import { Transaction } from './transaction';
+import {
+  getAssociatedTokenAccountAddress,
+  getSolTokenFromTokenName,
+  isValidAmount,
+  validateAddress,
+  validateMintAddress,
+  validateOwnerAddress,
+} from './utils';
+import { BaseCoin as CoinConfig, SolCoin } from '@bitgo/statics';
 import assert from 'assert';
-import { TokenTransfer, Transfer } from './iface';
+import { AtaInit, TokenAssociateRecipient, TokenTransfer, Transfer } from './iface';
 import { InstructionBuilderTypes } from './constants';
+import _ from 'lodash';
 
 export interface SendParams {
   address: string;
@@ -17,9 +25,10 @@ const UNSIGNED_BIGINT_MAX = BigInt('18446744073709551615');
 
 export class TransferBuilderV2 extends TransactionBuilder {
   private _sendParams: SendParams[] = [];
-
+  private _createAtaParams: TokenAssociateRecipient[];
   constructor(_coinConfig: Readonly<CoinConfig>) {
     super(_coinConfig);
+    this._createAtaParams = [];
   }
 
   protected get transactionType(): TransactionType {
@@ -45,6 +54,12 @@ export class TransferBuilderV2 extends TransactionBuilder {
           amount: transferInstruction.params.amount,
           tokenName: transferInstruction.params.tokenName,
         });
+      } else if (instruction.type === InstructionBuilderTypes.CreateAssociatedTokenAccount) {
+        const ataInitInstruction: AtaInit = instruction;
+        this._createAtaParams.push({
+          ownerAddress: ataInitInstruction.params.ownerAddress,
+          tokenName: ataInitInstruction.params.tokenName,
+        });
       }
     }
   }
@@ -82,11 +97,29 @@ export class TransferBuilderV2 extends TransactionBuilder {
     return this;
   }
 
+  /**
+   *
+   * @param {TokenAssociateRecipient} recipient - recipient of the associated token account creation
+   * @param {string} recipient.ownerAddress - owner of the associated token account
+   * @param {string} recipient.tokenName - name of the token that is intended to associate
+   * @returns {TransactionBuilder} This transaction builder
+   */
+  createAssociatedTokenAccount(recipient: TokenAssociateRecipient): this {
+    validateOwnerAddress(recipient.ownerAddress);
+    const token = getSolTokenFromTokenName(recipient.tokenName);
+    if (!token) {
+      throw new BuildTransactionError('Invalid token name, got: ' + recipient.tokenName);
+    }
+    validateMintAddress(token.tokenAddress);
+
+    this._createAtaParams.push(recipient);
+    return this;
+  }
+
   /** @inheritdoc */
   protected async buildImplementation(): Promise<Transaction> {
     assert(this._sender, 'Sender must be set before building the transaction');
-
-    this._instructionsData = await Promise.all(
+    const sendInstructions = await Promise.all(
       this._sendParams.map(async (sendParams: SendParams): Promise<Transfer | TokenTransfer> => {
         if (sendParams.tokenName) {
           const coin = getSolTokenFromTokenName(sendParams.tokenName);
@@ -98,7 +131,7 @@ export class TransferBuilderV2 extends TransactionBuilder {
               fromAddress: this._sender,
               toAddress: sendParams.address,
               amount: sendParams.amount,
-              tokenName: sendParams.tokenName,
+              tokenName: coin.name,
               sourceAddress: sourceAddress,
             },
           };
@@ -114,7 +147,28 @@ export class TransferBuilderV2 extends TransactionBuilder {
         }
       })
     );
-
+    const uniqueCreateAtaParams = _.uniqBy(this._createAtaParams, (recipient: TokenAssociateRecipient) => {
+      return recipient.ownerAddress + recipient.tokenName;
+    });
+    const createAtaInstructions = await Promise.all(
+      uniqueCreateAtaParams.map(async (recipient: TokenAssociateRecipient): Promise<AtaInit> => {
+        const coin = getSolTokenFromTokenName(recipient.tokenName);
+        assert(coin instanceof SolCoin);
+        const recipientTokenAddress = await getAssociatedTokenAccountAddress(coin.tokenAddress, recipient.ownerAddress);
+        return {
+          type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
+          params: {
+            ownerAddress: recipient.ownerAddress,
+            tokenName: coin.name,
+            mintAddress: coin.tokenAddress,
+            ataAddress: recipientTokenAddress,
+            payerAddress: this._sender,
+          },
+        };
+      })
+    );
+    // order is important, createAtaInstructions must be before sendInstructions
+    this._instructionsData = [...createAtaInstructions, ...sendInstructions];
     return await super.buildImplementation();
   }
 }
@@ -378,6 +378,8 @@ export const TOKEN_TRANSFERV2_UNSIGNED_TX_WITH_MEMO_AND_DURABLE_NONCE =
   'AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgEFCtPhdfVnHvE4DTw6gLpHWddJ/hL8IhXVdS3cKq2ekMAfAGymKVqOJEQemBHH67uu8ISJV4rtwTejLrjw7VSeW6dv+hKJ+pxZaLwHGEyk2Svp5PfAC5ZEi/wYI1tPTHHhbpXS8VwMObd6fTnfCKrnxvwQ5LFhipVbiG+aiTNM1eFsqRgbUvftkOrrDa1Lcb60U4NhZ2ExcX6TCVgQX7ejidUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANEDifvO5SjyCGEdzMc0sxCSVAyyuNWNEA8uqizttNpeBUpTWpkpIQZNJOhxYNo4fHw1td28kruB5B+oQEEFRI0Gp9UXGSxWjuCKhF9z0peIzwNcMUWyGrNE2AYuqUAAAAbd9uHXZaGT2cvhRs7reawctIXtX1s3kTqM9YV+/wCp4zLa+S+r7Oi2P/ekQAXl/f2a+hWHVrYcWpX5BLO40IEDBQMCCAEEBAAAAAkEAwYEAQoM4JMEAAAAAAAJBwAJdGVzdCBtZW1v';
 export const TOKEN_TRANSFERV2_UNSIGNED_TX_WITHOUT_MEMO =
   'AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgECBtPhdfVnHvE4DTw6gLpHWddJ/hL8IhXVdS3cKq2ekMAfAGymKVqOJEQemBHH67uu8ISJV4rtwTejLrjw7VSeW6eV0vFcDDm3en053wiq58b8EOSxYYqVW4hvmokzTNXhbKkYG1L37ZDq6w2tS3G+tFODYWdhMXF+kwlYEF+3o4nV0QOJ+87lKPIIYR3MxzSzEJJUDLK41Y0QDy6qLO202l4G3fbh12Whk9nL4UbO63msHLSF7V9bN5E6jPWFfv8AqeMy2vkvq+zotj/3pEAF5f39mvoVh1a2HFqV+QSzuNCBAQUEAgQDAQoM4JMEAAAAAAAJ';
+export const TOKEN_TRANSFERV2_SIGNED_TX_WITH_WITH_CREATE_ATA_AND_MEMO_AND_DURABLE_NONCE =
+  'AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACEsiaRee0gd6PEhtWkfWgn8n0Zeg8zK4VjCmEvOrnRG+1Kv4V8ip+Obi7MsLfjQugD30xguBAy2XySPRClXgwJAgAHDdPhdfVnHvE4DTw6gLpHWddJ/hL8IhXVdS3cKq2ekMAfAGymKVqOJEQemBHH67uu8ISJV4rtwTejLrjw7VSeW6dJP1r3FCCofxhWCBv6dIPMrDQrI7IXh1M2k+lzW80KTW/6Eon6nFlovAcYTKTZK+nk98ALlkSL/BgjW09MceFuldLxXAw5t3p9Od8IqufG/BDksWGKlVuIb5qJM0zV4WypGBtS9+2Q6usNrUtxvrRTg2FnYTFxfpMJWBBft6OJ1QAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAjJclj04kifG7PRApFI4NgwtaE5na/xCEBI572Nvp+FnRA4n7zuUo8ghhHczHNLMQklQMsrjVjRAPLqos7bTaXgVKU1qZKSEGTSTocWDaOHx8NbXdvJK7geQfqEBBBUSNBqfVFxksVo7gioRfc9KXiM8DXDFFshqzRNgGLqlAAAAGp9UXGSxcUSGMyUw9SvF/WNruCJuh/UTj29mKAAAAAAbd9uHXZaGT2cvhRs7reawctIXtX1s3kTqM9YV+/wCp4zLa+S+r7Oi2P/ekQAXl/f2a+hWHVrYcWpX5BLO40IEEBgMDCgEEBAAAAAcHAQIFCAYMCwAMBAQIBQEKDOCTBAAAAAAACQkACXRlc3QgbWVtbw==';
 export const TOKEN_TRANSFERV2_SIGNED_TX_WITH_MEMO_AND_DURABLE_NONCE =
   'AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA4Q7ZSEAQc2GAfbm2ZCQyA6jji+6BIwpUlPX65ziac2oWuvpT42uQ1ovlfC/2Chb9RyTcoitdkkLYym0t93jMFAgEFCtPhdfVnHvE4DTw6gLpHWddJ/hL8IhXVdS3cKq2ekMAfAGymKVqOJEQemBHH67uu8ISJV4rtwTejLrjw7VSeW6dv+hKJ+pxZaLwHGEyk2Svp5PfAC5ZEi/wYI1tPTHHhbpXS8VwMObd6fTnfCKrnxvwQ5LFhipVbiG+aiTNM1eFsqRgbUvftkOrrDa1Lcb60U4NhZ2ExcX6TCVgQX7ejidUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAANEDifvO5SjyCGEdzMc0sxCSVAyyuNWNEA8uqizttNpeBUpTWpkpIQZNJOhxYNo4fHw1td28kruB5B+oQEEFRI0Gp9UXGSxWjuCKhF9z0peIzwNcMUWyGrNE2AYuqUAAAAbd9uHXZaGT2cvhRs7reawctIXtX1s3kTqM9YV+/wCp4zLa+S+r7Oi2P/ekQAXl/f2a+hWHVrYcWpX5BLO40IEDBQMCCAEEBAAAAAkEAwYEAQoM4JMEAAAAAAAJBwAJdGVzdCBtZW1v';
 export const MULTI_TOKEN_TRANSFERV2_SIGNED =

@@ -4,6 +4,7 @@ import * as testData from '../../resources/sol';
 import should from 'should';
 
 describe('Sol Transfer Builder V2', () => {
+  let ataAddress;
   const factory = getBuilderFactory('tsol');
 
   const authAccount = new KeyPair(testData.authAccount).getKeys();
@@ -15,6 +16,7 @@ describe('Sol Transfer Builder V2', () => {
   const amount = '300000';
   const memo = 'test memo';
   const nameUSDC = testData.tokenTransfers.nameUSDC;
+  const mintUSDC = testData.tokenTransfers.mintUSDC;
   const owner = testData.tokenTransfers.owner;
   const walletPK = testData.associatedTokenAccounts.accounts[0].pub;
   const walletSK = testData.associatedTokenAccounts.accounts[0].prv;
@@ -275,6 +277,9 @@ describe('Sol Transfer Builder V2', () => {
     });
   });
   describe('Succeed Token Transfer', () => {
+    before(async () => {
+      ataAddress = await Utils.getAssociatedTokenAccountAddress(mintUSDC, otherAccount.pub);
+    });
     it('build a token transfer tx unsigned with memo', async () => {
       const txBuilder = factory.getTransferBuilderV2();
       txBuilder.nonce(recentBlockHash);
@@ -408,6 +413,58 @@ describe('Sol Transfer Builder V2', () => {
       should.equal(rawTx, testData.TOKEN_TRANSFERV2_SIGNED_TX_WITH_MEMO_AND_DURABLE_NONCE);
     });
 
+    it('build a token transfer tx unsigned with create ATA, memo and durable nonce', async () => {
+      const txBuilder = factory.getTransferBuilderV2();
+      txBuilder.nonce(recentBlockHash, {
+        walletNonceAddress: nonceAccount.pub,
+        authWalletAddress: walletPK,
+      });
+      txBuilder.feePayer(feePayerAccount.pub);
+      txBuilder.sender(walletPK);
+      txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
+      txBuilder.memo(memo);
+      txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
+      txBuilder.sign({ key: walletSK });
+      const tx = await txBuilder.build();
+      tx.id.should.not.equal(undefined);
+      tx.inputs.length.should.equal(1);
+      tx.inputs[0].should.deepEqual({
+        address: walletPK,
+        value: amount,
+        coin: nameUSDC,
+      });
+      tx.outputs.length.should.equal(1);
+      tx.outputs[0].should.deepEqual({
+        address: otherAccount.pub,
+        value: amount,
+        coin: nameUSDC,
+      });
+      const txJson = tx.toJson();
+      txJson.instructionsData.length.should.equal(3);
+      txJson.instructionsData[0].type.should.equal('CreateAssociatedTokenAccount');
+      txJson.instructionsData[0].params.should.deepEqual({
+        mintAddress: mintUSDC,
+        ataAddress: ataAddress,
+        ownerAddress: otherAccount.pub,
+        payerAddress: walletPK,
+        tokenName: nameUSDC,
+      });
+      txJson.instructionsData[1].type.should.equal('TokenTransfer');
+      txJson.instructionsData[1].params.should.deepEqual({
+        fromAddress: walletPK,
+        toAddress: otherAccount.pub,
+        amount: amount,
+        tokenName: nameUSDC,
+        sourceAddress: 'B5rJjuVi7En63iK6o3ijKdJwAoTe2gwCYmJsVdHQ2aKV',
+      });
+      txJson.instructionsData[2].type.should.equal('Memo');
+      txJson.instructionsData[2].params.memo.should.equal(memo);
+
+      const rawTx = tx.toBroadcastFormat();
+      should.equal(Utils.isValidRawTransaction(rawTx), true);
+      should.equal(rawTx, testData.TOKEN_TRANSFERV2_SIGNED_TX_WITH_WITH_CREATE_ATA_AND_MEMO_AND_DURABLE_NONCE);
+    });
+
     it('build a token multi transfer tx signed with memo and durable nonce', async () => {
       const account1 = new KeyPair({ prv: testData.extraAccounts.prv1 }).getKeys();
       const account2 = new KeyPair({ prv: testData.extraAccounts.prv2 }).getKeys();