feat: add admin withdraw ix for vault token a,b

sol-mocha · sol-mocha · commit ac7519a3b51b · 2024-03-07T07:57:44.000+06:00
diff --git a/programs/drip/src/actions/admin.rs b/programs/drip/src/actions/admin.rs
@@ -1,6 +1,7 @@
 use crate::errors::DripError;
 use crate::instruction_accounts::{
-    ClosePositionAccountAccounts, InitializeVaultAccountsBumps, WithdrawAAccounts,
+    AdminWithdrawAccounts, ClosePositionAccountAccounts, InitializeVaultAccountsBumps,
+    WithdrawAAccounts,
 };
 use crate::interactions::executor::CpiExecutor;
 use crate::interactions::transfer_token::TransferToken;
@@ -31,6 +32,9 @@ pub enum Admin<'a, 'info> {
     WithdrawA {
         accounts: &'a mut WithdrawAAccounts<'info>,
     },
+    AdminWithdraw {
+        accounts: &'a mut AdminWithdrawAccounts<'info>,
+    },
     ClosePositionAccount {
         accounts: &'a mut ClosePositionAccountAccounts<'info>,
     },
@@ -127,6 +131,32 @@ impl<'a, 'info> Validatable for Admin<'a, 'info> {
                     DripError::VaultTokenAAccountIsEmpty
                 );
             }
+            Admin::AdminWithdraw { accounts } => {
+                validate!(
+                    accounts.admin.key() == accounts.vault_proto_config.admin,
+                    DripError::SignerIsNotAdmin
+                );
+
+                validate!(
+                    accounts.vault_proto_config.key() == accounts.vault.proto_config,
+                    DripError::InvalidVaultProtoConfigReference
+                );
+
+                validate!(
+                    accounts.vault_token_account.owner.key() == accounts.vault.key(),
+                    DripError::IncorrectVaultTokenAccount
+                );
+
+                validate!(
+                    accounts.vault.drip_amount == 0,
+                    DripError::CannotWithdrawAWithNonZeroDripAmount
+                );
+
+                validate!(
+                    accounts.vault_token_account.amount > 0,
+                    DripError::VaultTokenAAccountIsEmpty
+                );
+            }
             Admin::ClosePositionAccount { accounts } => {
                 validate!(
                     accounts.admin.key() == accounts.vault_proto_config.admin,
@@ -191,6 +221,18 @@ impl<'a, 'info> Executable for Admin<'a, 'info> {
                 let signer: &Vault = &accounts.vault;
                 cpi_executor.execute_all(vec![&Some(&transfer_a_to_admin)], signer)?;
             }
+            Admin::AdminWithdraw { accounts } => {
+                let withdrawal_amount = accounts.vault_token_account.amount;
+                let transfer = TransferToken::new(
+                    &accounts.token_program,
+                    &accounts.vault_token_account,
+                    &accounts.destination_token_account,
+                    &accounts.vault.to_account_info(),
+                    withdrawal_amount,
+                );
+                let signer: &Vault = &accounts.vault;
+                cpi_executor.execute_all(vec![&Some(&transfer)], signer)?;
+            }
             Admin::ClosePositionAccount { accounts } => {
                 accounts
                     .position
diff --git a/programs/drip/src/instruction_accounts/admin.rs b/programs/drip/src/instruction_accounts/admin.rs
@@ -105,6 +105,25 @@ pub struct WithdrawAAccounts<'info> {
     pub token_program: Program<'info, Token>,
 }
 
+#[derive(Accounts)]
+pub struct AdminWithdrawAccounts<'info> {
+    pub admin: Signer<'info>,
+
+    pub vault_proto_config: Account<'info, VaultProtoConfig>,
+
+    pub vault: Account<'info, Vault>,
+
+    // mut needed because we are changing state
+    #[account(mut)]
+    pub vault_token_account: Account<'info, TokenAccount>,
+
+    // mut needed because we are changing state
+    #[account(mut)]
+    pub destination_token_account: Account<'info, TokenAccount>,
+
+    pub token_program: Program<'info, Token>,
+}
+
 #[derive(Accounts)]
 pub struct ClosePositionAccountAccounts<'info> {
     pub admin: Signer<'info>,
diff --git a/programs/drip/src/lib.rs b/programs/drip/src/lib.rs
@@ -108,12 +108,21 @@ pub mod drip {
         })
     }
 
+    /*
+       DEPRECATED: USE admin_withdraw
+    */
     pub fn withdraw_a(ctx: Context<WithdrawAAccounts>) -> Result<()> {
         handle_action(Admin::WithdrawA {
             accounts: ctx.accounts,
         })
     }
 
+    pub fn admin_withdraw(ctx: Context<AdminWithdrawAccounts>) -> Result<()> {
+        handle_action(Admin::AdminWithdraw {
+            accounts: ctx.accounts,
+        })
+    }
+
     pub fn admin_close_position_account(ctx: Context<ClosePositionAccountAccounts>) -> Result<()> {
         handle_action(Admin::ClosePositionAccount {
             accounts: ctx.accounts,
diff --git a/tests/integration-tests/batch4/adminWithdraw.ts b/tests/integration-tests/batch4/adminWithdraw.ts
@@ -0,0 +1,143 @@
+import "should";
+import { initLog } from "../../utils/log.util";
+import { before } from "mocha";
+import { Mint } from "@solana/spl-token";
+import { TokenUtil } from "../../utils/token.util";
+import { Keypair, PublicKey } from "@solana/web3.js";
+import { TestUtil } from "../../utils/config.util";
+import { deployVault, deployVaultProtoConfig } from "../../utils/setup.util";
+import { findAssociatedTokenAddress } from "../../utils/common.util";
+import { VaultUtil } from "../../utils/vault.util";
+import { SolUtil } from "../../utils/sol.util";
+
+describe("#adminWithdraw", () => {
+  initLog();
+
+  let tokensAuthority: Keypair;
+  let vaultAdmin: Keypair;
+  let tokenA: Mint, tokenB: Mint;
+  let vaultProtoConfig: PublicKey;
+  let vault: PublicKey;
+  let vaultTokenAAccount: PublicKey, vaultTokenBAccount: PublicKey;
+
+  before(async () => {
+    tokensAuthority = Keypair.generate();
+    vaultAdmin = Keypair.generate();
+
+    await SolUtil.fundAccount(
+      tokensAuthority.publicKey,
+      SolUtil.solToLamports(0.1),
+    );
+    await SolUtil.fundAccount(vaultAdmin.publicKey, SolUtil.solToLamports(0.1));
+
+    [tokenA, tokenB] = await TokenUtil.createMints(
+      [tokensAuthority.publicKey, tokensAuthority.publicKey],
+      [6, 9],
+    );
+  });
+
+  beforeEach(async () => {
+    vaultProtoConfig = await deployVaultProtoConfig(
+      1,
+      5,
+      5,
+      0,
+      vaultAdmin.publicKey,
+    );
+    const vaultTreasuryTokenBAccount = await TokenUtil.createTokenAccount(
+      tokenB,
+      TestUtil.provider.publicKey,
+      tokensAuthority,
+    );
+
+    const vaultPDA = await deployVault(
+      tokenA.address,
+      tokenB.address,
+      vaultTreasuryTokenBAccount,
+      vaultProtoConfig,
+      undefined,
+      vaultAdmin,
+    );
+
+    vault = vaultPDA.publicKey;
+    vaultTokenAAccount = await findAssociatedTokenAddress(
+      vaultPDA.publicKey,
+      tokenA.address,
+    );
+    vaultTokenBAccount = await findAssociatedTokenAddress(
+      vaultPDA.publicKey,
+      tokenB.address,
+    );
+    await TokenUtil.mintTo({
+      payer: tokensAuthority,
+      token: tokenA,
+      mintAuthority: tokensAuthority,
+      recipient: vaultTokenAAccount,
+      amount: BigInt(1_000_000_000),
+    });
+    await TokenUtil.mintTo({
+      payer: tokensAuthority,
+      token: tokenB,
+      mintAuthority: tokensAuthority,
+      recipient: vaultTokenBAccount,
+      amount: BigInt(1_000_000_000),
+    });
+  });
+
+  it("allows admin to withdraw funds to admin's token B account", async () => {
+    const adminTokenAccount = await TokenUtil.getOrCreateAssociatedTokenAccount(
+      tokenB,
+      vaultAdmin.publicKey,
+      tokensAuthority,
+    );
+    const adminTokenBBalanceBefore =
+      await TokenUtil.getTokenAccount(adminTokenAccount);
+    adminTokenBBalanceBefore.amount.toString().should.equal("0");
+    await VaultUtil.adminWithdraw(
+      vault,
+      vaultTokenBAccount,
+      adminTokenAccount,
+      vaultProtoConfig,
+      vaultAdmin,
+    );
+    const adminTokenBBalanceAfter =
+      await TokenUtil.getTokenAccount(adminTokenAccount);
+    adminTokenBBalanceAfter.amount.toString().should.equal("1000000000");
+  });
+
+  it("allows admin to withdraw funds to admin's token A account", async () => {
+    const adminTokenAccount = await TokenUtil.getOrCreateAssociatedTokenAccount(
+      tokenA,
+      vaultAdmin.publicKey,
+      tokensAuthority,
+    );
+    const adminTokenABalanceBefore =
+      await TokenUtil.getTokenAccount(adminTokenAccount);
+    adminTokenABalanceBefore.amount.toString().should.equal("0");
+    await VaultUtil.adminWithdraw(
+      vault,
+      vaultTokenAAccount,
+      adminTokenAccount,
+      vaultProtoConfig,
+      vaultAdmin,
+    );
+    const adminTokenABalanceAfter =
+      await TokenUtil.getTokenAccount(adminTokenAccount);
+    adminTokenABalanceAfter.amount.toString().should.equal("1000000000");
+  });
+
+  it("does not allow non-admin to withdraw", async () => {
+    const adminTokenAccount = await TokenUtil.getOrCreateAssociatedTokenAccount(
+      tokenB,
+      vaultAdmin.publicKey,
+      tokensAuthority,
+    );
+    await VaultUtil.adminWithdraw(
+      vault,
+      vaultTokenBAccount,
+      adminTokenAccount,
+      vaultProtoConfig,
+      tokensAuthority,
+    ).should.be.rejectedWith(/0x1785/);
+  });
+});
diff --git a/tests/utils/vault.util.ts b/tests/utils/vault.util.ts
@@ -16,7 +16,7 @@ import {
   Granularity,
   PDA,
 } from "./common.util";
-import { BN } from "@coral-xyz/anchor";
+import { BN, Program } from "@coral-xyz/anchor";
 import { TokenUtil } from "./token.util";
 import { deployVaultPeriod, depositToVault } from "./setup.util";
 import { SolUtil } from "./sol.util";
@@ -192,9 +192,7 @@ export class VaultUtil extends TestUtil {
       .transaction();
     if (admin) {
       const blockhash = await TestUtil.provider.connection.getLatestBlockhash();
-      const txId = await TestUtil.provider.connection.sendTransaction(tx, [
-        admin,
-      ]);
+      const txId = await TestUtil.provider.sendAndConfirm(tx, [admin]);
       await TestUtil.provider.connection.confirmTransaction(
         {
           signature: txId,
@@ -469,6 +467,27 @@ export class VaultUtil extends TestUtil {
     return this.provider.sendAndConfirm(tx, [withdrawer]);
   }
 
+  static async adminWithdraw(
+    vault: PublicKey,
+    vaultTokenAccount: PublicKey,
+    destinationTokenAccount: PublicKey,
+    vaultProtoConfig: PublicKey,
+    admin: Keypair | Signer,
+  ): Promise<TransactionSignature> {
+    const tx = await ProgramUtil.dripProgram.methods
+      .adminWithdraw()
+      .accounts({
+        admin: admin.publicKey,
+        vaultProtoConfig,
+        vault,
+        vaultTokenAccount,
+        destinationTokenAccount,
+        tokenProgram: ProgramUtil.tokenProgram.programId,
+      })
+      .transaction();
+    return this.provider.sendAndConfirm(tx, [admin]);
+  }
+
   static async withdrawA(
     vault: PublicKey,
     vaultTokenAAccount: PublicKey,
@@ -491,7 +510,7 @@ export class VaultUtil extends TestUtil {
       .transaction();
 
     if (admin) {
-      return this.provider.connection.sendTransaction(tx, [admin]);
+      return this.provider.sendAndConfirm(tx, [admin]);
     }
 
     return this.provider.sendAndConfirm(tx);
