Skip to content

Commit bee12da

Browse files
sergeytimoshinsergeytimoshin
committed
fix: harden runtime safety fallout
1 parent d436767 commit bee12da

18 files changed

Lines changed: 271 additions & 193 deletions

File tree

‎forester/src/forester_status.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -670,8 +670,8 @@ fn parse_tree_status(
670670
let fullness = next_index as f64 / capacity as f64 * 100.0;
671671

672672
let (queue_len, queue_cap) = queue_account
673-
.map(|acc| {
674-
match unsafe { parse_hash_set_from_bytes::<QueueAccount>(&acc.data) } {
673+
.map(
674+
|acc| match unsafe { parse_hash_set_from_bytes::<QueueAccount>(&acc.data) } {
675675
Ok(hs) => {
676676
let len = hs
677677
.iter()
@@ -684,8 +684,8 @@ fn parse_tree_status(
684684
warn!(?error, "Failed to parse StateV1 queue hash set");
685685
(None, None)
686686
}
687-
}
688-
})
687+
},
688+
)
689689
.unwrap_or((None, None));
690690

691691
(
@@ -727,8 +727,8 @@ fn parse_tree_status(
727727
let fullness = next_index as f64 / capacity as f64 * 100.0;
728728

729729
let (queue_len, queue_cap) = queue_account
730-
.map(|acc| {
731-
match unsafe { parse_hash_set_from_bytes::<QueueAccount>(&acc.data) } {
730+
.map(
731+
|acc| match unsafe { parse_hash_set_from_bytes::<QueueAccount>(&acc.data) } {
732732
Ok(hs) => {
733733
let len = hs
734734
.iter()
@@ -741,8 +741,8 @@ fn parse_tree_status(
741741
warn!(?error, "Failed to parse AddressV1 queue hash set");
742742
(None, None)
743743
}
744-
}
745-
})
744+
},
745+
)
746746
.unwrap_or((None, None));
747747

748748
(

‎forester/src/processor/v2/proof_worker.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,9 @@ impl ProofClients {
164164
}
165165
}
166166

167-
pub fn spawn_proof_workers(config: &ProverConfig) -> crate::Result<async_channel::Sender<ProofJob>> {
167+
pub fn spawn_proof_workers(
168+
config: &ProverConfig,
169+
) -> crate::Result<async_channel::Sender<ProofJob>> {
168170
let (job_tx, job_rx) = async_channel::bounded::<ProofJob>(256);
169171
let clients = Arc::new(ProofClients::new(config)?);
170172
tokio::spawn(async move { run_proof_pipeline(job_rx, clients).await });

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -209,9 +209,7 @@ impl<const HEIGHT: usize> MockBatchedForester<HEIGHT> {
209209
&[],
210210
)?;
211211
let proof_client = ProofClient::local()?;
212-
let proof_result = proof_client
213-
.generate_batch_update_proof(inputs)
214-
.await?;
212+
let proof_result = proof_client.generate_batch_update_proof(inputs).await?;
215213
let new_root = self.merkle_tree.root();
216214
let proof = CompressedProof {
217215
a: proof_result.0.proof.a,

‎prover/client/src/proof.rs‎

Lines changed: 20 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ use solana_bn254::compression::prelude::{
1212
convert_endianness,
1313
};
1414

15+
pub type CompressedProofBytes = ([u8; 32], [u8; 64], [u8; 32]);
16+
pub type UncompressedProofBytes = ([u8; 64], [u8; 128], [u8; 64]);
17+
1518
#[derive(Debug, Clone, Copy)]
1619
pub struct ProofCompressed {
1720
pub a: [u8; 32],
@@ -71,9 +74,8 @@ pub fn deserialize_hex_string_to_be_bytes(hex_str: &str) -> Result<[u8; 32], Pro
7174
.strip_prefix("0x")
7275
.or_else(|| hex_str.strip_prefix("0X"))
7376
.unwrap_or(hex_str);
74-
let big_uint = num_bigint::BigUint::from_str_radix(trimmed_str, 16).map_err(|error| {
75-
ProverClientError::InvalidHexString(format!("{hex_str}: {error}"))
76-
})?;
77+
let big_uint = num_bigint::BigUint::from_str_radix(trimmed_str, 16)
78+
.map_err(|error| ProverClientError::InvalidHexString(format!("{hex_str}: {error}")))?;
7779
let big_uint_bytes = big_uint.to_bytes_be();
7880
if big_uint_bytes.len() > 32 {
7981
return Err(ProverClientError::InvalidHexString(format!(
@@ -95,7 +97,7 @@ pub fn compress_proof(
9597
proof_a: &[u8; 64],
9698
proof_b: &[u8; 128],
9799
proof_c: &[u8; 64],
98-
) -> Result<([u8; 32], [u8; 64], [u8; 32]), ProverClientError> {
100+
) -> Result<CompressedProofBytes, ProverClientError> {
99101
let proof_a = alt_bn128_g1_compress(proof_a)?;
100102
let proof_b = alt_bn128_g2_compress(proof_b)?;
101103
let proof_c = alt_bn128_g1_compress(proof_c)?;
@@ -104,7 +106,7 @@ pub fn compress_proof(
104106

105107
pub fn proof_from_json_struct(
106108
json: GnarkProofJson,
107-
) -> Result<([u8; 64], [u8; 128], [u8; 64]), ProverClientError> {
109+
) -> Result<UncompressedProofBytes, ProverClientError> {
108110
let proof_a_x = deserialize_hex_string_to_be_bytes(json.ar.first().ok_or_else(|| {
109111
ProverClientError::InvalidProofData("missing proof A x coordinate".to_string())
110112
})?)?;
@@ -117,37 +119,24 @@ pub fn proof_from_json_struct(
117119
.map_err(|_| ProverClientError::InvalidProofData("invalid proof A length".to_string()))?;
118120
let proof_a = negate_g1(&proof_a)?;
119121
let proof_b_x_0 = deserialize_hex_string_to_be_bytes(
120-
json.bs
121-
.first()
122-
.and_then(|row| row.first())
123-
.ok_or_else(|| {
124-
ProverClientError::InvalidProofData("missing proof B x0 coordinate".to_string())
125-
})?,
122+
json.bs.first().and_then(|row| row.first()).ok_or_else(|| {
123+
ProverClientError::InvalidProofData("missing proof B x0 coordinate".to_string())
124+
})?,
126125
)?;
127126
let proof_b_x_1 = deserialize_hex_string_to_be_bytes(
128-
json.bs
129-
.first()
130-
.and_then(|row| row.get(1))
131-
.ok_or_else(|| {
132-
ProverClientError::InvalidProofData("missing proof B x1 coordinate".to_string())
133-
})?,
127+
json.bs.first().and_then(|row| row.get(1)).ok_or_else(|| {
128+
ProverClientError::InvalidProofData("missing proof B x1 coordinate".to_string())
129+
})?,
134130
)?;
135131
let proof_b_y_0 = deserialize_hex_string_to_be_bytes(
136-
json.bs
137-
.get(1)
138-
.and_then(|row| row.first())
139-
.ok_or_else(|| {
140-
ProverClientError::InvalidProofData("missing proof B y0 coordinate".to_string())
141-
})?,
142-
)?;
143-
let proof_b_y_1 = deserialize_hex_string_to_be_bytes(
144-
json.bs
145-
.get(1)
146-
.and_then(|row| row.get(1))
147-
.ok_or_else(|| {
148-
ProverClientError::InvalidProofData("missing proof B y1 coordinate".to_string())
149-
})?,
132+
json.bs.get(1).and_then(|row| row.first()).ok_or_else(|| {
133+
ProverClientError::InvalidProofData("missing proof B y0 coordinate".to_string())
134+
})?,
150135
)?;
136+
let proof_b_y_1 =
137+
deserialize_hex_string_to_be_bytes(json.bs.get(1).and_then(|row| row.get(1)).ok_or_else(
138+
|| ProverClientError::InvalidProofData("missing proof B y1 coordinate".to_string()),
139+
)?)?;
151140
let proof_b: [u8; 128] = [proof_b_x_0, proof_b_x_1, proof_b_y_0, proof_b_y_1]
152141
.concat()
153142
.try_into()

‎prover/client/tests/batch_address_append.rs‎

Lines changed: 33 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -38,40 +38,39 @@ async fn prove_batch_address_append() {
3838
IndexedMerkleTree::<Poseidon, usize>::new(DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize, 0)
3939
.unwrap();
4040

41-
let collect_non_inclusion_data =
42-
|tree: &IndexedMerkleTree<Poseidon, usize>, values: &[BigUint]| {
43-
let mut low_element_values = Vec::with_capacity(values.len());
44-
let mut low_element_indices = Vec::with_capacity(values.len());
45-
let mut low_element_next_indices = Vec::with_capacity(values.len());
46-
let mut low_element_next_values = Vec::with_capacity(values.len());
47-
let mut low_element_proofs: Vec<
48-
[[u8; 32]; DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize],
49-
> = Vec::with_capacity(values.len());
50-
51-
for new_element_value in values {
52-
let non_inclusion_proof = tree.get_non_inclusion_proof(new_element_value).unwrap();
53-
54-
low_element_values.push(non_inclusion_proof.leaf_lower_range_value);
55-
low_element_indices.push(non_inclusion_proof.leaf_index);
56-
low_element_next_indices.push(non_inclusion_proof.next_index);
57-
low_element_next_values.push(non_inclusion_proof.leaf_higher_range_value);
58-
low_element_proofs.push(
59-
non_inclusion_proof
60-
.merkle_proof
61-
.as_slice()
62-
.try_into()
63-
.unwrap(),
64-
);
65-
}
66-
67-
(
68-
low_element_values,
69-
low_element_indices,
70-
low_element_next_indices,
71-
low_element_next_values,
72-
low_element_proofs,
73-
)
74-
};
41+
let collect_non_inclusion_data = |tree: &IndexedMerkleTree<Poseidon, usize>,
42+
values: &[BigUint]| {
43+
let mut low_element_values = Vec::with_capacity(values.len());
44+
let mut low_element_indices = Vec::with_capacity(values.len());
45+
let mut low_element_next_indices = Vec::with_capacity(values.len());
46+
let mut low_element_next_values = Vec::with_capacity(values.len());
47+
let mut low_element_proofs: Vec<[[u8; 32]; DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize]> =
48+
Vec::with_capacity(values.len());
49+
50+
for new_element_value in values {
51+
let non_inclusion_proof = tree.get_non_inclusion_proof(new_element_value).unwrap();
52+
53+
low_element_values.push(non_inclusion_proof.leaf_lower_range_value);
54+
low_element_indices.push(non_inclusion_proof.leaf_index);
55+
low_element_next_indices.push(non_inclusion_proof.next_index);
56+
low_element_next_values.push(non_inclusion_proof.leaf_higher_range_value);
57+
low_element_proofs.push(
58+
non_inclusion_proof
59+
.merkle_proof
60+
.as_slice()
61+
.try_into()
62+
.unwrap(),
63+
);
64+
}
65+
66+
(
67+
low_element_values,
68+
low_element_indices,
69+
low_element_next_indices,
70+
low_element_next_values,
71+
low_element_proofs,
72+
)
73+
};
7574

7675
let initial_start_index = relayer_merkle_tree.merkle_tree.rightmost_index;
7776
let initial_root = relayer_merkle_tree.root();

‎sdk-libs/client/src/local_test_validator.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -110,10 +110,7 @@ pub async fn spawn_validator(config: LightValidatorConfig) {
110110
.stderr(Stdio::inherit())
111111
.spawn()
112112
.expect("Failed to start server process");
113-
let status = child
114-
.wait()
115-
.await
116-
.expect("Failed to wait for CLI process");
113+
let status = child.wait().await.expect("Failed to wait for CLI process");
117114
assert!(status.success(), "CLI exited with error: {}", status);
118115
} else {
119116
let _child = Command::new("sh")

‎sdk-tests/sdk-anchor-test/programs/sdk-anchor-test/tests/read_only.rs‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ async fn create_compressed_account(
127127
)
128128
.await?
129129
.value;
130-
let packed_accounts = rpc_result.pack_tree_infos(&mut remaining_accounts);
130+
let packed_accounts = rpc_result
131+
.pack_tree_infos(&mut remaining_accounts)
132+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?;
131133

132134
let output_tree_index = rpc
133135
.get_random_state_tree_info()
@@ -178,6 +180,7 @@ async fn read_sha256_light_system_cpi(
178180

179181
let packed_tree_accounts = rpc_result
180182
.pack_tree_infos(&mut remaining_accounts)
183+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
181184
.state_trees
182185
.unwrap();
183186

@@ -231,6 +234,7 @@ async fn read_sha256_lowlevel(
231234

232235
let packed_tree_accounts = rpc_result
233236
.pack_tree_infos(&mut remaining_accounts)
237+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
234238
.state_trees
235239
.unwrap();
236240

@@ -289,7 +293,9 @@ async fn create_compressed_account_poseidon(
289293
)
290294
.await?
291295
.value;
292-
let packed_accounts = rpc_result.pack_tree_infos(&mut remaining_accounts);
296+
let packed_accounts = rpc_result
297+
.pack_tree_infos(&mut remaining_accounts)
298+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?;
293299

294300
let output_tree_index = rpc
295301
.get_random_state_tree_info()
@@ -340,6 +346,7 @@ async fn read_poseidon_light_system_cpi(
340346

341347
let packed_tree_accounts = rpc_result
342348
.pack_tree_infos(&mut remaining_accounts)
349+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
343350
.state_trees
344351
.unwrap();
345352

@@ -393,6 +400,7 @@ async fn read_poseidon_lowlevel(
393400

394401
let packed_tree_accounts = rpc_result
395402
.pack_tree_infos(&mut remaining_accounts)
403+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
396404
.state_trees
397405
.unwrap();
398406

‎sdk-tests/sdk-anchor-test/programs/sdk-anchor-test/tests/test.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,9 @@ async fn create_compressed_account(
171171
)
172172
.await?
173173
.value;
174-
let packed_accounts = rpc_result.pack_tree_infos(&mut remaining_accounts);
174+
let packed_accounts = rpc_result
175+
.pack_tree_infos(&mut remaining_accounts)
176+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?;
175177

176178
let output_tree_index = rpc
177179
.get_random_state_tree_info()
@@ -223,6 +225,7 @@ async fn update_compressed_account(
223225

224226
let packed_tree_accounts = rpc_result
225227
.pack_tree_infos(&mut remaining_accounts)
228+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
226229
.state_trees
227230
.unwrap();
228231

@@ -277,6 +280,7 @@ async fn close_compressed_account(
277280

278281
let packed_tree_accounts = rpc_result
279282
.pack_tree_infos(&mut remaining_accounts)
283+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
280284
.state_trees
281285
.unwrap();
282286

@@ -340,6 +344,7 @@ async fn reinit_closed_account(
340344

341345
let packed_tree_accounts = rpc_result
342346
.pack_tree_infos(&mut remaining_accounts)
347+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
343348
.state_trees
344349
.unwrap();
345350

@@ -388,6 +393,7 @@ async fn close_compressed_account_permanent(
388393

389394
let packed_tree_accounts = rpc_result
390395
.pack_tree_infos(&mut remaining_accounts)
396+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
391397
.state_trees
392398
.unwrap();
393399

‎sdk-tests/sdk-native-test/tests/test.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,10 @@ pub async fn create_pda(
103103
.value;
104104

105105
let output_merkle_tree_index = accounts.insert_or_get(*merkle_tree_pubkey);
106-
let packed_address_tree_info = rpc_result.pack_tree_infos(&mut accounts).address_trees[0];
106+
let packed_address_tree_info = rpc_result
107+
.pack_tree_infos(&mut accounts)
108+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
109+
.address_trees[0];
107110
let (accounts, system_accounts_offset, tree_accounts_offset) = accounts.to_account_metas();
108111

109112
let instruction_data = CreatePdaInstructionData {
@@ -147,6 +150,7 @@ pub async fn update_pda(
147150

148151
let packed_accounts = rpc_result
149152
.pack_tree_infos(&mut accounts)
153+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
150154
.state_trees
151155
.unwrap();
152156

‎sdk-tests/sdk-pinocchio-v1-test/tests/test.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,10 @@ pub async fn create_pda(
101101
.value;
102102

103103
let output_merkle_tree_index = accounts.insert_or_get(*merkle_tree_pubkey);
104-
let packed_address_tree_info = rpc_result.pack_tree_infos(&mut accounts).address_trees[0];
104+
let packed_address_tree_info = rpc_result
105+
.pack_tree_infos(&mut accounts)
106+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
107+
.address_trees[0];
105108
let (accounts, system_accounts_offset, tree_accounts_offset) = accounts.to_account_metas();
106109
let instruction_data = CreatePdaInstructionData {
107110
proof: rpc_result.proof,
@@ -145,6 +148,7 @@ pub async fn update_pda(
145148

146149
let packed_accounts = rpc_result
147150
.pack_tree_infos(&mut accounts)
151+
.map_err(|error| RpcError::CustomError(format!("Failed to pack tree infos: {error}")))?
148152
.state_trees
149153
.unwrap();
150154

0 commit comments

Comments
 (0)