Skip to content

Commit 3cb13b3

Browse files
lgalabrulgalabru
committed
feat(associated-token-account): add token program aware ata helpers
1 parent 46f9602 commit 3cb13b3

4 files changed

Lines changed: 157 additions & 17 deletions

File tree

‎keys.go‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -697,22 +697,38 @@ func FindProgramAddress(seed [][]byte, programID PublicKey) (PublicKey, uint8, e
697697
func FindAssociatedTokenAddress(
698698
wallet PublicKey,
699699
mint PublicKey,
700+
) (PublicKey, uint8, error) {
701+
return FindAssociatedTokenAddressWithProgram(
702+
wallet,
703+
mint,
704+
TokenProgramID,
705+
)
706+
}
707+
708+
// FindAssociatedTokenAddressWithProgram returns the associated token account PDA
709+
// for the provided wallet, mint, and token program.
710+
func FindAssociatedTokenAddressWithProgram(
711+
wallet PublicKey,
712+
mint PublicKey,
713+
tokenProgram PublicKey,
700714
) (PublicKey, uint8, error) {
701715
return findAssociatedTokenAddressAndBumpSeed(
702716
wallet,
703717
mint,
718+
tokenProgram,
704719
SPLAssociatedTokenAccountProgramID,
705720
)
706721
}
707722

708723
func findAssociatedTokenAddressAndBumpSeed(
709724
walletAddress PublicKey,
710725
splTokenMintAddress PublicKey,
726+
tokenProgram PublicKey,
711727
programID PublicKey,
712728
) (PublicKey, uint8, error) {
713729
return FindProgramAddress([][]byte{
714730
walletAddress[:],
715-
TokenProgramID[:],
731+
tokenProgram[:],
716732
splTokenMintAddress[:],
717733
},
718734
programID,

‎programs/associated-token-account/Create.go‎

Lines changed: 42 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,10 @@ import (
2525
)
2626

2727
type Create struct {
28-
Payer solana.PublicKey `bin:"-" borsh_skip:"true"`
29-
Wallet solana.PublicKey `bin:"-" borsh_skip:"true"`
30-
Mint solana.PublicKey `bin:"-" borsh_skip:"true"`
28+
Payer solana.PublicKey `bin:"-" borsh_skip:"true"`
29+
Wallet solana.PublicKey `bin:"-" borsh_skip:"true"`
30+
Mint solana.PublicKey `bin:"-" borsh_skip:"true"`
31+
TokenProgram solana.PublicKey `bin:"-" borsh_skip:"true"`
3132

3233
// [0] = [WRITE, SIGNER] Payer
3334
// ··········· Funding account
@@ -51,7 +52,9 @@ type Create struct {
5152

5253
// NewCreateInstructionBuilder creates a new `Create` instruction builder.
5354
func NewCreateInstructionBuilder() *Create {
54-
nd := &Create{}
55+
nd := &Create{
56+
TokenProgram: solana.TokenProgramID,
57+
}
5558
return nd
5659
}
5760

@@ -70,6 +73,11 @@ func (inst *Create) SetMint(mint solana.PublicKey) *Create {
7073
return inst
7174
}
7275

76+
func (inst *Create) SetTokenProgram(tokenProgram solana.PublicKey) *Create {
77+
inst.TokenProgram = tokenProgram
78+
return inst
79+
}
80+
7381
func (inst *Create) SetAccounts(accounts []*solana.AccountMeta) error {
7482
inst.AccountMetaSlice = accounts
7583
if len(accounts) < 6 {
@@ -78,14 +86,20 @@ func (inst *Create) SetAccounts(accounts []*solana.AccountMeta) error {
7886
inst.Payer = accounts[0].PublicKey
7987
inst.Wallet = accounts[2].PublicKey
8088
inst.Mint = accounts[3].PublicKey
89+
inst.TokenProgram = accounts[5].PublicKey
8190
return nil
8291
}
8392

8493
func (inst Create) Build() *Instruction {
8594
// Find the associatedTokenAddress;
86-
associatedTokenAddress, _, _ := solana.FindAssociatedTokenAddress(
95+
tokenProgram := inst.TokenProgram
96+
if tokenProgram.IsZero() {
97+
tokenProgram = solana.TokenProgramID
98+
}
99+
associatedTokenAddress, _, _ := solana.FindAssociatedTokenAddressWithProgram(
87100
inst.Wallet,
88101
inst.Mint,
102+
tokenProgram,
89103
)
90104

91105
keys := []*solana.AccountMeta{
@@ -115,7 +129,7 @@ func (inst Create) Build() *Instruction {
115129
IsWritable: false,
116130
},
117131
{
118-
PublicKey: solana.TokenProgramID,
132+
PublicKey: tokenProgram,
119133
IsSigner: false,
120134
IsWritable: false,
121135
},
@@ -149,9 +163,14 @@ func (inst *Create) Validate() error {
149163
if inst.Mint.IsZero() {
150164
return errors.New("mint not set")
151165
}
152-
_, _, err := solana.FindAssociatedTokenAddress(
166+
tokenProgram := inst.TokenProgram
167+
if tokenProgram.IsZero() {
168+
tokenProgram = solana.TokenProgramID
169+
}
170+
_, _, err := solana.FindAssociatedTokenAddressWithProgram(
153171
inst.Wallet,
154172
inst.Mint,
173+
tokenProgram,
155174
)
156175
if err != nil {
157176
return fmt.Errorf("error while FindAssociatedTokenAddress: %w", err)
@@ -194,11 +213,26 @@ func NewCreateInstruction(
194213
payer solana.PublicKey,
195214
walletAddress solana.PublicKey,
196215
splTokenMintAddress solana.PublicKey,
216+
) *Create {
217+
return NewCreateInstructionWithTokenProgram(
218+
payer,
219+
walletAddress,
220+
splTokenMintAddress,
221+
solana.TokenProgramID,
222+
)
223+
}
224+
225+
func NewCreateInstructionWithTokenProgram(
226+
payer solana.PublicKey,
227+
walletAddress solana.PublicKey,
228+
splTokenMintAddress solana.PublicKey,
229+
tokenProgram solana.PublicKey,
197230
) *Create {
198231
return NewCreateInstructionBuilder().
199232
SetPayer(payer).
200233
SetWallet(walletAddress).
201-
SetMint(splTokenMintAddress)
234+
SetMint(splTokenMintAddress).
235+
SetTokenProgram(tokenProgram)
202236
}
203237

204238
func (inst *Create) GetPayerAccount() *solana.AccountMeta {

‎programs/associated-token-account/CreateIdempotent.go‎

Lines changed: 42 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,10 @@ import (
2828
// address and token mint, if it doesn't already exist.
2929
// Returns an error if the account exists but with a different owner.
3030
type CreateIdempotent struct {
31-
Payer solana.PublicKey `bin:"-" borsh_skip:"true"`
32-
Wallet solana.PublicKey `bin:"-" borsh_skip:"true"`
33-
Mint solana.PublicKey `bin:"-" borsh_skip:"true"`
31+
Payer solana.PublicKey `bin:"-" borsh_skip:"true"`
32+
Wallet solana.PublicKey `bin:"-" borsh_skip:"true"`
33+
Mint solana.PublicKey `bin:"-" borsh_skip:"true"`
34+
TokenProgram solana.PublicKey `bin:"-" borsh_skip:"true"`
3435

3536
// [0] = [WRITE, SIGNER] Payer
3637
// ··········· Funding account
@@ -54,7 +55,9 @@ type CreateIdempotent struct {
5455

5556
// NewCreateIdempotentInstructionBuilder creates a new `CreateIdempotent` instruction builder.
5657
func NewCreateIdempotentInstructionBuilder() *CreateIdempotent {
57-
return &CreateIdempotent{}
58+
return &CreateIdempotent{
59+
TokenProgram: solana.TokenProgramID,
60+
}
5861
}
5962

6063
func (inst *CreateIdempotent) SetPayer(payer solana.PublicKey) *CreateIdempotent {
@@ -72,6 +75,11 @@ func (inst *CreateIdempotent) SetMint(mint solana.PublicKey) *CreateIdempotent {
7275
return inst
7376
}
7477

78+
func (inst *CreateIdempotent) SetTokenProgram(tokenProgram solana.PublicKey) *CreateIdempotent {
79+
inst.TokenProgram = tokenProgram
80+
return inst
81+
}
82+
7583
func (inst *CreateIdempotent) SetAccounts(accounts []*solana.AccountMeta) error {
7684
inst.AccountMetaSlice = accounts
7785
if len(accounts) < 6 {
@@ -80,13 +88,19 @@ func (inst *CreateIdempotent) SetAccounts(accounts []*solana.AccountMeta) error
8088
inst.Payer = accounts[0].PublicKey
8189
inst.Wallet = accounts[2].PublicKey
8290
inst.Mint = accounts[3].PublicKey
91+
inst.TokenProgram = accounts[5].PublicKey
8392
return nil
8493
}
8594

8695
func (inst CreateIdempotent) Build() *Instruction {
87-
associatedTokenAddress, _, _ := solana.FindAssociatedTokenAddress(
96+
tokenProgram := inst.TokenProgram
97+
if tokenProgram.IsZero() {
98+
tokenProgram = solana.TokenProgramID
99+
}
100+
associatedTokenAddress, _, _ := solana.FindAssociatedTokenAddressWithProgram(
88101
inst.Wallet,
89102
inst.Mint,
103+
tokenProgram,
90104
)
91105

92106
keys := []*solana.AccountMeta{
@@ -116,7 +130,7 @@ func (inst CreateIdempotent) Build() *Instruction {
116130
IsWritable: false,
117131
},
118132
{
119-
PublicKey: solana.TokenProgramID,
133+
PublicKey: tokenProgram,
120134
IsSigner: false,
121135
IsWritable: false,
122136
},
@@ -150,9 +164,14 @@ func (inst *CreateIdempotent) Validate() error {
150164
if inst.Mint.IsZero() {
151165
return errors.New("mint not set")
152166
}
153-
_, _, err := solana.FindAssociatedTokenAddress(
167+
tokenProgram := inst.TokenProgram
168+
if tokenProgram.IsZero() {
169+
tokenProgram = solana.TokenProgramID
170+
}
171+
_, _, err := solana.FindAssociatedTokenAddressWithProgram(
154172
inst.Wallet,
155173
inst.Mint,
174+
tokenProgram,
156175
)
157176
if err != nil {
158177
return fmt.Errorf("error while FindAssociatedTokenAddress: %w", err)
@@ -192,11 +211,26 @@ func NewCreateIdempotentInstruction(
192211
payer solana.PublicKey,
193212
walletAddress solana.PublicKey,
194213
splTokenMintAddress solana.PublicKey,
214+
) *CreateIdempotent {
215+
return NewCreateIdempotentInstructionWithTokenProgram(
216+
payer,
217+
walletAddress,
218+
splTokenMintAddress,
219+
solana.TokenProgramID,
220+
)
221+
}
222+
223+
func NewCreateIdempotentInstructionWithTokenProgram(
224+
payer solana.PublicKey,
225+
walletAddress solana.PublicKey,
226+
splTokenMintAddress solana.PublicKey,
227+
tokenProgram solana.PublicKey,
195228
) *CreateIdempotent {
196229
return NewCreateIdempotentInstructionBuilder().
197230
SetPayer(payer).
198231
SetWallet(walletAddress).
199-
SetMint(splTokenMintAddress)
232+
SetMint(splTokenMintAddress).
233+
SetTokenProgram(tokenProgram)
200234
}
201235

202236
func (inst *CreateIdempotent) GetPayerAccount() *solana.AccountMeta {

‎programs/associated-token-account/instruction_test.go‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,62 @@ func TestDecodeSetsAccountsAndGetters(t *testing.T) {
229229
assert.Equal(t, ata, ci.GetAssociatedTokenAddressAccount().PublicKey)
230230
})
231231

232+
t.Run("Create token-2022", func(t *testing.T) {
233+
payer := solana.NewWallet().PublicKey()
234+
wallet := solana.NewWallet().PublicKey()
235+
mint := solana.NewWallet().PublicKey()
236+
237+
ix := NewCreateInstructionBuilder().
238+
SetPayer(payer).
239+
SetWallet(wallet).
240+
SetMint(mint).
241+
SetTokenProgram(solana.Token2022ProgramID).
242+
Build()
243+
244+
accounts := ix.Accounts()
245+
data, err := ix.Data()
246+
require.NoError(t, err)
247+
248+
decoded, err := DecodeInstruction(accounts, data)
249+
require.NoError(t, err)
250+
251+
create, ok := decoded.Impl.(*Create)
252+
require.True(t, ok)
253+
assert.Equal(t, solana.Token2022ProgramID, create.TokenProgram)
254+
255+
ata, _, err := solana.FindAssociatedTokenAddressWithProgram(wallet, mint, solana.Token2022ProgramID)
256+
require.NoError(t, err)
257+
assert.Equal(t, ata, create.GetAssociatedTokenAddressAccount().PublicKey)
258+
})
259+
260+
t.Run("CreateIdempotent token-2022", func(t *testing.T) {
261+
payer := solana.NewWallet().PublicKey()
262+
wallet := solana.NewWallet().PublicKey()
263+
mint := solana.NewWallet().PublicKey()
264+
265+
ix := NewCreateIdempotentInstructionBuilder().
266+
SetPayer(payer).
267+
SetWallet(wallet).
268+
SetMint(mint).
269+
SetTokenProgram(solana.Token2022ProgramID).
270+
Build()
271+
272+
accounts := ix.Accounts()
273+
data, err := ix.Data()
274+
require.NoError(t, err)
275+
276+
decoded, err := DecodeInstruction(accounts, data)
277+
require.NoError(t, err)
278+
279+
ci, ok := decoded.Impl.(*CreateIdempotent)
280+
require.True(t, ok)
281+
assert.Equal(t, solana.Token2022ProgramID, ci.TokenProgram)
282+
283+
ata, _, err := solana.FindAssociatedTokenAddressWithProgram(wallet, mint, solana.Token2022ProgramID)
284+
require.NoError(t, err)
285+
assert.Equal(t, ata, ci.GetAssociatedTokenAddressAccount().PublicKey)
286+
})
287+
232288
t.Run("RecoverNested", func(t *testing.T) {
233289
wallet := solana.NewWallet().PublicKey()
234290
nestedMint := solana.NewWallet().PublicKey()

0 commit comments

Comments
 (0)