token-swap: Fix withdraw instruction (#405)

* Include pool mint in instruction, required for burning
* Add simple Rust test
* Fix js withdraw test

token-swap/js/cli/token-swap-test.js

@@ -240,8 +240,12 @@ export async function deposit(): Promise<void> {
 }
 
 export async function withdraw(): Promise<void> {
+  console.log('Creating withdraw token A account');
   let userAccountA = await mintA.createAccount(owner.publicKey);
+  console.log('Creating withdraw token B account');
   let userAccountB = await mintB.createAccount(owner.publicKey);
+
+  console.log('Approving withdrawal from pool account');
   await tokenPool.approve(
     tokenAccountPool,
     authority,
@@ -251,8 +255,10 @@ export async function withdraw(): Promise<void> {
   );
   const [tokenProgramId,] = await GetPrograms(connection);
 
+  console.log('Withdrawing pool tokens for A and B tokens');
   await tokenSwap.withdraw(
     authority,
+    tokenPool.publicKey,
     tokenAccountPool,
     tokenAccountA,
     tokenAccountB,
@@ -262,21 +268,15 @@ export async function withdraw(): Promise<void> {
     USER_AMOUNT
   );
 
-  let info;
+  let info = await tokenPool.getAccountInfo(tokenAccountPool);
-  info = await tokenPool.getAccountInfo(tokenAccountPool);
-  console.log('tokenAccountPool', info.amount.toNumber());
   assert(info.amount.toNumber() == BASE_AMOUNT - USER_AMOUNT);
   info = await mintA.getAccountInfo(tokenAccountA);
-  console.log('tokenAccountA', info.amount.toNumber());
   assert(info.amount.toNumber() == BASE_AMOUNT);
   info = await mintB.getAccountInfo(tokenAccountB);
-  console.log('tokenAccountB', info.amount.toNumber());
   assert(info.amount.toNumber() == BASE_AMOUNT);
   info = await mintA.getAccountInfo(userAccountA);
-  console.log('userAccountA', info.amount.toNumber());
   assert(info.amount.toNumber() == USER_AMOUNT);
   info = await mintB.getAccountInfo(userAccountB);
-  console.log('userAccountB', info.amount.toNumber());
   assert(info.amount.toNumber() == USER_AMOUNT);
 }
 

token-swap/js/client/token-swap.js

@@ -467,6 +467,7 @@ export class TokenSwap {
    */
   async withdraw(
     authority: PublicKey,
+    poolMint: PublicKey,
     sourcePoolAccount: PublicKey,
     fromA: PublicKey,
     fromB: PublicKey,
@@ -481,6 +482,7 @@ export class TokenSwap {
       new Transaction().add(
         this.withdrawInstruction(
           authority,
+          poolMint,
           sourcePoolAccount,
           fromA,
           fromB,
@@ -493,8 +495,10 @@ export class TokenSwap {
       this.payer,
     );
   }
+
   withdrawInstruction(
     authority: PublicKey,
+    poolMint: PublicKey,
     sourcePoolAccount: PublicKey,
     fromA: PublicKey,
     fromB: PublicKey,
@@ -520,6 +524,7 @@ export class TokenSwap {
     const keys = [
       {pubkey: this.tokenSwap, isSigner: false, isWritable: false},
       {pubkey: authority, isSigner: false, isWritable: false},
+      {pubkey: poolMint, isSigner: false, isWritable: true},
       {pubkey: sourcePoolAccount, isSigner: false, isWritable: true},
       {pubkey: fromA, isSigner: false, isWritable: true},
       {pubkey: fromB, isSigner: false, isWritable: true},

token-swap/program/src/instruction.rs

@@ -78,12 +78,13 @@ pub enum SwapInstruction {
     ///
     ///   0. `[]` Token-swap
     ///   1. `[]` $authority
-    ///   2. `[writable]` SOURCE Pool account, amount is transferable by $authority.
+    ///   2. `[writable]` Pool mint account, $authority is the owner
-    ///   5. `[writable]` token_a Account to withdraw FROM.
+    ///   3. `[writable]` SOURCE Pool account, amount is transferable by $authority.
-    ///   6. `[writable]` token_b Account to withdraw FROM.
+    ///   4. `[writable]` token_a Swap Account to withdraw FROM.
-    ///   7. `[writable]` token_a user Account.
+    ///   5. `[writable]` token_b Swap Account to withdraw FROM.
-    ///   8. `[writable]` token_b user Account.
+    ///   6. `[writable]` token_a user Account to credit.
-    ///   9. '[]` Token program id
+    ///   7. `[writable]` token_b user Account to credit.
+    ///   8. '[]` Token program id
     ///   userdata: SOURCE amount of pool tokens to transfer. User receives an output based on the
     ///   percentage of the pool tokens that are returned.
     Withdraw(u64),
@@ -222,6 +223,41 @@ pub fn deposit(
     })
 }
 
+/// Creates a 'withdraw' instruction.
+pub fn withdraw(
+    program_id: &Pubkey,
+    token_program_id: &Pubkey,
+    swap_pubkey: &Pubkey,
+    authority_pubkey: &Pubkey,
+    pool_mint_pubkey: &Pubkey,
+    source_pubkey: &Pubkey,
+    swap_token_a_pubkey: &Pubkey,
+    swap_token_b_pubkey: &Pubkey,
+    output_token_a_pubkey: &Pubkey,
+    output_token_b_pubkey: &Pubkey,
+    amount: u64,
+) -> Result<Instruction, ProgramError> {
+    let data = SwapInstruction::Withdraw(amount).serialize()?;
+
+    let accounts = vec![
+        AccountMeta::new(*swap_pubkey, false),
+        AccountMeta::new(*authority_pubkey, false),
+        AccountMeta::new(*pool_mint_pubkey, false),
+        AccountMeta::new(*source_pubkey, false),
+        AccountMeta::new(*swap_token_a_pubkey, false),
+        AccountMeta::new(*swap_token_b_pubkey, false),
+        AccountMeta::new(*output_token_a_pubkey, false),
+        AccountMeta::new(*output_token_b_pubkey, false),
+        AccountMeta::new(*token_program_id, false),
+    ];
+
+    Ok(Instruction {
+        program_id: *program_id,
+        accounts,
+        data,
+    })
+}
+
 /// Unpacks a reference from a bytes buffer.
 /// TODO actually pack / unpack instead of relying on normal memory layout.
 pub fn unpack<T>(input: &[u8]) -> Result<&T, ProgramError> {

token-swap/program/src/processor.rs

@@ -48,27 +48,35 @@ impl Processor {
         Pubkey::create_program_address(&[&my_info.to_bytes()[..32], &[nonce]], program_id)
             .or(Err(Error::InvalidProgramAddress))
     }
+
     /// Issue a spl_token `Burn` instruction.
-    pub fn token_burn(
+    pub fn token_burn<'a>(
-        accounts: &[AccountInfo],
-        token_program_id: &Pubkey,
         swap: &Pubkey,
-        burn_account: &Pubkey,
+        token_program: AccountInfo<'a>,
-        mint: &Pubkey,
+        burn_account: AccountInfo<'a>,
-        authority: &Pubkey,
+        mint: AccountInfo<'a>,
+        authority: AccountInfo<'a>,
+        nonce: u8,
         amount: u64,
     ) -> Result<(), ProgramError> {
         let swap_bytes = swap.to_bytes();
-        let signers = &[&[&swap_bytes[..32]][..]];
+        let authority_signature_seeds = [&swap_bytes[..32], &[nonce]];
+        let signers = &[&authority_signature_seeds[..]];
+
         let ix = spl_token::instruction::burn(
-            token_program_id,
+            token_program.key,
-            burn_account,
+            burn_account.key,
-            mint,
+            mint.key,
-            authority,
+            authority.key,
             &[],
             amount,
         )?;
-        invoke_signed(&ix, accounts, signers)
+
+        invoke_signed(
+            &ix,
+            &[burn_account, mint, authority, token_program],
+            signers,
+        )
     }
 
     /// Issue a spl_token `MintTo` instruction.
@@ -338,6 +346,7 @@ impl Processor {
         let account_info_iter = &mut accounts.iter();
         let swap_info = next_account_info(account_info_iter)?;
         let authority_info = next_account_info(account_info_iter)?;
+        let pool_mint_info = next_account_info(account_info_iter)?;
         let source_info = next_account_info(account_info_iter)?;
         let token_a_info = next_account_info(account_info_iter)?;
         let token_b_info = next_account_info(account_info_iter)?;
@@ -389,16 +398,17 @@ impl Processor {
             b_amount,
         )?;
         Self::token_burn(
-            accounts,
-            token_program_info.key,
             swap_info.key,
-            source_info.key,
+            token_program_info.clone(),
-            &token_swap.pool_mint,
+            source_info.clone(),
-            authority_info.key,
+            pool_mint_info.clone(),
+            authority_info.clone(),
+            token_swap.nonce,
             amount,
         )?;
         Ok(())
     }
+
     /// Processes an [Instruction](enum.Instruction.html).
     pub fn process(program_id: &Pubkey, accounts: &[AccountInfo], input: &[u8]) -> ProgramResult {
         let instruction = SwapInstruction::deserialize(input)?;
@@ -491,7 +501,7 @@ solana_sdk::program_stubs!();
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::instruction::{deposit, initialize};
+    use crate::instruction::{deposit, initialize, withdraw};
     use solana_sdk::{
         account::Account, account_info::create_is_signer_account_infos, instruction::Instruction,
         rent::Rent, sysvar::rent,
@@ -510,7 +520,7 @@ mod tests {
         swap_account: Account,
         pool_mint_key: Pubkey,
         pool_mint_account: Account,
-        _pool_token_key: Pubkey,
+        pool_token_key: Pubkey,
         pool_token_account: Account,
         token_a_key: Pubkey,
         token_a_account: Account,
@@ -724,7 +734,7 @@ mod tests {
             swap_account,
             pool_mint_key,
             pool_mint_account,
-            _pool_token_key: pool_token_key,
+            pool_token_key,
             pool_token_account,
             token_a_key,
             token_a_account,
@@ -852,4 +862,76 @@ mod tests {
             pool_account.amount + depositor_pool_account.amount
         );
     }
+
+    #[test]
+    fn test_withdraw() {
+        let numerator = 1;
+        let denominator = 2;
+        let token_a_amount = 1000;
+        let token_b_amount = 2000;
+        let mut accounts = initialize_swap(numerator, denominator, token_a_amount, token_b_amount);
+        let seeds = [&accounts.swap_key.to_bytes()[..32], &[accounts.nonce]];
+        let authority_key = Pubkey::create_program_address(&seeds, &SWAP_PROGRAM_ID).unwrap();
+        let initial_a = token_a_amount / 10;
+        let (withdraw_token_a_key, mut withdraw_token_a_account) = mint_token(
+            &TOKEN_PROGRAM_ID,
+            &accounts.token_a_mint_key,
+            &mut accounts.token_a_mint_account,
+            &authority_key,
+            initial_a,
+        );
+        let initial_b = token_b_amount / 10;
+        let (withdraw_token_b_key, mut withdraw_token_b_account) = mint_token(
+            &TOKEN_PROGRAM_ID,
+            &accounts.token_b_mint_key,
+            &mut accounts.token_b_mint_account,
+            &authority_key,
+            initial_b,
+        );
+
+        let withdraw_amount = token_a_amount / 4;
+        do_process_instruction(
+            withdraw(
+                &SWAP_PROGRAM_ID,
+                &TOKEN_PROGRAM_ID,
+                &accounts.swap_key,
+                &authority_key,
+                &accounts.pool_mint_key,
+                &accounts.pool_token_key,
+                &accounts.token_a_key,
+                &accounts.token_b_key,
+                &withdraw_token_a_key,
+                &withdraw_token_b_key,
+                withdraw_amount,
+            )
+            .unwrap(),
+            vec![
+                &mut accounts.swap_account,
+                &mut Account::default(),
+                &mut accounts.pool_mint_account,
+                &mut accounts.pool_token_account,
+                &mut accounts.token_a_account,
+                &mut accounts.token_b_account,
+                &mut withdraw_token_a_account,
+                &mut withdraw_token_b_account,
+                &mut Account::default(),
+            ],
+        )
+        .unwrap();
+
+        let token_a = Processor::token_account_deserialize(&accounts.token_a_account.data).unwrap();
+        assert_eq!(token_a.amount, token_a_amount - withdraw_amount);
+        let token_b = Processor::token_account_deserialize(&accounts.token_b_account.data).unwrap();
+        assert_eq!(token_b.amount, token_b_amount - (withdraw_amount * 2));
+        let withdraw_token_a =
+            Processor::token_account_deserialize(&withdraw_token_a_account.data).unwrap();
+        assert_eq!(withdraw_token_a.amount, initial_a + withdraw_amount);
+        let withdraw_token_b =
+            Processor::token_account_deserialize(&withdraw_token_b_account.data).unwrap();
+        assert_eq!(withdraw_token_b.amount, initial_b + (withdraw_amount * 2));
+        let pool_account =
+            Processor::token_account_deserialize(&accounts.pool_token_account.data).unwrap();
+        let pool_mint = Processor::mint_deserialize(&accounts.pool_mint_account.data).unwrap();
+        assert_eq!(pool_mint.supply, pool_account.amount);
+    }
 }