Skip to content

Commit 15ec5a5

Browse files
hrishikeshjainhrishikeshjain
authored
Merge pull request #6028 from BitGo/COIN-3141
chore(sdk-coin-sol): Add support for pre-calculated recipient token addresses in TokenTransferBuilder
2 parents d3408e5 + 6d62150 commit 15ec5a5

File tree

7 files changed

+71
-14
lines changed

7 files changed

+71
-14
lines changed

‎modules/sdk-coin-sol/src/lib/ataInitializationBuilder.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ export class AtaInitializationBuilder extends TransactionBuilder {
130130
if (this._tokenAssociateRecipients.length === 0) {
131131
assert(this._mint && this._tokenName, 'Mint must be set before building the transaction');
132132
this._owner = this._owner || this._sender;
133+
133134
this._tokenAssociateRecipients.push({
134135
ownerAddress: this._owner,
135136
tokenName: this._tokenName,
@@ -143,7 +144,13 @@ export class AtaInitializationBuilder extends TransactionBuilder {
143144
if (!token) {
144145
throw new BuildTransactionError('Invalid transaction: invalid token name, got: ' + recipient.tokenName);
145146
}
146-
const ataPk = await getAssociatedTokenAccountAddress(token.tokenAddress, recipient.ownerAddress);
147+
148+
// Use the provided ataAddress if it exists, otherwise calculate it
149+
let ataPk = recipient.ataAddress;
150+
if (!ataPk) {
151+
ataPk = await getAssociatedTokenAccountAddress(token.tokenAddress, recipient.ownerAddress);
152+
}
153+
147154
this._instructionsData.push({
148155
type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
149156
params: {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,4 +175,5 @@ export interface TransactionExplanation extends BaseTransactionExplanation {
175175
export class TokenAssociateRecipient {
176176
ownerAddress: string;
177177
tokenName: string;
178+
ataAddress?: string;
178179
}

‎modules/sdk-coin-sol/src/lib/tokenTransferBuilder.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ export class TokenTransferBuilder extends TransactionBuilder {
5454
this._createAtaParams.push({
5555
ownerAddress: ataInitInstruction.params.ownerAddress,
5656
tokenName: ataInitInstruction.params.tokenName,
57+
ataAddress: ataInitInstruction.params.ataAddress,
5758
});
5859
}
5960
}
@@ -127,14 +128,18 @@ export class TokenTransferBuilder extends TransactionBuilder {
127128
uniqueCreateAtaParams.map(async (recipient: TokenAssociateRecipient): Promise<AtaInit> => {
128129
const coin = getSolTokenFromTokenName(recipient.tokenName);
129130
assert(coin instanceof SolCoin);
130-
const recipientTokenAddress = await getAssociatedTokenAccountAddress(coin.tokenAddress, recipient.ownerAddress);
131+
// Use the provided ataAddress if it exists, otherwise calculate it
132+
let ataAddress = recipient.ataAddress;
133+
if (!ataAddress) {
134+
ataAddress = await getAssociatedTokenAccountAddress(coin.tokenAddress, recipient.ownerAddress);
135+
}
131136
return {
132137
type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
133138
params: {
134139
ownerAddress: recipient.ownerAddress,
135140
tokenName: coin.name,
136141
mintAddress: coin.tokenAddress,
137-
ataAddress: recipientTokenAddress,
142+
ataAddress,
138143
payerAddress: this._sender,
139144
},
140145
};

‎modules/sdk-coin-sol/src/lib/transactionBuilder.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import { KeyPair } from '.';
2525
import { InstructionBuilderTypes } from './constants';
2626
import { solInstructionFactory } from './solInstructionFactory';
2727
import assert from 'assert';
28-
import { DurableNonceParams, InstructionParams, Memo, Nonce, Transfer } from './iface';
28+
import { DurableNonceParams, InstructionParams, Memo, Nonce, SetPriorityFee, Transfer } from './iface';
2929
import { instructionParamsFactory } from './instructionParamsFactory';
3030

3131
export abstract class TransactionBuilder extends BaseTransactionBuilder {
@@ -94,7 +94,7 @@ export abstract class TransactionBuilder extends BaseTransactionBuilder {
9494

9595
// If prio fee instruction exists, set the priority fee variable
9696
if (instruction.type === InstructionBuilderTypes.SetPriorityFee) {
97-
const priorityFeeInstructionsData = filteredPriorityFeeInstructionsData[0];
97+
const priorityFeeInstructionsData = filteredPriorityFeeInstructionsData[0] as SetPriorityFee;
9898
this.setPriorityFee({ amount: Number(priorityFeeInstructionsData.params.fee) });
9999
}
100100
}

‎modules/sdk-coin-sol/src/lib/transferBuilderV2.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ export class TransferBuilderV2 extends TransactionBuilder {
5959
this._createAtaParams.push({
6060
ownerAddress: ataInitInstruction.params.ownerAddress,
6161
tokenName: ataInitInstruction.params.tokenName,
62+
ataAddress: ataInitInstruction.params.ataAddress,
6263
});
6364
}
6465
}

‎modules/sdk-coin-sol/test/unit/transactionBuilder/ataInitBuilder.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,10 +251,12 @@ describe('Sol Associated Token Account Builder', () => {
251251
{
252252
ownerAddress: sender.pubkey,
253253
tokenName: mint,
254+
ataAddress: sender.ataPubkey,
254255
},
255256
{
256257
ownerAddress: ownerPubkeys.pubkey,
257258
tokenName: 'sol:ray',
259+
ataAddress: 'ACEuzYtR4gBFt6HLQTYisg2T7k8Vh4ss1SpnqmbVQSNy',
258260
},
259261
];
260262
const multiAtaInitBuilder = (recipients) => {
@@ -325,6 +327,7 @@ describe('Sol Associated Token Account Builder', () => {
325327
const duplicateRecipient = {
326328
ownerAddress: sender.pubkey,
327329
tokenName: mint,
330+
ataAddress: sender.ataPubkey,
328331
};
329332
const txBuilder = multiAtaInitBuilder(recipients);
330333
should(() => txBuilder.enableToken(duplicateRecipient)).throwError(
@@ -336,6 +339,7 @@ describe('Sol Associated Token Account Builder', () => {
336339
const errorMintRecipient = {
337340
ownerAddress: ownerPubkeys.pubkey,
338341
tokenName: 'invalidToken',
342+
ataAddress: ownerPubkeys.ataPubkey,
339343
};
340344
const txBuilder = multiAtaInitBuilder(recipients);
341345
should(() => txBuilder.enableToken(errorMintRecipient)).throwError(
@@ -356,6 +360,7 @@ describe('Sol Associated Token Account Builder', () => {
356360
txBuilder.enableToken({
357361
ownerAddress: account.pub,
358362
tokenName: mint,
363+
ataAddress: sender.ataPubkey,
359364
});
360365
txBuilder.sign({ key: account.prv });
361366
await txBuilder.build().should.rejectedWith('Invalid transaction: missing nonce blockhash');
@@ -367,6 +372,7 @@ describe('Sol Associated Token Account Builder', () => {
367372
txBuilder.enableToken({
368373
ownerAddress: account.pub,
369374
tokenName: mint,
375+
ataAddress: sender.ataPubkey,
370376
});
371377
txBuilder.sign({ key: account.prv });
372378
await txBuilder.build().should.rejectedWith('Invalid transaction: missing sender');
@@ -394,6 +400,7 @@ describe('Sol Associated Token Account Builder', () => {
394400
const invalidOwner = {
395401
ownerAddress: 'invalid owner',
396402
tokenName: mint,
403+
ataAddress: sender.ataPubkey,
397404
};
398405
const txBuilder = multiAtaInitBuilder(recipients);
399406
should(() => txBuilder.enableToken(invalidOwner)).throwError(

‎modules/sdk-coin-sol/test/unit/transactionBuilder/tokenTransferBuilder.ts‎

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -355,7 +355,11 @@ describe('Sol Token Transfer Builder', () => {
355355
txBuilder.sender(walletPK);
356356
txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
357357
txBuilder.memo(memo);
358-
txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
358+
txBuilder.createAssociatedTokenAccount({
359+
ownerAddress: otherAccount.pub,
360+
tokenName: nameUSDC,
361+
ataAddress: ataAddress,
362+
});
359363
const prioFeeMicroLamports = '0';
360364
const priorityFee: FeeOptions = {
361365
amount: prioFeeMicroLamports,
@@ -414,9 +418,21 @@ describe('Sol Token Transfer Builder', () => {
414418
txBuilder.send({ address: account1.pub, amount, tokenName: nameUSDC });
415419
txBuilder.send({ address: account2.pub, amount, tokenName: nameUSDC });
416420
txBuilder.memo(memo);
417-
txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
418-
txBuilder.createAssociatedTokenAccount({ ownerAddress: account1.pub, tokenName: nameUSDC });
419-
txBuilder.createAssociatedTokenAccount({ ownerAddress: account2.pub, tokenName: nameUSDC });
421+
txBuilder.createAssociatedTokenAccount({
422+
ownerAddress: otherAccount.pub,
423+
tokenName: nameUSDC,
424+
ataAddress: ataAddress,
425+
});
426+
txBuilder.createAssociatedTokenAccount({
427+
ownerAddress: account1.pub,
428+
tokenName: nameUSDC,
429+
ataAddress: ataAddress1,
430+
});
431+
txBuilder.createAssociatedTokenAccount({
432+
ownerAddress: account2.pub,
433+
tokenName: nameUSDC,
434+
ataAddress: ataAddress2,
435+
});
420436
txBuilder.setPriorityFee(priorityFee);
421437
const tx = await txBuilder.build();
422438
tx.inputs.length.should.equal(3);
@@ -519,9 +535,21 @@ describe('Sol Token Transfer Builder', () => {
519535
txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
520536
txBuilder.send({ address: otherAccount.pub, amount, tokenName: nameUSDC });
521537
txBuilder.memo(memo);
522-
txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
523-
txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
524-
txBuilder.createAssociatedTokenAccount({ ownerAddress: otherAccount.pub, tokenName: nameUSDC });
538+
txBuilder.createAssociatedTokenAccount({
539+
ownerAddress: otherAccount.pub,
540+
tokenName: nameUSDC,
541+
ataAddress: ataAddress,
542+
});
543+
txBuilder.createAssociatedTokenAccount({
544+
ownerAddress: otherAccount.pub,
545+
tokenName: nameUSDC,
546+
ataAddress: ataAddress,
547+
});
548+
txBuilder.createAssociatedTokenAccount({
549+
ownerAddress: otherAccount.pub,
550+
tokenName: nameUSDC,
551+
ataAddress: ataAddress,
552+
});
525553
txBuilder.setPriorityFee(priorityFee);
526554
const tx = await txBuilder.build();
527555
tx.inputs.length.should.equal(3);
@@ -649,15 +677,23 @@ describe('Sol Token Transfer Builder', () => {
649677
it('for invalid ownerAddress', () => {
650678
const txBuilder = tokenTransferBuilder();
651679
should(() =>
652-
txBuilder.createAssociatedTokenAccount({ ownerAddress: invalidPubKey, tokenName: nameUSDC })
680+
txBuilder.createAssociatedTokenAccount({
681+
ownerAddress: invalidPubKey,
682+
tokenName: nameUSDC,
683+
ataAddress: ataAddress,
684+
})
653685
).throwError('Invalid or missing ownerAddress, got: ' + invalidPubKey);
654686
});
655687

656688
it('for invalid tokenName', () => {
657689
const invalidTokenName = 'tsol:random';
658690
const txBuilder = tokenTransferBuilder();
659691
should(() =>
660-
txBuilder.createAssociatedTokenAccount({ ownerAddress: nonceAccount.pub, tokenName: invalidTokenName })
692+
txBuilder.createAssociatedTokenAccount({
693+
ownerAddress: nonceAccount.pub,
694+
tokenName: invalidTokenName,
695+
ataAddress: ataAddress,
696+
})
661697
).throwError('Invalid token name, got: ' + invalidTokenName);
662698
});
663699
});

0 commit comments

Comments
 (0)