Merge pull request #6201 from BitGo/WIN-5670

fix(sdk-coin-sol): fix ata creation for sol 2022 token

Author: abhishekagrawal080

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

@@ -12,6 +12,7 @@ import {
   decodeCloseAccountInstruction,
   getAssociatedTokenAddress,
   TOKEN_PROGRAM_ID,
+  TOKEN_2022_PROGRAM_ID,
 } from '@solana/spl-token';
 import {
   Keypair,
@@ -475,6 +476,29 @@ export function getSolTokenFromAddress(tokenAddress: string, network: BaseNetwor
   return undefined;
 }
 
+/**
+ * Get the statics coin object matching a given Solana token address if it exists
+ *
+ * @param tokenAddress The token address to match against
+ * @param network Solana Mainnet or Testnet
+ * @returns statics BaseCoin object for the matching token
+ */
+export function getSolTokenFromAddressOnly(tokenAddress: string): Readonly<BaseCoin> | undefined {
+  const tokens = coins.filter((coin) => {
+    if (coin instanceof SolCoin) {
+      return coin.tokenAddress.toLowerCase() === tokenAddress.toLowerCase();
+    }
+    return false;
+  });
+  const tokensArray = tokens.map((token) => token);
+  if (tokensArray.length >= 1) {
+    // there should never be two tokens with the same contract address, so we assert that here
+    assert(tokensArray.length === 1);
+    return tokensArray[0];
+  }
+  return undefined;
+}
+
 /**
  * Get the solana token object from token name
  * @param tokenName The token name to match against
@@ -505,17 +529,26 @@ export async function getAssociatedTokenAccountAddress(
   ownerAddress: string,
   allowOwnerOffCurve = false
 ): Promise<string> {
+  const mintPublicKey = new PublicKey(tokenMintAddress);
   const ownerPublicKey = new PublicKey(ownerAddress);
 
   // tokenAddress are not on ed25519 curve, so they can't be used as ownerAddress
   if (!allowOwnerOffCurve && !PublicKey.isOnCurve(ownerPublicKey.toBuffer())) {
     throw new UtilsError('Invalid ownerAddress - address off ed25519 curve, got: ' + ownerAddress);
   }
-  const ataAddress = await getAssociatedTokenAddress(
-    new PublicKey(tokenMintAddress),
-    ownerPublicKey,
-    allowOwnerOffCurve
-  );
+
+  const coin = getSolTokenFromAddressOnly(tokenMintAddress);
+  if (!coin || !(coin instanceof SolCoin)) {
+    throw new UtilsError(`Token not found or not a Solana token: ${tokenMintAddress}`);
+  }
+
+  let ataAddress: PublicKey;
+  const programId = coin.programId;
+  if (programId === TOKEN_2022_PROGRAM_ID.toString()) {
+    ataAddress = await getAssociatedTokenAddress(mintPublicKey, ownerPublicKey, false, TOKEN_2022_PROGRAM_ID);
+  } else {
+    ataAddress = await getAssociatedTokenAddress(mintPublicKey, ownerPublicKey, allowOwnerOffCurve);
+  }
   return ataAddress.toString();
 }
 

modules/sdk-coin-sol/test/resources/sol.ts

@@ -115,6 +115,11 @@ export const authAccount = {
   prv: 'Vd4PFt2YsjZFbyW6ej5bCK46QRptKyBe6Aq2FPxpynBzqXL5RUt3L3Y4FXyKq9igfshKFskv5yGqdrgvRknKbLk',
 };
 
+export const sol2022AuthAccount = {
+  pub: 'GZmjWFJT9dUnpWNqJabnRMq4twyqCr9kXFc3zMGePma9',
+  prv: 'Vd4PFt2YsjZFbyW6ej5bCK46QRptKyBe6Aq2FPxpynBzqXL5RUt3L3Y4FXyKq9igfshKFskv5yGqdrgvRknKbLk',
+};
+
 export const authAccount2 = {
   pub: '7nSM6jwFRGP5BDZQBcgBgpPNfSB2PZKUyLBS6SrJP9kx',
   prv: '5ecNbERFDGwouBRNv8G5P2e4rwQpKK7DeveahDoz168hkiuUNt6MdmzVDLPGtnSX9hugboCzbKLoNsuaWDFEs7UY',
@@ -325,6 +330,7 @@ export const tokenTransfers = {
   owner: 'i34VuFZY1Lz8Tg9ajAqq5MNjE3ES32uVR8oPFFtLd9F',
   nameUSDC: 'tsol:usdc',
   mintUSDC: 'F4uLeXJoFz3hw13MposuwaQbMcZbCjqvEGPPeRRB1Byf',
+  mintAI16Z: 'HeLp6NuQkmYB4pYWo2zYs22mESHXPQYzXbB8n4V98jwC',
   sourceUSDC: '2fyhC1YbqaYszkUQw2YGNRVkr2abr69UwFXVCjz4Q5f5',
   sourceRAY: 'EF3VmejtM8sAznwsiYAwF1LA7N2tRAgFUvoVjF3MMwzG',
   sourceSRM: '2P6cevaY3RLoQdW6xiduwePFqZCCqM6mTzNr3rpJbzKX',

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

@@ -315,6 +315,22 @@ describe('SOL util library', function () {
     });
   });
 
+  describe('getAssociatedTokenAccountAddress for sol 2022 token', function () {
+    const ai16ZMintAddress = testData.tokenTransfers.mintAI16Z;
+    const tokenAddress = 'JC2HAixRNrhBgNjCsV2PBfZEdyitEnWqJPNGokFD8xSA';
+    it('should succeed for native address as owner address', async function () {
+      const ownerAddress = testData.sol2022AuthAccount.pub;
+      const result = await Utils.getAssociatedTokenAccountAddress(ai16ZMintAddress, ownerAddress);
+      result.should.be.equal(tokenAddress);
+    });
+    it('should fail for token address as owner address', async function () {
+      const invalidOwnerAddress = tokenAddress;
+      await Utils.getAssociatedTokenAccountAddress(ai16ZMintAddress, invalidOwnerAddress).should.be.rejectedWith(
+        'Invalid ownerAddress - address off ed25519 curve, got: ' + invalidOwnerAddress
+      );
+    });
+  });
+
   describe('matchTransactionTypeByInstructionsOrder', function () {
     describe('Activate stake instructions', function () {
       it('should match staking activate instructions', function () {

modules/statics/src/account.ts

@@ -8,6 +8,14 @@ import {
   CELO_TOKEN_FEATURES,
 } from './coinFeatures';
 
+/**
+ * This is the program id against sol tokem program.
+ */
+export enum ProgramID {
+  TokenProgramId = 'TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA',
+  Token2022ProgramId = 'TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb',
+}
+
 export interface AccountConstructorOptions {
   id: string;
   fullName: string;
@@ -97,6 +105,7 @@ export interface EosCoinConstructorOptions extends AccountConstructorOptions {
 export interface SolCoinConstructorOptions extends AccountConstructorOptions {
   tokenAddress: string;
   contractAddress: string;
+  programId: string;
 }
 
 export interface AdaCoinConstructorOptions extends AccountConstructorOptions {
@@ -366,13 +375,15 @@ export class EosCoin extends AccountCoinToken {
 export class SolCoin extends AccountCoinToken {
   public tokenAddress: string;
   public contractAddress: string;
+  public programId: string;
   constructor(options: SolCoinConstructorOptions) {
     super({
       ...options,
     });
 
     this.tokenAddress = options.contractAddress;
     this.contractAddress = options.contractAddress;
+    this.programId = options.programId;
   }
 }
 
@@ -1610,6 +1621,7 @@ export function solToken(
   contractAddress: string,
   asset: UnderlyingAsset,
   features: CoinFeature[] = [...AccountCoin.DEFAULT_FEATURES, CoinFeature.REQUIRES_RESERVE],
+  programId: ProgramID = ProgramID.TokenProgramId,
   prefix = '',
   suffix: string = name.toUpperCase(),
   network: AccountNetwork = Networks.main.sol,
@@ -1626,6 +1638,7 @@ export function solToken(
       prefix,
       suffix,
       features,
+      programId,
       decimalPlaces,
       asset,
       isToken: true,
@@ -1658,6 +1671,7 @@ export function tsolToken(
   contractAddress: string,
   asset: UnderlyingAsset,
   features: CoinFeature[] = [...AccountCoin.DEFAULT_FEATURES, CoinFeature.REQUIRES_RESERVE],
+  programId = ProgramID.TokenProgramId,
   prefix = '',
   suffix: string = name.toUpperCase(),
   network: AccountNetwork = Networks.test.sol
@@ -1671,6 +1685,7 @@ export function tsolToken(
     contractAddress,
     asset,
     features,
+    programId,
     prefix,
     suffix,
     network

modules/statics/src/coins/solTokens.ts

@@ -1,4 +1,4 @@
-import { solToken, tsolToken } from '../account';
+import { ProgramID, solToken, tsolToken } from '../account';
 import { CoinFeature, UnderlyingAsset } from '../base';
 import {
   SOL_TOKEN_FEATURES,
@@ -75,7 +75,8 @@ export const solTokens = [
     'HeLp6NuQkmYB4pYWo2zYs22mESHXPQYzXbB8n4V98jwC', // https://solscan.io/token/HeLp6NuQkmYB4pYWo2zYs22mESHXPQYzXbB8n4V98jwC
     'HeLp6NuQkmYB4pYWo2zYs22mESHXPQYzXbB8n4V98jwC',
     UnderlyingAsset['sol:ai16z'],
-    SOL_TOKEN_FEATURES
+    SOL_TOKEN_FEATURES,
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     'd6bee044-de5b-4ce2-8ed2-dbae6f23e728',
@@ -2471,7 +2472,8 @@ export const solTokens = [
       CoinFeature.CUSTODY_BITGO_FRANKFURT,
       CoinFeature.CUSTODY_BITGO_GERMANY,
       CoinFeature.CUSTODY_BITGO_SINGAPORE,
-    ]
+    ],
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     'b4e9b281-952e-4b75-8d57-8666e0b4a33e',
@@ -2486,7 +2488,8 @@ export const solTokens = [
       CoinFeature.CUSTODY_BITGO_FRANKFURT,
       CoinFeature.CUSTODY_BITGO_GERMANY,
       CoinFeature.CUSTODY_BITGO_SINGAPORE,
-    ]
+    ],
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     'cf7d8a35-11ea-46dc-80a5-e0bed443a18c',
@@ -2501,7 +2504,8 @@ export const solTokens = [
       CoinFeature.CUSTODY_BITGO_FRANKFURT,
       CoinFeature.CUSTODY_BITGO_GERMANY,
       CoinFeature.CUSTODY_BITGO_SINGAPORE,
-    ]
+    ],
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     '694e2cc5-cc12-4fa9-9447-3904ac7a2c57',
@@ -2516,7 +2520,8 @@ export const solTokens = [
       CoinFeature.CUSTODY_BITGO_FRANKFURT,
       CoinFeature.CUSTODY_BITGO_GERMANY,
       CoinFeature.CUSTODY_BITGO_SINGAPORE,
-    ]
+    ],
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     '6dd706b0-4764-48b1-b45d-e78a98e609b3',
@@ -2531,7 +2536,8 @@ export const solTokens = [
       CoinFeature.CUSTODY_BITGO_FRANKFURT,
       CoinFeature.CUSTODY_BITGO_GERMANY,
       CoinFeature.CUSTODY_BITGO_SINGAPORE,
-    ]
+    ],
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     '679055e6-e8cc-4b91-906c-0689de806a59',
@@ -2566,7 +2572,8 @@ export const solTokens = [
     '4MmJVdwYN8LwvbGeCowYjSx7KoEi6BJWg8XXnW4fDDp6',
     '4MmJVdwYN8LwvbGeCowYjSx7KoEi6BJWg8XXnW4fDDp6',
     UnderlyingAsset['sol:tbill'],
-    SOL_TOKEN_FEATURES
+    SOL_TOKEN_FEATURES,
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     '88814831-9cae-4174-8d1c-380aeb3fe55a',
@@ -2576,7 +2583,8 @@ export const solTokens = [
     'DghpMkatCiUsofbTmid3M3kAbDTPqDwKiYHnudXeGG52',
     'DghpMkatCiUsofbTmid3M3kAbDTPqDwKiYHnudXeGG52',
     UnderlyingAsset['sol:eurcv'],
-    SOL_TOKEN_FEATURES
+    SOL_TOKEN_FEATURES,
+    ProgramID.Token2022ProgramId
   ),
   solToken(
     '9a89cbd8-f097-4ab8-a5ea-f594c44cb33e',