From 5b97bab172f91049f6b0583344fa7330cdfcffea Mon Sep 17 00:00:00 2001
From: taobig <studentht@hotmail.com>
Date: Mon, 17 Feb 2025 22:10:39 +0800
Subject: [PATCH] Add FindAssociatedToken2022Address

---
 keys.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/keys.go b/keys.go
index 306c4c98a..9a216823b 100644
--- a/keys.go
+++ b/keys.go
@@ -691,14 +691,21 @@ func FindProgramAddress(seed [][]byte, programID PublicKey) (PublicKey, uint8, e
 	return PublicKey{}, bumpSeed, errors.New("unable to find a valid program address")
 }
 
-func FindAssociatedTokenAddress(
-	wallet PublicKey,
-	mint PublicKey,
-) (PublicKey, uint8, error) {
+func FindAssociatedTokenAddress(wallet PublicKey, mint PublicKey) (PublicKey, uint8, error) {
+	return findAssociatedTokenAddressAndBumpSeed(
+		wallet,
+		mint,
+		SPLAssociatedTokenAccountProgramID,
+		TokenProgramID,
+	)
+}
+
+func FindAssociatedToken2022Address(wallet PublicKey, mint PublicKey) (PublicKey, uint8, error) {
 	return findAssociatedTokenAddressAndBumpSeed(
 		wallet,
 		mint,
 		SPLAssociatedTokenAccountProgramID,
+		Token2022ProgramID,
 	)
 }
 
@@ -706,10 +713,11 @@ func findAssociatedTokenAddressAndBumpSeed(
 	walletAddress PublicKey,
 	splTokenMintAddress PublicKey,
 	programID PublicKey,
+	tokenProgramID PublicKey,
 ) (PublicKey, uint8, error) {
 	return FindProgramAddress([][]byte{
 		walletAddress[:],
-		TokenProgramID[:],
+		tokenProgramID[:],
 		splTokenMintAddress[:],
 	},
 		programID,