Skip to content

Commit 1ee21c0

Browse files
ananas-blockananas-block
authored
chore: limit v1 state tree, v2 state&address tree creations to protocol authority (#2325)
* chore: limit v1 state tree, v2 state&address tree creations to protocol authority Entire-Checkpoint: beb8e7abdeef * chore: use bigger stack in test Entire-Checkpoint: ed8d74024c1f
1 parent 93a04cb commit 1ee21c0

16 files changed

Lines changed: 153 additions & 31 deletions

File tree

‎program-tests/account-compression-test/tests/batched_merkle_tree_test.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,7 @@ async fn test_batch_state_merkle_tree() {
379379
let new_keypair_queue = Keypair::new();
380380
let payer = context.get_payer().insecure_clone();
381381
create_batched_state_merkle_tree(
382+
&payer,
382383
&payer,
383384
false,
384385
&mut context,

‎program-tests/justfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ test-account-compression:
1616
RUSTFLAGS="-D warnings" cargo test-sbf -p account-compression-test
1717

1818
test-registry:
19-
RUSTFLAGS="-D warnings" cargo test-sbf -p registry-test
19+
RUST_MIN_STACK=16777216 RUSTFLAGS="-D warnings" cargo test-sbf -p registry-test
2020

2121
# System program tests
2222
test-system: test-system-address test-system-compression test-system-re-init

‎program-tests/registry-test/tests/tests.rs‎

Lines changed: 71 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use light_program_test::{
3131
deregister_program_with_registry_program, register_program_with_registry_program,
3232
},
3333
state_tree::create_state_merkle_tree_and_queue_account,
34+
state_tree_v2::create_batched_state_merkle_tree,
3435
test_accounts::{TestAccounts, NOOP_PROGRAM_ID},
3536
test_keypairs::{GROUP_PDA_SEED_TEST_KEYPAIR, OLD_REGISTRY_ID_TEST_KEYPAIR},
3637
},
@@ -387,6 +388,7 @@ async fn test_initialize_protocol_config() {
387388
let nullifier_queue_keypair = Keypair::new();
388389
let cpi_context_keypair = Keypair::new();
389390
create_state_merkle_tree_and_queue_account(
391+
&payer,
390392
&payer,
391393
true,
392394
&mut rpc,
@@ -411,6 +413,7 @@ async fn test_initialize_protocol_config() {
411413
let nullifier_queue_keypair = Keypair::new();
412414
let cpi_context_keypair = Keypair::new();
413415
let result = create_state_merkle_tree_and_queue_account(
416+
&payer,
414417
&payer,
415418
true,
416419
&mut rpc,
@@ -436,6 +439,7 @@ async fn test_initialize_protocol_config() {
436439
let cpi_context_keypair = Keypair::new();
437440
create_state_merkle_tree_and_queue_account(
438441
&payer,
442+
&updated_keypair,
439443
true,
440444
&mut rpc,
441445
&merkle_tree_keypair,
@@ -453,12 +457,38 @@ async fn test_initialize_protocol_config() {
453457
.await
454458
.unwrap();
455459
}
460+
// FAIL: initialize a Merkle tree with network_fee but wrong authority
461+
{
462+
let merkle_tree_keypair = Keypair::new();
463+
let nullifier_queue_keypair = Keypair::new();
464+
let cpi_context_keypair = Keypair::new();
465+
let result = create_state_merkle_tree_and_queue_account(
466+
&payer,
467+
&payer,
468+
true,
469+
&mut rpc,
470+
&merkle_tree_keypair,
471+
&nullifier_queue_keypair,
472+
Some(&cpi_context_keypair),
473+
None,
474+
None,
475+
1,
476+
&StateMerkleTreeConfig {
477+
network_fee: Some(5000),
478+
..Default::default()
479+
},
480+
&NullifierQueueConfig::default(),
481+
)
482+
.await;
483+
assert_rpc_error(result, 3, RegistryError::InvalidSigner.into()).unwrap();
484+
}
456485
// FAIL: initialize a Merkle tree with network_fee + forester (must be rejected)
457486
{
458487
let merkle_tree_keypair = Keypair::new();
459488
let nullifier_queue_keypair = Keypair::new();
460489
let cpi_context_keypair = Keypair::new();
461490
let result = create_state_merkle_tree_and_queue_account(
491+
&payer,
462492
&payer,
463493
true,
464494
&mut rpc,
@@ -483,6 +513,7 @@ async fn test_initialize_protocol_config() {
483513
let nullifier_queue_keypair = Keypair::new();
484514
let cpi_context_keypair = Keypair::new();
485515
let result = create_state_merkle_tree_and_queue_account(
516+
&payer,
486517
&payer,
487518
true,
488519
&mut rpc,
@@ -1875,11 +1906,48 @@ async fn test_batch_address_tree() {
18751906
network_fee: Some(1),
18761907
..Default::default()
18771908
};
1878-
let result =
1879-
create_batch_address_merkle_tree(&mut rpc, &payer, &new_merkle_tree, test_tree_params)
1880-
.await;
1909+
let result = create_batch_address_merkle_tree(
1910+
&mut rpc,
1911+
&payer,
1912+
&payer,
1913+
&new_merkle_tree,
1914+
test_tree_params,
1915+
)
1916+
.await;
18811917
assert_rpc_error(result, 1, RegistryError::InvalidNetworkFee.into()).unwrap();
18821918
}
1919+
// FAIL: create batched address tree with correct fee but wrong authority
1920+
{
1921+
let new_merkle_tree = Keypair::new();
1922+
let test_tree_params = InitAddressTreeAccountsInstructionData::default();
1923+
let result = create_batch_address_merkle_tree(
1924+
&mut rpc,
1925+
&payer,
1926+
&payer,
1927+
&new_merkle_tree,
1928+
test_tree_params,
1929+
)
1930+
.await;
1931+
assert_rpc_error(result, 1, RegistryError::InvalidSigner.into()).unwrap();
1932+
}
1933+
// FAIL: create batched state tree with correct fee but wrong authority
1934+
{
1935+
let merkle_tree_keypair = Keypair::new();
1936+
let queue_keypair = Keypair::new();
1937+
let cpi_context_keypair = Keypair::new();
1938+
let result = create_batched_state_merkle_tree(
1939+
&payer,
1940+
&payer,
1941+
true,
1942+
&mut rpc,
1943+
&merkle_tree_keypair,
1944+
&queue_keypair,
1945+
&cpi_context_keypair,
1946+
InitStateTreeAccountsInstructionData::default(),
1947+
)
1948+
.await;
1949+
assert_rpc_error(result, 3, RegistryError::InvalidSigner.into()).unwrap();
1950+
}
18831951

18841952
for i in 0..tree_params.input_queue_batch_size * 2 {
18851953
println!("tx {}", i);

‎program-tests/system-test/tests/test.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ use light_compressed_account::{
2020
},
2121
TreeType,
2222
};
23-
use light_merkle_tree_metadata::errors::MerkleTreeMetadataError;
2423
use light_program_test::{
2524
accounts::test_accounts::TestAccounts,
2625
indexer::{TestIndexer, TestIndexerExtensions},
@@ -623,7 +622,9 @@ pub async fn failing_transaction_address<R: Rpc>(
623622
.unwrap();
624623
}
625624

626-
// invalid address queue account
625+
// invalid address queue account (replaced with nullifier queue)
626+
// The system program hits MissingLegacyMerkleContext (error code 6066)
627+
// when trying to look up the legacy context for the replaced queue account.
627628
{
628629
let inputs_struct = inputs_struct.clone();
629630
let mut remaining_accounts = remaining_accounts.clone();
@@ -639,7 +640,7 @@ pub async fn failing_transaction_address<R: Rpc>(
639640
payer,
640641
inputs_struct,
641642
remaining_accounts.clone(),
642-
MerkleTreeMetadataError::InvalidQueueType.into(),
643+
6066, // SystemProgramError::MissingLegacyMerkleContext
643644
)
644645
.await
645646
.unwrap();

‎program-tests/utils/src/e2e_test_env.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1245,6 +1245,7 @@ where
12451245
let forester = Pubkey::new_unique();
12461246
println!("queue config: {:?}", queue_config);
12471247
create_state_merkle_tree_and_queue_account(
1248+
&self.payer,
12481249
&self.payer,
12491250
true,
12501251
&mut self.rpc,

‎programs/registry/src/account_compression_cpi/sdk.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ pub fn create_initialize_address_merkle_tree_and_queue_instruction(
289289
}
290290

291291
pub fn create_initialize_merkle_tree_instruction(
292-
payer: Pubkey,
292+
authority: Pubkey,
293293
merkle_tree_pubkey: Pubkey,
294294
nullifier_queue_pubkey: Pubkey,
295295
cpi_context_pubkey: Pubkey,
@@ -309,7 +309,7 @@ pub fn create_initialize_merkle_tree_instruction(
309309
queue_config: nullifier_queue_config,
310310
};
311311
let accounts = crate::accounts::InitializeMerkleTreeAndQueue {
312-
authority: payer,
312+
authority,
313313
registered_program_pda: register_program_pda,
314314
merkle_tree: merkle_tree_pubkey,
315315
queue: nullifier_queue_pubkey,
@@ -327,7 +327,7 @@ pub fn create_initialize_merkle_tree_instruction(
327327
}
328328

329329
pub fn create_initialize_batched_merkle_tree_instruction(
330-
payer: Pubkey,
330+
authority: Pubkey,
331331
merkle_tree_pubkey: Pubkey,
332332
queue_pubkey: Pubkey,
333333
cpi_context_pubkey: Pubkey,
@@ -341,7 +341,7 @@ pub fn create_initialize_batched_merkle_tree_instruction(
341341
params: params.try_to_vec().unwrap(),
342342
};
343343
let accounts = crate::accounts::InitializeBatchedStateMerkleTreeAndQueue {
344-
authority: payer,
344+
authority,
345345
registered_program_pda: register_program_pda,
346346
merkle_tree: merkle_tree_pubkey,
347347
queue: queue_pubkey,
@@ -460,7 +460,7 @@ pub fn create_rollover_batch_state_tree_instruction(
460460
}
461461

462462
pub fn create_initialize_batched_address_merkle_tree_instruction(
463-
payer: Pubkey,
463+
authority: Pubkey,
464464
merkle_tree_pubkey: Pubkey,
465465
params: InitAddressTreeAccountsInstructionData,
466466
) -> Instruction {
@@ -473,7 +473,7 @@ pub fn create_initialize_batched_address_merkle_tree_instruction(
473473
};
474474
let protocol_config_pda = get_protocol_config_pda_address().0;
475475
let accounts = crate::accounts::InitializeBatchedAddressTree {
476-
authority: payer,
476+
authority,
477477
registered_program_pda: register_program_pda,
478478
merkle_tree: merkle_tree_pubkey,
479479
cpi_authority,

‎programs/registry/src/lib.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,9 @@ pub mod light_registry {
345345
msg!("Forester pubkey must not be defined for trees serviced by light foresters.");
346346
return err!(RegistryError::ForesterDefined);
347347
}
348+
if ctx.accounts.authority.key() != ctx.accounts.protocol_config_pda.authority {
349+
return err!(RegistryError::InvalidSigner);
350+
}
348351
} else if forester.is_none() {
349352
msg!("Forester pubkey required for trees without a network fee.");
350353
msg!("Trees without a network fee will not be serviced by light foresters.");
@@ -518,6 +521,9 @@ pub mod light_registry {
518521
} else {
519522
return err!(RegistryError::InvalidNetworkFee);
520523
}
524+
if ctx.accounts.authority.key() != ctx.accounts.protocol_config_pda.authority {
525+
return err!(RegistryError::InvalidSigner);
526+
}
521527
check_cpi_context(
522528
ctx.accounts.cpi_context_account.to_account_info(),
523529
&ctx.accounts.protocol_config_pda.config,
@@ -606,6 +612,9 @@ pub mod light_registry {
606612
} else {
607613
return err!(RegistryError::InvalidNetworkFee);
608614
}
615+
if ctx.accounts.authority.key() != ctx.accounts.protocol_config_pda.authority {
616+
return err!(RegistryError::InvalidSigner);
617+
}
609618
process_initialize_batched_address_merkle_tree(&ctx, bump, params.try_to_vec()?)
610619
}
611620

‎sdk-libs/program-test/src/accounts/address_tree_v2.rs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use crate::utils::create_account::create_account_instruction;
1111
pub async fn create_batch_address_merkle_tree<R: Rpc>(
1212
rpc: &mut R,
1313
payer: &Keypair,
14+
authority: &Keypair,
1415
new_address_merkle_tree_keypair: &Keypair,
1516
address_tree_params: InitAddressTreeAccountsInstructionData,
1617
) -> Result<Signature, RpcError> {
@@ -34,14 +35,18 @@ pub async fn create_batch_address_merkle_tree<R: Rpc>(
3435
);
3536

3637
let instruction = create_initialize_batched_address_merkle_tree_instruction(
37-
payer.pubkey(),
38+
authority.pubkey(),
3839
new_address_merkle_tree_keypair.pubkey(),
3940
address_tree_params,
4041
);
42+
let mut signers: Vec<&Keypair> = vec![payer, new_address_merkle_tree_keypair];
43+
if authority.pubkey() != payer.pubkey() {
44+
signers.push(authority);
45+
}
4146
rpc.create_and_send_transaction(
4247
&[create_mt_account_ix, instruction],
4348
&payer.pubkey(),
44-
&[payer, new_address_merkle_tree_keypair],
49+
&signers,
4550
)
4651
.await
4752
}

‎sdk-libs/program-test/src/accounts/initialize.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ pub async fn initialize_accounts<R: Rpc + TestRpc>(
117117
}
118118
if !config.skip_v1_trees {
119119
create_state_merkle_tree_and_queue_account(
120+
&keypairs.governance_authority,
120121
&keypairs.governance_authority,
121122
true,
122123
context,
@@ -133,6 +134,7 @@ pub async fn initialize_accounts<R: Rpc + TestRpc>(
133134

134135
if !skip_second_v1_tree {
135136
create_state_merkle_tree_and_queue_account(
137+
&keypairs.governance_authority,
136138
&keypairs.governance_authority,
137139
true,
138140
context,
@@ -168,6 +170,7 @@ pub async fn initialize_accounts<R: Rpc + TestRpc>(
168170
#[cfg(feature = "v2")]
169171
if let Some(v2_state_tree_config) = _v2_state_tree_config {
170172
create_batched_state_merkle_tree(
173+
&keypairs.governance_authority,
171174
&keypairs.governance_authority,
172175
true,
173176
context,
@@ -183,6 +186,7 @@ pub async fn initialize_accounts<R: Rpc + TestRpc>(
183186
create_batch_address_merkle_tree(
184187
context,
185188
&keypairs.governance_authority,
189+
&keypairs.governance_authority,
186190
&keypairs.batch_address_merkle_tree,
187191
*params,
188192
)

‎sdk-libs/program-test/src/accounts/state_tree.rs‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ pub fn create_insert_leaves_instruction(
105105
#[allow(clippy::too_many_arguments)]
106106
pub async fn create_state_merkle_tree_and_queue_account<R: Rpc>(
107107
payer: &Keypair,
108+
authority: &Keypair,
108109
registry: bool,
109110
rpc: &mut R,
110111
merkle_tree_keypair: &Keypair,
@@ -163,7 +164,7 @@ pub async fn create_state_merkle_tree_and_queue_account<R: Rpc>(
163164
);
164165

165166
let instruction = create_initialize_merkle_tree_instruction_registry(
166-
payer.pubkey(),
167+
authority.pubkey(),
167168
merkle_tree_keypair.pubkey(),
168169
nullifier_queue_keypair.pubkey(),
169170
cpi_context_keypair.pubkey(),
@@ -172,6 +173,15 @@ pub async fn create_state_merkle_tree_and_queue_account<R: Rpc>(
172173
program_owner,
173174
forester,
174175
);
176+
let mut signers = vec![
177+
payer,
178+
merkle_tree_keypair,
179+
nullifier_queue_keypair,
180+
cpi_context_keypair,
181+
];
182+
if authority.pubkey() != payer.pubkey() {
183+
signers.push(authority);
184+
}
175185
Transaction::new_signed_with_payer(
176186
&[
177187
create_cpi_context_instruction,
@@ -180,12 +190,7 @@ pub async fn create_state_merkle_tree_and_queue_account<R: Rpc>(
180190
instruction,
181191
],
182192
Some(&payer.pubkey()),
183-
&vec![
184-
payer,
185-
merkle_tree_keypair,
186-
nullifier_queue_keypair,
187-
cpi_context_keypair,
188-
],
193+
&signers,
189194
rpc.get_latest_blockhash().await?.0,
190195
)
191196
} else {

0 commit comments

Comments
 (0)