Skip to content

fix: revert compressed token account type matching based on tree type #2158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ananas-block wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: revert compressed token account type matching based on tree type #2158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions program-tests/compressed-token-test/tests/v1.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1235,8 +1235,9 @@ async fn test_mint_to_failing() {
.create_and_send_transaction(&[instruction], &payer_1.pubkey(), &[&payer_1])
.await;
assert_rpc_error(
result, 0,
21, //SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
result,
0,
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
)
.unwrap();
}
Expand Down Expand Up @@ -2270,8 +2271,9 @@ async fn test_approve_failing() {
)
.await;
assert_rpc_error(
result, 0,
21, // SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
result,
0,
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
)
.unwrap();
}
Expand Down Expand Up @@ -2315,8 +2317,9 @@ async fn test_approve_failing() {
)
.await;
assert_rpc_error(
result, 0,
21, //SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
result,
0,
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
)
.unwrap();
}
Expand Down Expand Up @@ -2744,8 +2747,9 @@ async fn test_revoke_failing() {
)
.await;
assert_rpc_error(
result, 0,
21, // SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
result,
0,
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
)
.unwrap();
}
Expand Down
2 changes: 0 additions & 2 deletions programs/compressed-token/anchor/src/burn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,6 @@ pub fn create_input_and_output_accounts_burn(
lamports,
&hashed_mint,
&[inputs.change_account_merkle_tree_index],
remaining_accounts,
)?;
output_compressed_accounts
} else {
Expand All @@ -181,7 +180,6 @@ pub fn create_input_and_output_accounts_burn(
&mut compressed_input_accounts,
input_token_data.as_slice(),
&hashed_mint,
remaining_accounts,
)?;
Ok((compressed_input_accounts, output_compressed_accounts))
}
Expand Down
4 changes: 0 additions & 4 deletions programs/compressed-token/anchor/src/delegation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,13 +158,11 @@ pub fn create_input_and_output_accounts_approve(
lamports,
&hashed_mint,
&merkle_tree_indices,
remaining_accounts,
)?;
add_data_hash_to_input_compressed_accounts::<NOT_FROZEN>(
&mut compressed_input_accounts,
input_token_data.as_slice(),
&hashed_mint,
remaining_accounts,
)?;
Ok((compressed_input_accounts, output_compressed_accounts))
}
Expand Down Expand Up @@ -250,13 +248,11 @@ pub fn create_input_and_output_accounts_revoke(
lamports,
&hashed_mint,
&[inputs.output_account_merkle_tree_index],
remaining_accounts,
)?;
add_data_hash_to_input_compressed_accounts::<NOT_FROZEN>(
&mut compressed_input_accounts,
input_token_data.as_slice(),
&hashed_mint,
remaining_accounts,
)?;
Ok((compressed_input_accounts, output_compressed_accounts))
}
Expand Down
24 changes: 4 additions & 20 deletions programs/compressed-token/anchor/src/freeze.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use account_compression::StateMerkleTreeAccount;
use anchor_lang::prelude::*;
use light_compressed_account::{
compressed_account::{CompressedAccount, CompressedAccountData},
Expand All @@ -11,10 +10,11 @@ use light_compressed_account::{
use light_ctoken_interface::state::CompressedTokenAccountState;

use crate::{
constants::TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR,
process_transfer::{
add_data_hash_to_input_compressed_accounts, cpi_execute_compressed_transaction_transfer,
get_input_compressed_accounts_with_merkle_context_and_check_signer,
get_token_account_discriminator, InputTokenDataWithContext, BATCHED_DISCRIMINATOR,
InputTokenDataWithContext,
},
FreezeInstruction, TokenData,
};
Expand Down Expand Up @@ -115,7 +115,6 @@ pub fn create_input_and_output_accounts_freeze_or_thaw<
&mut compressed_input_accounts,
input_token_data.as_slice(),
&hashed_mint,
remaining_accounts,
)?;
Ok((compressed_input_accounts, output_compressed_accounts))
}
Expand Down Expand Up @@ -164,25 +163,10 @@ fn create_token_output_accounts<const IS_FROZEN: bool>(
};
token_data.serialize(&mut token_data_bytes)?;

let discriminator_bytes = &remaining_accounts[token_data_with_context
.merkle_context
.merkle_tree_pubkey_index
as usize]
.try_borrow_data()?[0..8];
use anchor_lang::Discriminator;
let data_hash = match discriminator_bytes {
StateMerkleTreeAccount::DISCRIMINATOR => token_data.hash_v1(),
BATCHED_DISCRIMINATOR => token_data.hash_v2(),
_ => panic!(), // TODO: throw error
}
.map_err(ProgramError::from)?;

let discriminator = get_token_account_discriminator(discriminator_bytes)?;

let data: CompressedAccountData = CompressedAccountData {
discriminator,
discriminator: TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR,
data: token_data_bytes,
data_hash,
data_hash: token_data.hash_v1().map_err(ProgramError::from)?,
};
output_compressed_accounts[i] = OutputCompressedAccountWithPackedContext {
compressed_account: CompressedAccount {
Expand Down
1 change: 0 additions & 1 deletion programs/compressed-token/anchor/src/process_mint.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,6 @@ pub fn process_mint_to_or_compress<'info, const IS_MINT_TO: bool>(
// We ensure that the Merkle tree account is the first
// remaining account in the cpi to the system program.
&vec![0; amounts.len()],
&[ctx.accounts.merkle_tree.to_account_info()],
)?;
bench_sbf_end!("tm_output_compressed_accounts");

Expand Down
104 changes: 9 additions & 95 deletions programs/compressed-token/anchor/src/process_transfer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
use account_compression::{utils::constants::CPI_AUTHORITY_PDA_SEED, StateMerkleTreeAccount};
use anchor_lang::{
prelude::*, solana_program::program_error::ProgramError, AnchorDeserialize, Discriminator,
};
use account_compression::utils::constants::CPI_AUTHORITY_PDA_SEED;
use anchor_lang::{prelude::*, solana_program::program_error::ProgramError, AnchorDeserialize};
use light_compressed_account::{
compressed_account::{CompressedAccount, CompressedAccountData, PackedMerkleContext},
hash_to_bn254_field_size_be,
Expand All @@ -15,17 +13,11 @@ use light_compressed_account::{
};
use light_ctoken_interface::state::{CompressedTokenAccountState, TokenData};
use light_heap::{bench_sbf_end, bench_sbf_start};
use light_system_program::{
account_traits::{InvokeAccounts, SignerAccounts},
errors::SystemProgramError,
};
use light_system_program::account_traits::{InvokeAccounts, SignerAccounts};
use light_zero_copy::num_trait::ZeroCopyNumTrait;

use crate::{
constants::{
BUMP_CPI_AUTHORITY, NOT_FROZEN, TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR,
TOKEN_COMPRESSED_ACCOUNT_V2_DISCRIMINATOR,
},
constants::{BUMP_CPI_AUTHORITY, NOT_FROZEN, TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR},
spl_compression::process_compression_or_decompression,
ErrorCode, TransferInstruction,
};
Expand Down Expand Up @@ -134,7 +126,6 @@ pub fn process_transfer<'a, 'b, 'c, 'info: 'b + 'c>(
.iter()
.map(|data| data.merkle_tree_index)
.collect::<Vec<u8>>(),
ctx.remaining_accounts,
)?;
bench_sbf_end!("t_create_output_compressed_accounts");

Expand All @@ -144,7 +135,6 @@ pub fn process_transfer<'a, 'b, 'c, 'info: 'b + 'c>(
&mut compressed_input_accounts,
input_token_data.as_slice(),
&hashed_mint,
ctx.remaining_accounts,
)?;
}
bench_sbf_end!("t_add_token_data_to_input_compressed_accounts");
Expand Down Expand Up @@ -183,19 +173,6 @@ pub fn process_transfer<'a, 'b, 'c, 'info: 'b + 'c>(
ctx.remaining_accounts,
)
}
pub const BATCHED_DISCRIMINATOR: &[u8] = b"BatchMta";
pub const OUTPUT_QUEUE_DISCRIMINATOR: &[u8] = b"queueacc";

/// Helper function to determine the appropriate token account discriminator based on tree type
pub fn get_token_account_discriminator(tree_discriminator: &[u8]) -> Result<[u8; 8]> {
match tree_discriminator {
StateMerkleTreeAccount::DISCRIMINATOR => Ok(TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR),
BATCHED_DISCRIMINATOR | OUTPUT_QUEUE_DISCRIMINATOR => {
Ok(TOKEN_COMPRESSED_ACCOUNT_V2_DISCRIMINATOR)
}
_ => err!(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch),
}
}

/// Creates output compressed accounts.
/// Steps:
Expand All @@ -214,7 +191,6 @@ pub fn create_output_compressed_accounts(
lamports: Option<Vec<Option<impl ZeroCopyNumTrait>>>,
hashed_mint: &[u8; 32],
merkle_tree_indices: &[u8],
remaining_accounts: &[AccountInfo<'_>],
) -> Result<u64> {
let mut sum_lamports = 0;
let hashed_delegate_store = if let Some(delegate) = delegate {
Expand Down Expand Up @@ -259,30 +235,7 @@ pub fn create_output_compressed_accounts(
let hashed_owner = hash_to_bn254_field_size_be(owner.to_pubkey_bytes().as_slice());

let mut amount_bytes = [0u8; 32];
let discriminator_bytes =
&remaining_accounts[merkle_tree_indices[i] as usize].try_borrow_data()?[0..8];
match discriminator_bytes {
StateMerkleTreeAccount::DISCRIMINATOR => {
amount_bytes[24..].copy_from_slice(amount.to_bytes_le().as_slice());
Ok(())
}
BATCHED_DISCRIMINATOR => {
amount_bytes[24..].copy_from_slice(amount.to_bytes_be().as_slice());
Ok(())
}
OUTPUT_QUEUE_DISCRIMINATOR => {
amount_bytes[24..].copy_from_slice(amount.to_bytes_be().as_slice());
Ok(())
}
_ => {
msg!(
"{} is no Merkle tree or output queue account. ",
remaining_accounts[merkle_tree_indices[i] as usize].key()
);
err!(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch)
}
}?;

amount_bytes[24..].copy_from_slice(amount.to_bytes_le().as_slice());
let data_hash = TokenData::hash_with_hashed_values(
hashed_mint,
&hashed_owner,
Expand All @@ -291,10 +244,8 @@ pub fn create_output_compressed_accounts(
)
.map_err(ProgramError::from)?;

let discriminator = get_token_account_discriminator(discriminator_bytes)?;

let data = CompressedAccountData {
discriminator,
discriminator: TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR,
data: token_data_bytes,
data_hash,
};
Expand Down Expand Up @@ -326,7 +277,6 @@ pub fn add_data_hash_to_input_compressed_accounts<const FROZEN_INPUTS: bool>(
input_compressed_accounts_with_merkle_context: &mut [InAccount],
input_token_data: &[TokenData],
hashed_mint: &[u8; 32],
remaining_accounts: &[AccountInfo<'_>],
) -> Result<()> {
for (i, compressed_account_with_context) in input_compressed_accounts_with_merkle_context
.iter_mut()
Expand All @@ -335,38 +285,8 @@ pub fn add_data_hash_to_input_compressed_accounts<const FROZEN_INPUTS: bool>(
let hashed_owner = hash_to_bn254_field_size_be(&input_token_data[i].owner.to_bytes());

let mut amount_bytes = [0u8; 32];
let discriminator_bytes = &remaining_accounts[compressed_account_with_context
.merkle_context
.merkle_tree_pubkey_index
as usize]
.try_borrow_data()?[0..8];
match discriminator_bytes {
StateMerkleTreeAccount::DISCRIMINATOR => {
amount_bytes[24..]
.copy_from_slice(input_token_data[i].amount.to_le_bytes().as_slice());
Ok(())
}
BATCHED_DISCRIMINATOR => {
amount_bytes[24..]
.copy_from_slice(input_token_data[i].amount.to_be_bytes().as_slice());
Ok(())
}
OUTPUT_QUEUE_DISCRIMINATOR => {
amount_bytes[24..]
.copy_from_slice(input_token_data[i].amount.to_be_bytes().as_slice());
Ok(())
}
_ => {
msg!(
"{} is no Merkle tree or output queue account. ",
remaining_accounts[compressed_account_with_context
.merkle_context
.merkle_tree_pubkey_index as usize]
.key()
);
err!(anchor_lang::error::ErrorCode::AccountDiscriminatorMismatch)
}
}?;
amount_bytes[24..].copy_from_slice(&input_token_data[i].amount.to_le_bytes());

let delegate_store;
let hashed_delegate = if let Some(delegate) = input_token_data[i].delegate {
delegate_store = hash_to_bn254_field_size_be(&delegate.to_bytes());
Expand Down Expand Up @@ -680,15 +600,9 @@ pub fn get_input_compressed_accounts_with_merkle_context_and_check_signer<const
return err!(ErrorCode::DelegateSignerCheckFailed);
}

// Determine discriminator based on tree type
let discriminator_bytes = &remaining_accounts
[input_token_data.merkle_context.merkle_tree_pubkey_index as usize]
.try_borrow_data()?[0..8];
let discriminator = get_token_account_discriminator(discriminator_bytes)?;

let compressed_account = InAccount {
lamports: input_token_data.lamports.unwrap_or_default(),
discriminator,
discriminator: TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR,
merkle_context: input_token_data.merkle_context,
root_index: input_token_data.root_index,
data_hash: [0u8; 32],
Expand Down
Loading