diff --git a/account-comparison/programs/account-comparison/src/lib.rs b/account-comparison/programs/account-comparison/src/lib.rs index 0877a5f..ebfe064 100644 --- a/account-comparison/programs/account-comparison/src/lib.rs +++ b/account-comparison/programs/account-comparison/src/lib.rs @@ -4,6 +4,7 @@ use anchor_lang::prelude::*; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -58,11 +59,18 @@ pub mod account_comparison { CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|err| ProgramError::from(LightSdkError::from(err)))?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"account", ctx.accounts.user.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|err| ProgramError::from(LightSdkError::from(err)))?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/anchor/burn/programs/burn/src/lib.rs b/basic-operations/anchor/burn/programs/burn/src/lib.rs index ea6dd1e..b4db4da 100644 --- a/basic-operations/anchor/burn/programs/burn/src/lib.rs +++ b/basic-operations/anchor/burn/programs/burn/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMetaBurn, PackedAddressTreeInfo, ValidityProof}, @@ -38,11 +39,18 @@ pub mod burn { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"message", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/anchor/close/programs/close/src/lib.rs b/basic-operations/anchor/close/programs/close/src/lib.rs index 83a37c4..070edc8 100644 --- a/basic-operations/anchor/close/programs/close/src/lib.rs +++ b/basic-operations/anchor/close/programs/close/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -38,11 +39,18 @@ pub mod close { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"message", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/anchor/create/programs/create/src/lib.rs b/basic-operations/anchor/create/programs/create/src/lib.rs index 7880d0c..dc08fcf 100644 --- a/basic-operations/anchor/create/programs/create/src/lib.rs +++ b/basic-operations/anchor/create/programs/create/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{PackedAddressTreeInfo, ValidityProof}, @@ -38,11 +39,18 @@ pub mod create { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"message", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/anchor/reinit/programs/reinit/src/lib.rs b/basic-operations/anchor/reinit/programs/reinit/src/lib.rs index 72b1787..90b314a 100644 --- a/basic-operations/anchor/reinit/programs/reinit/src/lib.rs +++ b/basic-operations/anchor/reinit/programs/reinit/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -38,11 +39,18 @@ pub mod reinit { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"message", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/anchor/update/programs/update/src/lib.rs b/basic-operations/anchor/update/programs/update/src/lib.rs index b9e619e..51a3867 100644 --- a/basic-operations/anchor/update/programs/update/src/lib.rs +++ b/basic-operations/anchor/update/programs/update/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -38,11 +39,18 @@ pub mod update { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"message", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/basic-operations/native/programs/burn/src/lib.rs b/basic-operations/native/programs/burn/src/lib.rs index b10d99f..24bb3fc 100644 --- a/basic-operations/native/programs/burn/src/lib.rs +++ b/basic-operations/native/programs/burn/src/lib.rs @@ -8,6 +8,7 @@ use light_macros::pubkey; use light_sdk::{ account::sha::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -81,12 +82,19 @@ fn create(accounts: &[AccountInfo], instruction_data: &[u8]) -> Result<(), Light let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instruction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(LightSdkError::ProgramError(ProgramError::InvalidAccountData)); + } + let (address, address_seed) = derive_address( &[b"message", signer.key.as_ref()], - &instruction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/basic-operations/native/programs/close/src/lib.rs b/basic-operations/native/programs/close/src/lib.rs index cee8620..ba9af6a 100644 --- a/basic-operations/native/programs/close/src/lib.rs +++ b/basic-operations/native/programs/close/src/lib.rs @@ -8,6 +8,7 @@ use light_macros::pubkey; use light_sdk::{ account::sha::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -81,12 +82,19 @@ fn create(accounts: &[AccountInfo], instruction_data: &[u8]) -> Result<(), Light let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instruction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(LightSdkError::ProgramError(ProgramError::InvalidAccountData)); + } + let (address, address_seed) = derive_address( &[b"message", signer.key.as_ref()], - &instruction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/basic-operations/native/programs/create/src/lib.rs b/basic-operations/native/programs/create/src/lib.rs index e07a9c9..f44e6ff 100644 --- a/basic-operations/native/programs/create/src/lib.rs +++ b/basic-operations/native/programs/create/src/lib.rs @@ -8,6 +8,7 @@ use light_macros::pubkey; use light_sdk::{ account::sha::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -93,12 +94,19 @@ pub fn create( let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instruction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData); + } + let (address, address_seed) = derive_address( &[b"message", signer.key.as_ref()], - &instruction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/basic-operations/native/programs/reinit/src/lib.rs b/basic-operations/native/programs/reinit/src/lib.rs index 49cbe8c..da3cac0 100644 --- a/basic-operations/native/programs/reinit/src/lib.rs +++ b/basic-operations/native/programs/reinit/src/lib.rs @@ -8,6 +8,7 @@ use light_macros::pubkey; use light_sdk::{ account::sha::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -89,12 +90,19 @@ fn create(accounts: &[AccountInfo], instruction_data: &[u8]) -> Result<(), Light let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instruction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(LightSdkError::ProgramError(ProgramError::InvalidAccountData)); + } + let (address, address_seed) = derive_address( &[b"message", signer.key.as_ref()], - &instruction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/basic-operations/native/programs/update/src/lib.rs b/basic-operations/native/programs/update/src/lib.rs index 3bdc147..4e51406 100644 --- a/basic-operations/native/programs/update/src/lib.rs +++ b/basic-operations/native/programs/update/src/lib.rs @@ -8,6 +8,7 @@ use light_macros::pubkey; use light_sdk::{ account::sha::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -109,12 +110,19 @@ pub fn create( let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instruction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData); + } + let (address, address_seed) = derive_address( &[b"message", signer.key.as_ref()], - &instruction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/counter/anchor/programs/counter/src/lib.rs b/counter/anchor/programs/counter/src/lib.rs index 4712f94..98074e1 100644 --- a/counter/anchor/programs/counter/src/lib.rs +++ b/counter/anchor/programs/counter/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, Discriminator}; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{v1::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -41,11 +42,18 @@ pub mod counter { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[b"counter", ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/counter/native/src/lib.rs b/counter/native/src/lib.rs index 7c4d51a..231783d 100644 --- a/counter/native/src/lib.rs +++ b/counter/native/src/lib.rs @@ -5,6 +5,7 @@ use light_macros::pubkey; use light_sdk::{ account::LightAccount, address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -165,12 +166,19 @@ pub fn create_counter( let light_cpi_accounts = CpiAccounts::new(signer, &accounts[1..], LIGHT_CPI_SIGNER); + let address_tree_pubkey = instuction_data + .address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ProgramError::NotEnoughAccountKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V1 { + solana_program::msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData); + } + let (address, address_seed) = derive_address( &[b"counter", signer.key.as_ref()], - &instuction_data - .address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ProgramError::NotEnoughAccountKeys)?, + &address_tree_pubkey, &ID, ); diff --git a/counter/pinocchio/src/lib.rs b/counter/pinocchio/src/lib.rs index e2a3cbd..367cd98 100644 --- a/counter/pinocchio/src/lib.rs +++ b/counter/pinocchio/src/lib.rs @@ -4,6 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use light_macros::pubkey_array; use light_sdk_pinocchio::{ address::v1::derive_address, + constants::ADDRESS_TREE_V1, cpi::{ v1::{CpiAccounts, LightSystemProgramCpi}, CpiAccountsConfig, CpiSigner, InvokeLightSystemProgram, LightCpiInstruction, @@ -24,11 +25,11 @@ pub const LIGHT_CPI_SIGNER: CpiSigner = entrypoint!(process_instruction); fn to_custom_error>(e: E) -> ProgramError { - ProgramError::Custom(u64::from(e.into()) as u32) + ProgramError::Custom(e.into() as u32) } fn to_custom_error_u32>(e: E) -> ProgramError { - ProgramError::Custom(u32::from(e.into())) + ProgramError::Custom(e.into()) } #[repr(u8)] @@ -186,10 +187,15 @@ pub fn create_counter( .map_err(to_custom_error_u32)? .key(); + if *tree_pubkey != ADDRESS_TREE_V1 { + pinocchio::log::sol_log("Invalid address tree"); + return Err(ProgramError::InvalidAccountData); + } + let program_id = Pubkey::from(ID); let (address, address_seed) = derive_address( &[b"counter", signer.key().as_ref()], - &tree_pubkey, + tree_pubkey, &program_id, ); diff --git a/create-and-update/programs/create-and-update/src/lib.rs b/create-and-update/programs/create-and-update/src/lib.rs index bb71c72..0082295 100644 --- a/create-and-update/programs/create-and-update/src/lib.rs +++ b/create-and-update/programs/create-and-update/src/lib.rs @@ -5,6 +5,7 @@ use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize}; use light_sdk::{ account::LightAccount, address::v2::derive_address, + constants::ADDRESS_TREE_V2, cpi::{v2::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{account_meta::CompressedAccountMeta, PackedAddressTreeInfo, ValidityProof}, @@ -44,6 +45,12 @@ pub mod create_and_update { let address_tree_pubkey = address_tree_info .get_tree_pubkey(&light_cpi_accounts) .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[FIRST_SEED, ctx.accounts.signer.key().as_ref()], &address_tree_pubkey, @@ -87,6 +94,11 @@ pub mod create_and_update { .get_tree_pubkey(&light_cpi_accounts) .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + if new_account_address_tree_pubkey.to_bytes() != ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + // Create new compressed account let (new_address, new_address_seed) = derive_address( &[SECOND_SEED, ctx.accounts.signer.key().as_ref()], @@ -199,12 +211,19 @@ pub mod create_and_update { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + // Create first compressed account let (first_address, first_address_seed) = derive_address( &[FIRST_SEED, ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); @@ -219,9 +238,7 @@ pub mod create_and_update { // Create second compressed account let (second_address, second_address_seed) = derive_address( &[SECOND_SEED, ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/read-only/src/lib.rs b/read-only/src/lib.rs index e5068d6..4a4e915 100644 --- a/read-only/src/lib.rs +++ b/read-only/src/lib.rs @@ -8,6 +8,7 @@ use light_sdk::cpi::{v2::LightSystemProgramCpi, InvokeLightSystemProgram, LightC use light_sdk::{ account::LightAccount, address::v2::derive_address, + constants::ADDRESS_TREE_V2, cpi::{v2::CpiAccounts, CpiSigner}, derive_light_cpi_signer, instruction::{ @@ -42,11 +43,18 @@ pub mod read_only { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[FIRST_SEED, ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); diff --git a/zk-id/src/lib.rs b/zk-id/src/lib.rs index a9d78bd..4164e0e 100644 --- a/zk-id/src/lib.rs +++ b/zk-id/src/lib.rs @@ -53,11 +53,18 @@ pub mod zk_id { crate::LIGHT_CPI_SIGNER, ); + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != light_sdk::constants::ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[ISSUER, ctx.accounts.signer.key().as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, ); msg!("address {:?}", address); @@ -116,11 +123,18 @@ pub mod zk_id { .checked_add(1) .ok_or(ProgramError::ArithmeticOverflow)?; + let address_tree_pubkey = address_tree_info + .get_tree_pubkey(&light_cpi_accounts) + .map_err(|_| ErrorCode::AccountNotEnoughKeys)?; + + if address_tree_pubkey.to_bytes() != light_sdk::constants::ADDRESS_TREE_V2 { + msg!("Invalid address tree"); + return Err(ProgramError::InvalidAccountData.into()); + } + let (address, address_seed) = derive_address( &[CREDENTIAL, credential_pubkey.as_ref()], - &address_tree_info - .get_tree_pubkey(&light_cpi_accounts) - .map_err(|_| ErrorCode::AccountNotEnoughKeys)?, + &address_tree_pubkey, &crate::ID, );