Support associated token for JS (Also, make the program testable) (#1364)

* Implement some js helpers for associated tokens

* Create integration test and fix hard-coding in spl-associated-token

* Run lint:fix and pretty:fix

* Run flow as well...

* More robust test fixture setup

* Revert api breaking part

* Fix tests...

* Populate ts/flow type definitions

* Improve test a bit

* More consistent arg order; docs; more tests

* lints and pretty

* type definition updates and test tweaks

* More simplification...

* More cleanup

* Address review comments and small cleanings

* Bump the version

Author: ryoqun

associated-token-account/program/src/lib.rs

@@ -21,13 +21,11 @@ pub(crate) fn get_associated_token_address_and_bump_seed(
     spl_token_mint_address: &Pubkey,
     program_id: &Pubkey,
 ) -> (Pubkey, u8) {
-    Pubkey::find_program_address(
-        &[
-            &wallet_address.to_bytes(),
-            &spl_token::id().to_bytes(),
-            &spl_token_mint_address.to_bytes(),
-        ],
+    get_associated_token_address_and_bump_seed_internal(
+        wallet_address,
+        spl_token_mint_address,
         program_id,
+        &spl_token::id(),
     )
 }
 
@@ -39,6 +37,22 @@ pub fn get_associated_token_address(
     get_associated_token_address_and_bump_seed(&wallet_address, &spl_token_mint_address, &id()).0
 }
 
+fn get_associated_token_address_and_bump_seed_internal(
+    wallet_address: &Pubkey,
+    spl_token_mint_address: &Pubkey,
+    program_id: &Pubkey,
+    token_program_id: &Pubkey,
+) -> (Pubkey, u8) {
+    Pubkey::find_program_address(
+        &[
+            &wallet_address.to_bytes(),
+            &token_program_id.to_bytes(),
+            &spl_token_mint_address.to_bytes(),
+        ],
+        program_id,
+    )
+}
+
 /// Create an associated token account for the given wallet address and token mint
 ///
 /// Accounts expected by this instruction:

associated-token-account/program/src/processor.rs

@@ -27,12 +27,14 @@ pub fn process_instruction(
     let spl_token_mint_info = next_account_info(account_info_iter)?;
     let system_program_info = next_account_info(account_info_iter)?;
     let spl_token_program_info = next_account_info(account_info_iter)?;
+    let spl_token_program_id = spl_token_program_info.key;
     let rent_sysvar_info = next_account_info(account_info_iter)?;
 
-    let (associated_token_address, bump_seed) = get_associated_token_address_and_bump_seed(
+    let (associated_token_address, bump_seed) = get_associated_token_address_and_bump_seed_internal(
         &wallet_account_info.key,
         &spl_token_mint_info.key,
         program_id,
+        &spl_token_program_id,
     );
     if associated_token_address != *associated_token_account_info.key {
         msg!("Error: Associated address does not match seed derivation");
@@ -41,7 +43,7 @@ pub fn process_instruction(
 
     let associated_token_account_signer_seeds: &[&[_]] = &[
         &wallet_account_info.key.to_bytes(),
-        &spl_token::id().to_bytes(),
+        &spl_token_program_id.to_bytes(),
         &spl_token_mint_info.key.to_bytes(),
         &[bump_seed],
     ];
@@ -87,7 +89,7 @@ pub fn process_instruction(
 
     msg!("Assign the associated token account to the SPL Token program");
     invoke_signed(
-        &system_instruction::assign(associated_token_account_info.key, &spl_token::id()),
+        &system_instruction::assign(associated_token_account_info.key, &spl_token_program_id),
         &[
             associated_token_account_info.clone(),
             system_program_info.clone(),
@@ -98,7 +100,7 @@ pub fn process_instruction(
     msg!("Initialize the associated token account");
     invoke(
         &spl_token::instruction::initialize_account(
-            &spl_token::id(),
+            &spl_token_program_id,
             associated_token_account_info.key,
             spl_token_mint_info.key,
             wallet_account_info.key,

ci/js-test-token.sh

@@ -10,5 +10,6 @@ npm install
 npm run lint
 npm run flow
 npm run defs
+npm run test
 npm run start-with-test-validator
 PROGRAM_VERSION=2.0.4 npm run start-with-test-validator

token/js/cli/main.js

@@ -8,8 +8,10 @@ import {
   loadTokenProgram,
   createMint,
   createAccount,
+  createAssociatedAccount,
   transfer,
   transferChecked,
+  transferCheckedAssociated,
   approveRevoke,
   failOnApproveOverspend,
   setAuthority,
@@ -30,6 +32,8 @@ async function main() {
   await createMint();
   console.log('Run test: createAccount');
   await createAccount();
+  console.log('Run test: createAssociatedAccount');
+  await createAssociatedAccount();
   console.log('Run test: mintTo');
   await mintTo();
   console.log('Run test: mintToChecked');
@@ -38,6 +42,8 @@ async function main() {
   await transfer();
   console.log('Run test: transferChecked');
   await transferChecked();
+  console.log('Run test: transferCheckedAssociated');
+  await transferCheckedAssociated();
   console.log('Run test: approveRevoke');
   await approveRevoke();
   console.log('Run test: failOnApproveOverspend');

token/js/cli/store/index.js

@@ -13,16 +13,20 @@ export class Store {
     return path.join(__dirname, 'store');
   }
 
+  static getFilename(uri: string): string {
+    return path.join(Store.getDir(), uri);
+  }
+
   async load(uri: string): Promise<Object> {
-    const filename = path.join(Store.getDir(), uri);
+    const filename = Store.getFilename(uri);
     const data = await fs.readFile(filename, 'utf8');
     const config = JSON.parse(data);
     return config;
   }
 
   async save(uri: string, config: Object): Promise<void> {
     await mkdirp(Store.getDir());
-    const filename = path.join(Store.getDir(), uri);
+    const filename = Store.getFilename(uri);
     await fs.writeFile(filename, JSON.stringify(config), 'utf8');
   }
 }

token/js/cli/token-test.js

@@ -9,18 +9,25 @@ import {
   BPF_LOADER_PROGRAM_ID,
 } from '@solana/web3.js';
 
-import {Token, NATIVE_MINT} from '../client/token';
+import {
+  Token,
+  TOKEN_PROGRAM_ID,
+  ASSOCIATED_TOKEN_PROGRAM_ID,
+  NATIVE_MINT,
+} from '../client/token';
 import {url} from '../url';
 import {newAccountWithLamports} from '../client/util/new-account-with-lamports';
 import {sleep} from '../client/util/sleep';
 import {Store} from './store';
 
 // Loaded token program's program id
 let programId: PublicKey;
+let associatedProgramId: PublicKey;
 
 // Accounts setup in createMint and used by all subsequent tests
 let testMintAuthority: Account;
 let testToken: Token;
+let testTokenDecimals: number = 2;
 
 // Accounts setup in createAccount and used by all subsequent tests
 let testAccountOwner: Account;
@@ -78,43 +85,60 @@ async function loadProgram(
   return program_account.publicKey;
 }
 
-async function GetPrograms(connection: Connection): Promise<PublicKey> {
+async function GetPrograms(connection: Connection): Promise<void> {
   const programVersion = process.env.PROGRAM_VERSION;
   if (programVersion) {
     switch (programVersion) {
       case '2.0.4':
-        return new PublicKey('TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA');
+        programId = TOKEN_PROGRAM_ID;
+        associatedProgramId = ASSOCIATED_TOKEN_PROGRAM_ID;
+        return;
       default:
         throw new Error('Unknown program version');
     }
   }
 
   const store = new Store();
-  let tokenProgramId = null;
   try {
     const config = await store.load('config.json');
-    console.log('Using pre-loaded Token program');
+    console.log('Using pre-loaded Token programs');
     console.log(
-      '  Note: To reload program remove client/util/store/config.json',
+      `  Note: To reload program remove ${Store.getFilename('config.json')}`,
     );
-    tokenProgramId = new PublicKey(config.tokenProgramId);
+    programId = new PublicKey(config.tokenProgramId);
+    associatedProgramId = new PublicKey(config.associatedTokenProgramId);
+    let info;
+    info = await connection.getAccountInfo(programId);
+    assert(info != null);
+    info = await connection.getAccountInfo(associatedProgramId);
+    assert(info != null);
   } catch (err) {
-    tokenProgramId = await loadProgram(
+    console.log(
+      'Checking pre-loaded Token programs failed, will load new programs:',
+    );
+    console.log({err});
+
+    programId = await loadProgram(
       connection,
       '../../target/bpfel-unknown-unknown/release/spl_token.so',
     );
+    associatedProgramId = await loadProgram(
+      connection,
+      '../../target/bpfel-unknown-unknown/release/spl_associated_token_account.so',
+    );
     await store.save('config.json', {
-      tokenProgramId: tokenProgramId.toString(),
+      tokenProgramId: programId.toString(),
+      associatedTokenProgramId: associatedProgramId.toString(),
     });
   }
-  return tokenProgramId;
 }
 
 export async function loadTokenProgram(): Promise<void> {
   const connection = await getConnection();
-  programId = await GetPrograms(connection);
+  await GetPrograms(connection);
 
   console.log('Token Program ID', programId.toString());
+  console.log('Associated Token Program ID', associatedProgramId.toString());
 }
 
 export async function createMint(): Promise<void> {
@@ -126,9 +150,12 @@ export async function createMint(): Promise<void> {
     payer,
     testMintAuthority.publicKey,
     testMintAuthority.publicKey,
-    2,
+    testTokenDecimals,
     programId,
   );
+  // HACK: override hard-coded ASSOCIATED_TOKEN_PROGRAM_ID with corresponding
+  // custom test fixture
+  testToken.associatedProgramId = associatedProgramId;
 
   const mintInfo = await testToken.getMintInfo();
   if (mintInfo.mintAuthority !== null) {
@@ -137,7 +164,7 @@ export async function createMint(): Promise<void> {
     assert(mintInfo.mintAuthority !== null);
   }
   assert(mintInfo.supply.toNumber() === 0);
-  assert(mintInfo.decimals === 2);
+  assert(mintInfo.decimals === testTokenDecimals);
   assert(mintInfo.isInitialized === true);
   if (mintInfo.freezeAuthority !== null) {
     assert(mintInfo.freezeAuthority.equals(testMintAuthority.publicKey));
@@ -160,6 +187,48 @@ export async function createAccount(): Promise<void> {
   assert(accountInfo.isNative === false);
   assert(accountInfo.rentExemptReserve === null);
   assert(accountInfo.closeAuthority === null);
+
+  // you can create as many accounts as with same owner
+  const testAccount2 = await testToken.createAccount(
+    testAccountOwner.publicKey,
+  );
+  assert(!testAccount2.equals(testAccount));
+}
+
+export async function createAssociatedAccount(): Promise<void> {
+  let info;
+  const connection = await getConnection();
+
+  const owner = new Account();
+  const associatedAddress = await Token.getAssociatedTokenAddress(
+    associatedProgramId,
+    programId,
+    testToken.publicKey,
+    owner.publicKey,
+  );
+
+  // associated account shouldn't exist
+  info = await connection.getAccountInfo(associatedAddress);
+  assert(info == null);
+
+  const createdAddress = await testToken.createAssociatedTokenAccount(
+    owner.publicKey,
+  );
+  assert(createdAddress.equals(associatedAddress));
+
+  // associated account should exist now
+  info = await testToken.getAccountInfo(associatedAddress);
+  assert(info != null);
+  assert(info.mint.equals(testToken.publicKey));
+  assert(info.owner.equals(owner.publicKey));
+  assert(info.amount.toNumber() === 0);
+
+  // creating again should cause TX error for the associated token account
+  assert(
+    await didThrow(testToken, testToken.createAssociatedTokenAccount, [
+      owner.publicKey,
+    ]),
+  );
 }
 
 export async function mintTo(): Promise<void> {
@@ -219,7 +288,7 @@ export async function transferChecked(): Promise<void> {
       testAccountOwner,
       [],
       100,
-      1,
+      testTokenDecimals - 1,
     ]),
   );
 
@@ -229,7 +298,7 @@ export async function transferChecked(): Promise<void> {
     testAccountOwner,
     [],
     100,
-    2,
+    testTokenDecimals,
   );
 
   const mintInfo = await testToken.getMintInfo();
@@ -242,6 +311,26 @@ export async function transferChecked(): Promise<void> {
   assert(testAccountInfo.amount.toNumber() === 1800);
 }
 
+export async function transferCheckedAssociated(): Promise<void> {
+  const dest = new Account().publicKey;
+  let associatedAccount;
+
+  associatedAccount = await testToken.getOrCreateAssociatedAccountInfo(dest);
+  assert(associatedAccount.amount.toNumber() === 0);
+
+  await testToken.transferChecked(
+    testAccount,
+    associatedAccount.address,
+    testAccountOwner,
+    [],
+    123,
+    testTokenDecimals,
+  );
+
+  associatedAccount = await testToken.getOrCreateAssociatedAccountInfo(dest);
+  assert(associatedAccount.amount.toNumber() === 123);
+}
+
 export async function approveRevoke(): Promise<void> {
   const delegate = new Account().publicKey;