From b4e0eb69e317ee563f62ab740f564ac04294081f Mon Sep 17 00:00:00 2001 From: ananas-block Date: Sun, 1 Dec 2024 20:34:07 +0000 Subject: [PATCH] feat: restrict allowed token 2022 extensions --- Cargo.lock | 1 + Cargo.toml | 1 + .../src/idl/light_compressed_token.ts | 96 ++++++++++++- .../src/idls/light_compressed_token.ts | 96 ++++++++++++- programs/compressed-token/Cargo.toml | 2 +- .../src/instructions/create_token_pool.rs | 27 ++++ programs/compressed-token/src/lib.rs | 6 +- programs/compressed-token/src/process_mint.rs | 23 ++- .../compressed-token-test/tests/test.rs | 131 +++++++++++++++++- test-utils/Cargo.toml | 2 +- 10 files changed, 361 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9c85d5bd62..9077d2527f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3545,6 +3545,7 @@ dependencies = [ "solana-sdk", "solana-security-txt", "spl-token", + "spl-token-2022 3.0.4", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 03b788d30b..95fdbbe41b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ solana-transaction-status = "=1.18.22" solana-account-decoder = "=1.18.22" solana-rpc = "=1.18.22" spl-token = "=4.0.0" +spl-token-2022 = {version="3.0.0", no-default-features = true, features = ["no-entrypoint"]} # Anchor anchor-lang = "=0.29.0" diff --git a/js/compressed-token/src/idl/light_compressed_token.ts b/js/compressed-token/src/idl/light_compressed_token.ts index 7e4de98b5f..5cc34dfcd8 100644 --- a/js/compressed-token/src/idl/light_compressed_token.ts +++ b/js/compressed-token/src/idl/light_compressed_token.ts @@ -45,6 +45,43 @@ export type LightCompressedToken = { ]; args: []; }, + { + name: 'createTokenPool2022'; + accounts: [ + { + name: 'feePayer'; + isMut: true; + isSigner: true; + docs: ['UNCHECKED: only pays fees.']; + }, + { + name: 'tokenPoolPda'; + isMut: true; + isSigner: false; + }, + { + name: 'systemProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'mint'; + isMut: true; + isSigner: false; + }, + { + name: 'tokenProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'cpiAuthorityPda'; + isMut: false; + isSigner: false; + }, + ]; + args: []; + }, { name: 'mintTo'; docs: [ @@ -77,14 +114,15 @@ export type LightCompressedToken = { name: 'mint'; isMut: true; isSigner: false; + docs: [ + 'mint is written to and we check the program id we invoke to be either', + 'spl_token::ID or token_2022::ID.', + ]; }, { name: 'tokenPoolPda'; isMut: true; isSigner: false; - docs: [ - 'account to a token account of a different mint will fail', - ]; }, { name: 'tokenProgram'; @@ -521,6 +559,7 @@ export type LightCompressedToken = { name: 'authority'; isMut: false; isSigner: true; + docs: ['check_mint_and_freeze_authority().']; }, { name: 'cpiAuthorityPda'; @@ -567,6 +606,7 @@ export type LightCompressedToken = { name: 'mint'; isMut: false; isSigner: false; + docs: ['check_mint_and_freeze_authority().']; }, ]; args: [ @@ -593,6 +633,7 @@ export type LightCompressedToken = { name: 'authority'; isMut: false; isSigner: true; + docs: ['check_mint_and_freeze_authority().']; }, { name: 'cpiAuthorityPda'; @@ -639,6 +680,7 @@ export type LightCompressedToken = { name: 'mint'; isMut: false; isSigner: false; + docs: ['check_mint_and_freeze_authority().']; }, ]; args: [ @@ -1716,6 +1758,43 @@ export const IDL: LightCompressedToken = { ], args: [], }, + { + name: 'createTokenPool2022', + accounts: [ + { + name: 'feePayer', + isMut: true, + isSigner: true, + docs: ['UNCHECKED: only pays fees.'], + }, + { + name: 'tokenPoolPda', + isMut: true, + isSigner: false, + }, + { + name: 'systemProgram', + isMut: false, + isSigner: false, + }, + { + name: 'mint', + isMut: true, + isSigner: false, + }, + { + name: 'tokenProgram', + isMut: false, + isSigner: false, + }, + { + name: 'cpiAuthorityPda', + isMut: false, + isSigner: false, + }, + ], + args: [], + }, { name: 'mintTo', docs: [ @@ -1748,14 +1827,15 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: true, isSigner: false, + docs: [ + 'mint is written to and we check the program id we invoke to be either', + 'spl_token::ID or token_2022::ID.', + ], }, { name: 'tokenPoolPda', isMut: true, isSigner: false, - docs: [ - 'account to a token account of a different mint will fail', - ], }, { name: 'tokenProgram', @@ -2192,6 +2272,7 @@ export const IDL: LightCompressedToken = { name: 'authority', isMut: false, isSigner: true, + docs: ['check_mint_and_freeze_authority().'], }, { name: 'cpiAuthorityPda', @@ -2238,6 +2319,7 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: false, isSigner: false, + docs: ['check_mint_and_freeze_authority().'], }, ], args: [ @@ -2264,6 +2346,7 @@ export const IDL: LightCompressedToken = { name: 'authority', isMut: false, isSigner: true, + docs: ['check_mint_and_freeze_authority().'], }, { name: 'cpiAuthorityPda', @@ -2310,6 +2393,7 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: false, isSigner: false, + docs: ['check_mint_and_freeze_authority().'], }, ], args: [ diff --git a/js/stateless.js/src/idls/light_compressed_token.ts b/js/stateless.js/src/idls/light_compressed_token.ts index 7e4de98b5f..5cc34dfcd8 100644 --- a/js/stateless.js/src/idls/light_compressed_token.ts +++ b/js/stateless.js/src/idls/light_compressed_token.ts @@ -45,6 +45,43 @@ export type LightCompressedToken = { ]; args: []; }, + { + name: 'createTokenPool2022'; + accounts: [ + { + name: 'feePayer'; + isMut: true; + isSigner: true; + docs: ['UNCHECKED: only pays fees.']; + }, + { + name: 'tokenPoolPda'; + isMut: true; + isSigner: false; + }, + { + name: 'systemProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'mint'; + isMut: true; + isSigner: false; + }, + { + name: 'tokenProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'cpiAuthorityPda'; + isMut: false; + isSigner: false; + }, + ]; + args: []; + }, { name: 'mintTo'; docs: [ @@ -77,14 +114,15 @@ export type LightCompressedToken = { name: 'mint'; isMut: true; isSigner: false; + docs: [ + 'mint is written to and we check the program id we invoke to be either', + 'spl_token::ID or token_2022::ID.', + ]; }, { name: 'tokenPoolPda'; isMut: true; isSigner: false; - docs: [ - 'account to a token account of a different mint will fail', - ]; }, { name: 'tokenProgram'; @@ -521,6 +559,7 @@ export type LightCompressedToken = { name: 'authority'; isMut: false; isSigner: true; + docs: ['check_mint_and_freeze_authority().']; }, { name: 'cpiAuthorityPda'; @@ -567,6 +606,7 @@ export type LightCompressedToken = { name: 'mint'; isMut: false; isSigner: false; + docs: ['check_mint_and_freeze_authority().']; }, ]; args: [ @@ -593,6 +633,7 @@ export type LightCompressedToken = { name: 'authority'; isMut: false; isSigner: true; + docs: ['check_mint_and_freeze_authority().']; }, { name: 'cpiAuthorityPda'; @@ -639,6 +680,7 @@ export type LightCompressedToken = { name: 'mint'; isMut: false; isSigner: false; + docs: ['check_mint_and_freeze_authority().']; }, ]; args: [ @@ -1716,6 +1758,43 @@ export const IDL: LightCompressedToken = { ], args: [], }, + { + name: 'createTokenPool2022', + accounts: [ + { + name: 'feePayer', + isMut: true, + isSigner: true, + docs: ['UNCHECKED: only pays fees.'], + }, + { + name: 'tokenPoolPda', + isMut: true, + isSigner: false, + }, + { + name: 'systemProgram', + isMut: false, + isSigner: false, + }, + { + name: 'mint', + isMut: true, + isSigner: false, + }, + { + name: 'tokenProgram', + isMut: false, + isSigner: false, + }, + { + name: 'cpiAuthorityPda', + isMut: false, + isSigner: false, + }, + ], + args: [], + }, { name: 'mintTo', docs: [ @@ -1748,14 +1827,15 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: true, isSigner: false, + docs: [ + 'mint is written to and we check the program id we invoke to be either', + 'spl_token::ID or token_2022::ID.', + ], }, { name: 'tokenPoolPda', isMut: true, isSigner: false, - docs: [ - 'account to a token account of a different mint will fail', - ], }, { name: 'tokenProgram', @@ -2192,6 +2272,7 @@ export const IDL: LightCompressedToken = { name: 'authority', isMut: false, isSigner: true, + docs: ['check_mint_and_freeze_authority().'], }, { name: 'cpiAuthorityPda', @@ -2238,6 +2319,7 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: false, isSigner: false, + docs: ['check_mint_and_freeze_authority().'], }, ], args: [ @@ -2264,6 +2346,7 @@ export const IDL: LightCompressedToken = { name: 'authority', isMut: false, isSigner: true, + docs: ['check_mint_and_freeze_authority().'], }, { name: 'cpiAuthorityPda', @@ -2310,6 +2393,7 @@ export const IDL: LightCompressedToken = { name: 'mint', isMut: false, isSigner: false, + docs: ['check_mint_and_freeze_authority().'], }, ], args: [ diff --git a/programs/compressed-token/Cargo.toml b/programs/compressed-token/Cargo.toml index c8112692a9..54e04a2859 100644 --- a/programs/compressed-token/Cargo.toml +++ b/programs/compressed-token/Cargo.toml @@ -35,7 +35,7 @@ solana-security-txt = "1.1.0" light-hasher = { version = "1.1.0", path = "../../merkle-tree/hasher" } light-heap = { version = "1.1.0", path = "../../heap", optional = true } light-utils = { version = "1.1.0", path = "../../utils" } - +spl-token-2022 = { workspace = true } [target.'cfg(not(target_os = "solana"))'.dependencies] solana-sdk = { workspace = true } diff --git a/programs/compressed-token/src/instructions/create_token_pool.rs b/programs/compressed-token/src/instructions/create_token_pool.rs index 319e2553b6..5cc428da3a 100644 --- a/programs/compressed-token/src/instructions/create_token_pool.rs +++ b/programs/compressed-token/src/instructions/create_token_pool.rs @@ -5,6 +5,11 @@ use anchor_spl::{ token_2022::Token2022, token_interface::{Mint as Mint22, TokenAccount as Token22Account}, }; +use spl_token_2022::{ + extension::{BaseStateWithExtensions, ExtensionType, PodStateWithExtensions}, + pod::PodMint, +}; + pub const POOL_SEED: &[u8] = b"pool"; /// Creates a spl token pool account which is owned by the token authority pda. @@ -66,3 +71,25 @@ pub fn get_token_pool_pda(mint: &Pubkey) -> Pubkey { let (address, _) = Pubkey::find_program_address(seeds, &crate::ID); address } + +// cpi guard could be ok but should be tested +const ALLOWED_EXTENSION_TYPES: [ExtensionType; 6] = [ + ExtensionType::MetadataPointer, + ExtensionType::InterestBearingConfig, + ExtensionType::GroupPointer, + ExtensionType::GroupMemberPointer, + ExtensionType::TokenGroup, + ExtensionType::TokenGroupMember, +]; + +pub fn assert_mint_extensions(account_data: &[u8]) -> Result<()> { + let mint = PodStateWithExtensions::::unpack(account_data).unwrap(); + let mint_extensions = mint.get_extension_types().unwrap(); + if !mint_extensions + .iter() + .all(|item| ALLOWED_EXTENSION_TYPES.contains(item)) + { + return err!(crate::ErrorCode::MintWithInvalidExtension); + } + Ok(()) +} diff --git a/programs/compressed-token/src/lib.rs b/programs/compressed-token/src/lib.rs index d24041421c..6b1dddbe26 100644 --- a/programs/compressed-token/src/lib.rs +++ b/programs/compressed-token/src/lib.rs @@ -45,8 +45,11 @@ pub mod light_compressed_token { } pub fn create_token_pool_2022<'info>( - _ctx: Context<'_, '_, '_, 'info, CreateTokenPoolInstruction2022<'info>>, + ctx: Context<'_, '_, '_, 'info, CreateTokenPoolInstruction2022<'info>>, ) -> Result<()> { + create_token_pool::assert_mint_extensions( + &ctx.accounts.mint.to_account_info().try_borrow_data()?, + )?; Ok(()) } @@ -210,4 +213,5 @@ pub enum ErrorCode { MintHasNoFreezeAuthority, InvalidTokenProgram, InvalidTokenMintOwner, + MintWithInvalidExtension, } diff --git a/programs/compressed-token/src/process_mint.rs b/programs/compressed-token/src/process_mint.rs index 32e372eb79..f0c71c6085 100644 --- a/programs/compressed-token/src/process_mint.rs +++ b/programs/compressed-token/src/process_mint.rs @@ -1,7 +1,9 @@ use account_compression::{program::AccountCompression, utils::constants::CPI_AUTHORITY_PDA_SEED}; -use anchor_lang::prelude::*; -use anchor_spl::{token::TokenAccount, token_2022}; - +use anchor_lang::{prelude::*, solana_program::program_option::COption}; +use anchor_spl::{ + token::{Mint, TokenAccount}, + token_2022, +}; use light_system_program::{program::LightSystemProgram, OutputCompressedAccountWithPackedContext}; use crate::{program::LightCompressedToken, spl_compression::spl_token_pool_derivation}; @@ -287,6 +289,17 @@ pub fn mint_spl_to_pool_pda(ctx: &Context, amounts: &[u64]) - .checked_add(*amount) .ok_or(crate::ErrorCode::MintTooLarge)?; } + + if let COption::Some(mint_authority) = + Mint::try_deserialize(&mut &ctx.accounts.mint.data.borrow()[..])?.mint_authority + { + if mint_authority != ctx.accounts.authority.key() { + return err!(crate::ErrorCode::InvalidAuthorityMint); + } + } else { + return err!(crate::ErrorCode::InvalidAuthorityMint); + } + let pre_token_balance = TokenAccount::try_deserialize(&mut &ctx.accounts.token_pool_pda.data.borrow()[..])?.amount; match ctx.accounts.token_program.key() { @@ -340,9 +353,7 @@ pub struct MintToInstruction<'info> { /// CHECK: #[account(seeds = [CPI_AUTHORITY_PDA_SEED], bump)] pub cpi_authority_pda: UncheckedAccount<'info>, - /// CHECK: by the token program token program internal checks, since the - /// mint is written to and we check the program id we invoke to be either - /// spl_token::ID or token_2022::ID. + /// CHECK: in mint_spl_to_pool_pda(). #[account(mut)] pub mint: AccountInfo<'info>, /// CHECK: with spl_token_pool_derivation(). diff --git a/test-programs/compressed-token-test/tests/test.rs b/test-programs/compressed-token-test/tests/test.rs index 73b0b00ec7..d492e2ca50 100644 --- a/test-programs/compressed-token-test/tests/test.rs +++ b/test-programs/compressed-token-test/tests/test.rs @@ -3,8 +3,9 @@ use anchor_lang::{ system_program, AnchorDeserialize, AnchorSerialize, InstructionData, ToAccountMetas, }; -use anchor_spl::token::Mint; +use anchor_spl::token::{Mint, TokenAccount}; use anchor_spl::token_2022::spl_token_2022; +use anchor_spl::token_2022::spl_token_2022::extension::ExtensionType; use light_compressed_token::delegation::sdk::{ create_approve_instruction, create_revoke_instruction, CreateApproveInstructionInputs, CreateRevokeInstructionInputs, @@ -17,6 +18,7 @@ use light_compressed_token::mint_sdk::{ }; use light_compressed_token::process_transfer::transfer_sdk::create_transfer_instruction; use light_compressed_token::process_transfer::{get_cpi_authority_pda, TokenTransferOutputData}; +use light_compressed_token::spl_compression::spl_token_pool_derivation; use light_compressed_token::token_data::AccountState; use light_compressed_token::{token_data::TokenData, ErrorCode}; use light_prover_client::gnark::helpers::{kill_prover, spawn_prover, ProofType, ProverConfig}; @@ -50,6 +52,7 @@ use light_test_utils::{ }; use light_verifier::VerifierError; use rand::Rng; +use solana_sdk::system_instruction; use solana_sdk::{ instruction::{Instruction, InstructionError}, pubkey::Pubkey, @@ -178,6 +181,118 @@ async fn test_failing_create_token_pool() { ) .unwrap(); } + // failing test try to create a token pool with mint with non-whitelisted token extension + { + let payer = rpc.get_payer().insecure_clone(); + let payer_pubkey = payer.pubkey(); + let mint = Keypair::new(); + let token_authority = payer.insecure_clone(); + let space = ExtensionType::try_calculate_account_len::(&[ + ExtensionType::MintCloseAuthority, + ]) + .unwrap(); + + let mut instructions = vec![system_instruction::create_account( + &payer.pubkey(), + &mint.pubkey(), + rpc.get_minimum_balance_for_rent_exemption(space) + .await + .unwrap(), + space as u64, + &spl_token_2022::ID, + )]; + let invalid_token_extension_ix = + spl_token_2022::instruction::initialize_mint_close_authority( + &spl_token_2022::ID, + &mint.pubkey(), + Some(&token_authority.pubkey()), + ) + .unwrap(); + instructions.push(invalid_token_extension_ix); + instructions.push( + spl_token_2022::instruction::initialize_mint( + &spl_token_2022::ID, + &mint.pubkey(), + &token_authority.pubkey(), + None, + 2, + ) + .unwrap(), + ); + instructions.push(create_create_token_pool_2022_instruction( + &payer_pubkey, + &mint.pubkey(), + )); + + let result = rpc + .create_and_send_transaction(&instructions, &payer_pubkey, &[&payer, &mint]) + .await; + assert_rpc_error( + result, + 3, + light_compressed_token::ErrorCode::MintWithInvalidExtension.into(), + ) + .unwrap(); + } + // functional create token pool account with token 2022 mint with allowed metadata pointer extension + { + let payer = rpc.get_payer().insecure_clone(); + // create_mint_helper(&mut rpc, &payer).await; + let payer_pubkey = payer.pubkey(); + + let mint = Keypair::new(); + let token_authority = payer.insecure_clone(); + let space = ExtensionType::try_calculate_account_len::(&[ + ExtensionType::MetadataPointer, + ]) + .unwrap(); + + let mut instructions = vec![system_instruction::create_account( + &payer.pubkey(), + &mint.pubkey(), + rpc.get_minimum_balance_for_rent_exemption(space) + .await + .unwrap(), + space as u64, + &spl_token_2022::ID, + )]; + let token_extension_ix = + spl_token_2022::extension::metadata_pointer::instruction::initialize( + &spl_token_2022::ID, + &mint.pubkey(), + Some(token_authority.pubkey()), + None, + ) + .unwrap(); + instructions.push(token_extension_ix); + instructions.push( + spl_token_2022::instruction::initialize_mint( + &spl_token_2022::ID, + &mint.pubkey(), + &token_authority.pubkey(), + None, + 2, + ) + .unwrap(), + ); + instructions.push(create_create_token_pool_2022_instruction( + &payer_pubkey, + &mint.pubkey(), + )); + rpc.create_and_send_transaction(&instructions, &payer_pubkey, &[&payer, &mint]) + .await + .unwrap(); + + let token_pool_pubkey = get_token_pool_pda(&mint.pubkey()); + let token_pool_account = rpc.get_account(token_pool_pubkey).await.unwrap().unwrap(); + spl_token_pool_derivation( + &mint.pubkey(), + &light_compressed_token::ID, + &token_pool_pubkey, + ) + .unwrap(); + assert_eq!(token_pool_account.data.len(), TokenAccount::LEN); + } } #[tokio::test] @@ -542,7 +657,12 @@ async fn test_mint_to_failing() { .create_and_send_transaction(&[instruction], &payer_2.pubkey(), &[&payer_2]) .await; // Owner doesn't match the mint authority. - assert_rpc_error(result, 0, 4).unwrap(); + assert_rpc_error( + result, + 0, + light_compressed_token::ErrorCode::InvalidAuthorityMint.into(), + ) + .unwrap(); } // 2. Try to mint token from `mint_2` and sign the transaction with `mint_1` // authority. @@ -561,7 +681,12 @@ async fn test_mint_to_failing() { .create_and_send_transaction(&[instruction], &payer_1.pubkey(), &[&payer_1]) .await; // Owner doesn't match the mint authority. - assert_rpc_error(result, 0, 4).unwrap(); + assert_rpc_error( + result, + 0, + light_compressed_token::ErrorCode::InvalidAuthorityMint.into(), + ) + .unwrap(); } // 3. Try to mint token to random token account. { diff --git a/test-utils/Cargo.toml b/test-utils/Cargo.toml index 94e50c581b..2931e44c03 100644 --- a/test-utils/Cargo.toml +++ b/test-utils/Cargo.toml @@ -46,7 +46,7 @@ log = "0.4" serde = { version = "1.0.197", features = ["derive"] } async-trait = "0.1.82" light-client = { workspace = true } -spl-token-2022 = "3.0.0" +spl-token-2022 = { workspace = true } [dev-dependencies] rand = "0.8"