From a8de2cb3b251db8b1f43bba4dac5d685ee5cbf4d Mon Sep 17 00:00:00 2001 From: ananas Date: Thu, 16 Oct 2025 17:37:17 +0100 Subject: [PATCH] fix: CToken checked deserialization --- .../src/state/ctoken/zero_copy.rs | 41 ++++++++++ .../ctoken-types/tests/ctoken/failing.rs | 79 ++++++++++++++++++- .../tests/ctoken/close.rs | 4 +- .../tests/ctoken/compress_and_close.rs | 4 +- .../compressed-token/program/src/claim.rs | 8 +- .../src/close_token_account/processor.rs | 5 +- .../program/src/ctoken_transfer.rs | 3 +- .../ctoken/compress_or_decompress_ctokens.rs | 4 +- 8 files changed, 132 insertions(+), 16 deletions(-) diff --git a/program-libs/ctoken-types/src/state/ctoken/zero_copy.rs b/program-libs/ctoken-types/src/state/ctoken/zero_copy.rs index 4a62f8abc5..a708c83844 100644 --- a/program-libs/ctoken-types/src/state/ctoken/zero_copy.rs +++ b/program-libs/ctoken-types/src/state/ctoken/zero_copy.rs @@ -433,6 +433,47 @@ impl<'a> ZeroCopyAt<'a> for CToken { } } +impl CToken { + /// Zero-copy deserialization with initialization check. + /// Returns an error if the account is not initialized (byte 108 must be 1). + #[profile] + pub fn zero_copy_at_checked( + bytes: &[u8], + ) -> Result<(ZCToken<'_>, &[u8]), crate::error::CTokenError> { + // Check minimum size for state field at byte 108 + if bytes.len() < 109 { + return Err(crate::error::CTokenError::InvalidAccountData); + } + + // Verify account is initialized (state byte at offset 108 must be 1) + if bytes[108] != 1 { + return Err(crate::error::CTokenError::InvalidAccountState); + } + + // Proceed with normal deserialization + Ok(CToken::zero_copy_at(bytes)?) + } + + /// Mutable zero-copy deserialization with initialization check. + /// Returns an error if the account is not initialized (byte 108 must be 1). + #[profile] + pub fn zero_copy_at_mut_checked( + bytes: &mut [u8], + ) -> Result<(ZCompressedTokenMut<'_>, &mut [u8]), crate::error::CTokenError> { + // Check minimum size for state field at byte 108 + if bytes.len() < 109 { + return Err(crate::error::CTokenError::InvalidAccountData); + } + + // Verify account is initialized (state byte at offset 108 must be 1) + if bytes[108] != 1 { + return Err(crate::error::CTokenError::InvalidAccountState); + } + + Ok(CToken::zero_copy_at_mut(bytes)?) + } +} + impl<'a> ZeroCopyAtMut<'a> for CToken { type ZeroCopyAtMut = ZCompressedTokenMut<'a>; diff --git a/program-libs/ctoken-types/tests/ctoken/failing.rs b/program-libs/ctoken-types/tests/ctoken/failing.rs index 0739596271..60e08e5653 100644 --- a/program-libs/ctoken-types/tests/ctoken/failing.rs +++ b/program-libs/ctoken-types/tests/ctoken/failing.rs @@ -1,4 +1,7 @@ -use light_ctoken_types::state::{CToken, CompressedTokenConfig}; +use light_ctoken_types::{ + error::CTokenError, + state::{CToken, CompressedTokenConfig}, +}; use light_zero_copy::ZeroCopyNew; #[test] @@ -17,3 +20,77 @@ fn test_compressed_token_new_zero_copy_buffer_too_small() { // Should fail with size error assert!(result.is_err()); } + +#[test] +fn test_zero_copy_at_checked_uninitialized_account() { + // Create a 165-byte buffer with all zeros (byte 108 = 0, uninitialized) + let buffer = vec![0u8; 165]; + + // This should fail because byte 108 is 0 (not initialized) + let result = CToken::zero_copy_at_checked(&buffer); + + // Assert it returns InvalidAccountState error + assert!(matches!(result, Err(CTokenError::InvalidAccountState))); +} + +#[test] +fn test_zero_copy_at_mut_checked_uninitialized_account() { + // Create a 165-byte mutable buffer with all zeros + let mut buffer = vec![0u8; 165]; + + // This should fail because byte 108 is 0 (not initialized) + let result = CToken::zero_copy_at_mut_checked(&mut buffer); + + // Assert it returns InvalidAccountState error + assert!(matches!(result, Err(CTokenError::InvalidAccountState))); +} + +#[test] +fn test_zero_copy_at_checked_frozen_account() { + // Create a 165-byte buffer with byte 108 = 2 (AccountState::Frozen) + let mut buffer = vec![0u8; 165]; + buffer[108] = 2; // AccountState::Frozen + + // This should fail because byte 108 is 2 (frozen, not initialized) + let result = CToken::zero_copy_at_checked(&buffer); + + // Assert it returns InvalidAccountState error + assert!(matches!(result, Err(CTokenError::InvalidAccountState))); +} + +#[test] +fn test_zero_copy_at_mut_checked_frozen_account() { + // Create a 165-byte mutable buffer with byte 108 = 2 + let mut buffer = vec![0u8; 165]; + buffer[108] = 2; // AccountState::Frozen + + // This should fail because byte 108 is 2 (frozen, not initialized) + let result = CToken::zero_copy_at_mut_checked(&mut buffer); + + // Assert it returns InvalidAccountState error + assert!(matches!(result, Err(CTokenError::InvalidAccountState))); +} + +#[test] +fn test_zero_copy_at_checked_buffer_too_small() { + // Create a 100-byte buffer (less than 109 bytes minimum) + let buffer = vec![0u8; 100]; + + // This should fail because buffer is too small + let result = CToken::zero_copy_at_checked(&buffer); + + // Assert it returns InvalidAccountData error + assert!(matches!(result, Err(CTokenError::InvalidAccountData))); +} + +#[test] +fn test_zero_copy_at_mut_checked_buffer_too_small() { + // Create a 100-byte mutable buffer + let mut buffer = vec![0u8; 100]; + + // This should fail because buffer is too small + let result = CToken::zero_copy_at_mut_checked(&mut buffer); + + // Assert it returns InvalidAccountData error + assert!(matches!(result, Err(CTokenError::InvalidAccountData))); +} diff --git a/program-tests/compressed-token-test/tests/ctoken/close.rs b/program-tests/compressed-token-test/tests/ctoken/close.rs index b10aad95a4..4d006beafa 100644 --- a/program-tests/compressed-token-test/tests/ctoken/close.rs +++ b/program-tests/compressed-token-test/tests/ctoken/close.rs @@ -264,7 +264,7 @@ async fn test_close_token_account_fails() { &owner_keypair, Some(rent_sponsor), "uninitialized_account", - 10, // ProgramError::UninitializedAccount + 18036, // CTokenError::InvalidAccountState ) .await; } @@ -317,7 +317,7 @@ async fn test_close_token_account_fails() { &owner_keypair, Some(rent_sponsor), "frozen_account", - 76, // ErrorCode::AccountFrozen + 18036, // CTokenError::InvalidAccountState ) .await; } diff --git a/program-tests/compressed-token-test/tests/ctoken/compress_and_close.rs b/program-tests/compressed-token-test/tests/ctoken/compress_and_close.rs index 4b725c091c..c9fd938261 100644 --- a/program-tests/compressed-token-test/tests/ctoken/compress_and_close.rs +++ b/program-tests/compressed-token-test/tests/ctoken/compress_and_close.rs @@ -840,8 +840,8 @@ async fn test_compress_and_close_output_validation_errors() { .await; // Assert that the transaction failed with account frozen error - // Error: AccountFrozen (76 = 0x4c) - light_program_test::utils::assert::assert_rpc_error(result, 0, 76).unwrap(); + // Error: InvalidAccountState (18036) + light_program_test::utils::assert::assert_rpc_error(result, 0, 18036).unwrap(); } } diff --git a/programs/compressed-token/program/src/claim.rs b/programs/compressed-token/program/src/claim.rs index 5224222549..5a6c1d18eb 100644 --- a/programs/compressed-token/program/src/claim.rs +++ b/programs/compressed-token/program/src/claim.rs @@ -1,10 +1,9 @@ use anchor_compressed_token::ErrorCode; use anchor_lang::prelude::ProgramError; -use light_account_checks::{AccountInfoTrait, AccountIterator}; +use light_account_checks::{checks::check_owner, AccountInfoTrait, AccountIterator}; use light_compressible::{compression_info::ClaimAndUpdate, config::CompressibleConfig}; use light_ctoken_types::state::{CToken, ZExtensionStructMut}; use light_program_profiler::profile; -use light_zero_copy::traits::ZeroCopyAtMut; use pinocchio::{account_info::AccountInfo, sysvars::Sysvar}; use spl_pod::solana_msg::msg; @@ -96,13 +95,16 @@ fn validate_and_claim( token_account: &AccountInfo, current_slot: u64, ) -> Result, ProgramError> { + // Verify the token account is owned by the compressed token program + check_owner(&crate::LIGHT_CPI_SIGNER.program_id, token_account)?; + // Get current lamports balance let current_lamports = AccountInfoTrait::lamports(token_account); // Claim rent for completed epochs let bytes = token_account.data_len() as u64; // Parse and process the token account let mut token_account_data = AccountInfoTrait::try_borrow_mut_data(token_account)?; - let (mut compressed_token, _) = CToken::zero_copy_at_mut(&mut token_account_data)?; + let (mut compressed_token, _) = CToken::zero_copy_at_mut_checked(&mut token_account_data)?; // Find compressible extension if let Some(extensions) = compressed_token.extensions.as_mut() { diff --git a/programs/compressed-token/program/src/close_token_account/processor.rs b/programs/compressed-token/program/src/close_token_account/processor.rs index 4e92d6a51e..39e8e26b26 100644 --- a/programs/compressed-token/program/src/close_token_account/processor.rs +++ b/programs/compressed-token/program/src/close_token_account/processor.rs @@ -4,7 +4,6 @@ use light_account_checks::{checks::check_signer, AccountInfoTrait}; use light_compressible::rent::{get_rent_exemption_lamports, AccountRentState}; use light_ctoken_types::state::{CToken, ZCompressedTokenMut, ZExtensionStructMut}; use light_program_profiler::profile; -use light_zero_copy::traits::{ZeroCopyAt, ZeroCopyAtMut}; use pinocchio::account_info::AccountInfo; #[cfg(target_os = "solana")] use pinocchio::sysvars::Sysvar; @@ -26,7 +25,7 @@ pub fn process_close_token_account( // Try to parse as CToken using zero-copy deserialization let token_account_data = &mut AccountInfoTrait::try_borrow_mut_data(accounts.token_account)?; - let (ctoken, _) = CToken::zero_copy_at_mut(token_account_data)?; + let (ctoken, _) = CToken::zero_copy_at_mut_checked(token_account_data)?; validate_token_account_close_instruction(&accounts, &ctoken)?; } close_token_account(&accounts)?; @@ -156,7 +155,7 @@ pub fn distribute_lamports(accounts: &CloseTokenAccountAccounts<'_>) -> Result<( // Check for compressible extension and handle lamport distribution let token_account_data = AccountInfoTrait::try_borrow_data(accounts.token_account)?; - let (ctoken, _) = CToken::zero_copy_at(&token_account_data)?; + let (ctoken, _) = CToken::zero_copy_at_checked(&token_account_data)?; if let Some(extensions) = ctoken.extensions.as_ref() { for extension in extensions { diff --git a/programs/compressed-token/program/src/ctoken_transfer.rs b/programs/compressed-token/program/src/ctoken_transfer.rs index 7687778111..58e1fabcf7 100644 --- a/programs/compressed-token/program/src/ctoken_transfer.rs +++ b/programs/compressed-token/program/src/ctoken_transfer.rs @@ -4,7 +4,6 @@ use light_ctoken_types::{ CTokenError, }; use light_program_profiler::profile; -use light_zero_copy::traits::ZeroCopyAt; use pinocchio::account_info::AccountInfo; use pinocchio_token_program::processor::transfer::process_transfer; @@ -60,7 +59,7 @@ fn calculate_and_execute_top_up_transfers( .account .try_borrow_data() .map_err(convert_program_error)?; - let (token, _) = CToken::zero_copy_at(&account_data)?; + let (token, _) = CToken::zero_copy_at_checked(&account_data)?; if let Some(extensions) = token.extensions.as_ref() { for extension in extensions.iter() { if let ZExtensionStruct::Compressible(compressible_extension) = extension { diff --git a/programs/compressed-token/program/src/transfer2/compression/ctoken/compress_or_decompress_ctokens.rs b/programs/compressed-token/program/src/transfer2/compression/ctoken/compress_or_decompress_ctokens.rs index ddc270371b..fb08b80060 100644 --- a/programs/compressed-token/program/src/transfer2/compression/ctoken/compress_or_decompress_ctokens.rs +++ b/programs/compressed-token/program/src/transfer2/compression/ctoken/compress_or_decompress_ctokens.rs @@ -7,7 +7,6 @@ use light_ctoken_types::{ CTokenError, }; use light_program_profiler::profile; -use light_zero_copy::traits::ZeroCopyAtMut; use pinocchio::{ account_info::AccountInfo, sysvars::{clock::Clock, Sysvar}, @@ -37,8 +36,7 @@ pub fn compress_or_decompress_ctokens( .try_borrow_mut_data() .map_err(|_| ProgramError::AccountBorrowFailed)?; - let (mut ctoken, _) = CToken::zero_copy_at_mut(&mut token_account_data) - .map_err(|_| ProgramError::InvalidAccountData)?; + let (mut ctoken, _) = CToken::zero_copy_at_mut_checked(&mut token_account_data)?; if ctoken.mint.to_bytes() != mint { msg!(