Fix comments

programs/locker/src/instructions/create_vesting_escrow.rs

@@ -2,11 +2,9 @@ use anchor_spl::token::{Token, TokenAccount};
 
 use crate::*;
 use crate::safe_math::SafeMath;
+use crate::TokenProgramFLag::UseSplToken;
 use crate::util::token::transfer_to_escrow;
 
-pub const USE_SPL_TOKEN_PROGRAM: u8 = 0;
-pub const USE_TOKEN_2022_PROGRAM: u8 = 1;
-
 #[derive(AnchorSerialize, AnchorDeserialize)]
 /// Accounts for [locker::create_vesting_escrow].
 pub struct CreateVestingEscrowParameters {
@@ -47,7 +45,7 @@ impl CreateVestingEscrowParameters {
         sender: Pubkey,
         base: Pubkey,
         escrow_bump: u8,
-        token_program_flag: u8,
+        token_program_flag: TokenProgramFLag,
     ) -> Result<()> {
         self.validate()?;
 
@@ -130,7 +128,7 @@ pub fn handle_create_vesting_escrow(
         ctx.accounts.sender.key(),
         ctx.accounts.base.key(),
         ctx.bumps.escrow,
-        USE_SPL_TOKEN_PROGRAM,
+        UseSplToken,
     )?;
 
     transfer_to_escrow(

programs/locker/src/instructions/v2/create_vesting_escrow.rs

@@ -1,8 +1,10 @@
+use anchor_spl::token::Token;
 use anchor_spl::token_interface::{
     Mint, TokenAccount, TokenInterface,
 };
 
 use crate::*;
+use crate::TokenProgramFLag::UseSplToken;
 use crate::util::{AccountsType, calculate_transfer_fee_included_amount, parse_remaining_accounts, ParsedRemainingAccounts, transfer_to_escrow_v2, validate_mint};
 
 #[event_cpi]
@@ -69,14 +71,20 @@ pub fn handle_create_vesting_escrow_v2<'c: 'info, 'info>(
         LockerError::InvalidEscrowTokenAddress
     );
 
+    let token_mint_info = ctx.accounts.mint.to_account_info();
+    let token_program_flag = if *token_mint_info.owner == Token::id() {
+        UseSplToken
+    } else {
+        TokenProgramFLag::UseToken2022
+    };
     params.init_escrow(
         &ctx.accounts.escrow,
         ctx.accounts.recipient.key(),
         ctx.accounts.sender_token.mint,
         ctx.accounts.sender.key(),
         ctx.accounts.base.key(),
         ctx.bumps.escrow,
-        USE_TOKEN_2022_PROGRAM,
+        token_program_flag,
     )?;
 
     // Process remaining accounts

programs/locker/src/state/vesting_escrow.rs

@@ -14,6 +14,14 @@ pub enum UpdateRecipientMode {
     EitherCreatorAndRecipient, //3
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, IntoPrimitive, TryFromPrimitive)]
+#[repr(u8)]
+pub enum TokenProgramFLag {
+    UseSplToken,   //0
+    UseToken2022,  //1
+}
+
+
 #[account(zero_copy)]
 #[derive(Default, InitSpace, Debug)]
 pub struct VestingEscrow {
@@ -68,7 +76,7 @@ impl VestingEscrow {
         base: Pubkey,
         escrow_bump: u8,
         update_recipient_mode: u8,
-        token_program_flag: u8,
+        token_program_flag: TokenProgramFLag,
     ) {
         self.vesting_start_time = vesting_start_time;
         self.cliff_time = cliff_time;
@@ -82,7 +90,7 @@ impl VestingEscrow {
         self.base = base;
         self.escrow_bump = escrow_bump;
         self.update_recipient_mode = update_recipient_mode;
-        self.token_program_flag = token_program_flag;
+        self.token_program_flag = token_program_flag as u8;
     }
 
     pub fn get_max_unlocked_amount(&self, current_ts: u64) -> Result<u64> {

tests/common/index.ts

@@ -1,5 +1,11 @@
-import { web3, Wallet } from "@coral-xyz/anchor";
-import { ASSOCIATED_TOKEN_PROGRAM_ID, createMint, getOrCreateAssociatedTokenAccount, mintTo, TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import { Wallet, web3 } from "@coral-xyz/anchor";
+import {
+  ASSOCIATED_TOKEN_PROGRAM_ID,
+  createMint,
+  getOrCreateAssociatedTokenAccount,
+  mintTo,
+  TOKEN_PROGRAM_ID,
+} from "@solana/spl-token";
 
 export * from "./asserter";
 
@@ -69,7 +75,13 @@ export const createMintIfNotExists = async (
   }
   return mintKeypair.publicKey;
 };
-export const mintTokenTo = async (connection: web3.Connection, mintAuthority: web3.Keypair, tokenMint: web3.PublicKey, owner: web3.PublicKey, amount: number) => {
+export const mintTokenTo = async (
+  connection: web3.Connection,
+  mintAuthority: web3.Keypair,
+  tokenMint: web3.PublicKey,
+  owner: web3.PublicKey,
+  amount: number
+) => {
   const userToken = await getOrCreateAssociatedTokenAccount(
     connection,
     mintAuthority,
@@ -97,11 +109,11 @@ export const mintTokenTo = async (connection: web3.Connection, mintAuthority: we
     },
     TOKEN_PROGRAM_ID
   );
-}
+};
 
 export const getCurrentBlockTime = async (connection: web3.Connection) => {
   // TODO fetch clock account can help to reduce rpc call
   const currentSlot = await connection.getSlot();
   const currentBlockTime = await connection.getBlockTime(currentSlot);
   return currentBlockTime;
-}
+};

tests/v2/locker-spl-token.ts

@@ -0,0 +1,153 @@
+import * as anchor from "@coral-xyz/anchor";
+import { web3 } from "@coral-xyz/anchor";
+import {
+  createAssociatedTokenAccountIdempotent,
+  createInitializeMint2Instruction,
+  mintTo,
+  TOKEN_PROGRAM_ID,
+} from "@solana/spl-token";
+import { BN } from "bn.js";
+import { createAndFundWallet, getCurrentBlockTime, sleep } from "../common";
+import {
+  claimToken,
+  createLockerProgram,
+  createVestingPlan,
+} from "./locker_utils";
+import {
+  sendAndConfirmTransaction,
+  SystemProgram,
+  Transaction,
+} from "@solana/web3.js";
+
+const provider = anchor.AnchorProvider.env();
+
+describe("[V2] Full flow With SPL Token", () => {
+  const tokenDecimal = 8;
+  let mintAuthority: web3.Keypair;
+  let mintKeypair: web3.Keypair;
+  let TOKEN: web3.PublicKey;
+
+  let UserKP: web3.Keypair;
+  let RecipientKP: web3.Keypair;
+  let RecipientToken: web3.PublicKey;
+
+  let mintAmount: bigint;
+
+  before(async () => {
+    {
+      const result = await createAndFundWallet(provider.connection);
+      UserKP = result.keypair;
+    }
+    {
+      const result = await createAndFundWallet(provider.connection);
+      RecipientKP = result.keypair;
+    }
+
+    mintAuthority = new web3.Keypair();
+    mintKeypair = new web3.Keypair();
+    TOKEN = mintKeypair.publicKey;
+
+    mintAmount = BigInt(1_000_000 * Math.pow(10, tokenDecimal)); // Mint 1,000,000 tokens
+
+    // Step 2 - Create a New Token
+    const mintLamports =
+      await provider.connection.getMinimumBalanceForRentExemption(82);
+    const mintTransaction = new Transaction().add(
+      SystemProgram.createAccount({
+        fromPubkey: UserKP.publicKey,
+        newAccountPubkey: TOKEN,
+        space: 82,
+        lamports: mintLamports,
+        programId: TOKEN_PROGRAM_ID,
+      }),
+      createInitializeMint2Instruction(
+        TOKEN, // Mint account
+        tokenDecimal, // Decimals
+        mintAuthority.publicKey, // Mint authority
+        null, // Freeze authority
+        TOKEN_PROGRAM_ID // Token program ID
+      )
+    );
+    await sendAndConfirmTransaction(
+      provider.connection,
+      mintTransaction,
+      [UserKP, mintKeypair],
+      undefined
+    );
+
+    const userToken = await createAssociatedTokenAccountIdempotent(
+      provider.connection,
+      UserKP,
+      TOKEN,
+      UserKP.publicKey,
+      {},
+      TOKEN_PROGRAM_ID
+    );
+
+    await mintTo(
+      provider.connection,
+      UserKP,
+      TOKEN,
+      userToken,
+      mintAuthority,
+      mintAmount,
+      [],
+      undefined,
+      TOKEN_PROGRAM_ID
+    );
+
+    RecipientToken = await createAssociatedTokenAccountIdempotent(
+      provider.connection,
+      UserKP,
+      TOKEN,
+      RecipientKP.publicKey,
+      {},
+      TOKEN_PROGRAM_ID
+    );
+  });
+
+  it("Full flow", async () => {
+    console.log("Create vesting plan");
+    const program = createLockerProgram(new anchor.Wallet(UserKP));
+    let currentBlockTime = await getCurrentBlockTime(
+      program.provider.connection
+    );
+    const cliffTime = new BN(currentBlockTime).add(new BN(5));
+    let escrow = await createVestingPlan({
+      ownerKeypair: UserKP,
+      vestingStartTime: new BN(0),
+      tokenMint: TOKEN,
+      isAssertion: true,
+      cliffTime,
+      frequency: new BN(1),
+      cliffUnlockAmount: new BN(100_000),
+      amountPerPeriod: new BN(50_000),
+      numberOfPeriod: new BN(2),
+      recipient: RecipientKP.publicKey,
+      updateRecipientMode: 0,
+      tokenProgram: TOKEN_PROGRAM_ID,
+    });
+    while (true) {
+      const currentBlockTime = await getCurrentBlockTime(
+        program.provider.connection
+      );
+      if (currentBlockTime > cliffTime.toNumber()) {
+        break;
+      } else {
+        await sleep(1000);
+        console.log("Wait until cliffTime");
+      }
+    }
+
+    console.log("Claim token");
+    await claimToken({
+      recipient: RecipientKP,
+      recipientToken: RecipientToken,
+      tokenMint: TOKEN,
+      escrow,
+      maxAmount: new BN(1_000_000),
+      isAssertion: true,
+      tokenProgram: TOKEN_PROGRAM_ID,
+    });
+  });
+});

tests/v2/locker.ts

@@ -25,7 +25,7 @@ import {
 
 const provider = anchor.AnchorProvider.env();
 
-describe("[V2] Test full flow", () => {
+describe("[V2] Test full flow With token 2022", () => {
   let TOKEN: web3.PublicKey;
   let UserKP: web3.Keypair;
   let RecipientKP: web3.Keypair;
@@ -61,7 +61,7 @@ describe("[V2] Test full flow", () => {
     );
   });
 
-  it("Full flow With token 2022", async () => {
+  it("Full flow", async () => {
     console.log("Create vesting plan");
     const program = createLockerProgram(new anchor.Wallet(UserKP));
     let currentBlockTime = await getCurrentBlockTime(

tests/v2/locker_utils/index.ts

@@ -4,6 +4,7 @@ import {
   ASSOCIATED_TOKEN_PROGRAM_ID,
   createAssociatedTokenAccountInstruction,
   getAssociatedTokenAddressSync,
+  TOKEN_PROGRAM_ID,
 } from "@solana/spl-token";
 import { expect } from "chai";
 
@@ -158,7 +159,9 @@ export async function createVestingPlan(params: CreateVestingPlanParams) {
     );
     expect(escrowState.base.toString()).eq(baseKP.publicKey.toString());
     expect(escrowState.updateRecipientMode).eq(updateRecipientMode);
-    expect(escrowState.tokenProgramFlag).eq(1);
+    expect(escrowState.tokenProgramFlag).eq(
+      tokenProgram == TOKEN_PROGRAM_ID ? 0 : 1
+    );
   }
 
   return escrow;