Skip to content

Commit ddda372

Browse files
ArniStarkwareArniStarkware
committed
feat: post compilation size limit validation
1 parent 1e568dc commit ddda372

File tree

4 files changed

+114
-9
lines changed

4 files changed

+114
-9
lines changed

‎crates/gateway/src/errors.rs

+13
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,19 @@ use crate::compiler_version::{VersionId, VersionIdError};
1818
/// Errors directed towards the end-user, as a result of gateway requests.
1919
#[derive(Debug, Error)]
2020
pub enum GatewayError {
21+
#[error(
22+
"Cannot declare Casm contract class with bytecode size of {bytecode_size}; max allowed \
23+
size: {max_bytecode_size}."
24+
)]
25+
CasmBytecodeSizeTooLarge { bytecode_size: usize, max_bytecode_size: usize },
26+
#[error(
27+
"Cannot declare Casm contract class with size of {contract_class_object_size}; max \
28+
allowed size: {max_contract_class_object_size}."
29+
)]
30+
CasmContractClassObjectSizeTooLarge {
31+
contract_class_object_size: usize,
32+
max_contract_class_object_size: usize,
33+
},
2134
#[error(transparent)]
2235
CompilationError(#[from] starknet_sierra_compile::compile::CompilationUtilError),
2336
#[error(

‎crates/gateway/src/gateway.rs

+35-2
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,14 @@ fn process_tx(
123123
stateless_tx_validator.validate(&tx)?;
124124

125125
// Compile Sierra to Casm.
126+
let sierra_to_casm_compilation_config = SierraToCasmCompilationConfig {
127+
max_bytecode_size: stateless_tx_validator.config.max_bytecode_size,
128+
max_raw_class_size: stateless_tx_validator.config.max_raw_class_size,
129+
};
126130
let optional_class_info = match &tx {
127-
RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx)?),
131+
RPCTransaction::Declare(declare_tx) => {
132+
Some(compile_contract_class(declare_tx, sierra_to_casm_compilation_config)?)
133+
}
128134
_ => None,
129135
};
130136

@@ -139,10 +145,20 @@ fn process_tx(
139145
})
140146
}
141147

148+
// TODO(Arni): Move the gateway compilation util to a dedicated file.
149+
// TODO(Arni): Find a better place for this config.
150+
pub struct SierraToCasmCompilationConfig {
151+
pub max_bytecode_size: usize,
152+
pub max_raw_class_size: usize,
153+
}
154+
142155
/// Formats the contract class for compilation, compiles it, and returns the compiled contract class
143156
/// wrapped in a [`ClassInfo`].
144157
/// Assumes the contract class is of a Sierra program which is compiled to Casm.
145-
pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResult<ClassInfo> {
158+
pub fn compile_contract_class(
159+
declare_tx: &RPCDeclareTransaction,
160+
config: SierraToCasmCompilationConfig,
161+
) -> GatewayResult<ClassInfo> {
146162
let RPCDeclareTransaction::V3(tx) = declare_tx;
147163
let starknet_api_contract_class = &tx.contract_class;
148164
let cairo_lang_contract_class =
@@ -159,6 +175,23 @@ pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResu
159175
}
160176
};
161177

178+
let bytecode_size = casm_contract_class.bytecode.len();
179+
if bytecode_size > config.max_bytecode_size {
180+
return Err(GatewayError::CasmBytecodeSizeTooLarge {
181+
bytecode_size,
182+
max_bytecode_size: config.max_bytecode_size,
183+
});
184+
}
185+
let contract_class_object_size = serde_json::to_string(&casm_contract_class)
186+
.expect("Unexpected error serializing Casm contract class.")
187+
.len();
188+
if contract_class_object_size > config.max_raw_class_size {
189+
return Err(GatewayError::CasmContractClassObjectSizeTooLarge {
190+
contract_class_object_size,
191+
max_contract_class_object_size: config.max_raw_class_size,
192+
});
193+
}
194+
162195
let hash_result =
163196
CompiledClassHash(felt_to_stark_felt(&casm_contract_class.compiled_class_hash()));
164197
if hash_result != tx.compiled_class_hash {

‎crates/gateway/src/gateway_test.rs

+55-5
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@ use tokio::task;
2222

2323
use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig};
2424
use crate::errors::GatewayError;
25-
use crate::gateway::{add_tx, compile_contract_class, AppState, SharedMempoolClient};
25+
use crate::gateway::{
26+
add_tx, compile_contract_class, AppState, SharedMempoolClient, SierraToCasmCompilationConfig,
27+
};
2628
use crate::starknet_api_test_utils::{declare_tx, deploy_account_tx, invoke_tx};
2729
use crate::state_reader_test_utils::{
2830
local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account,
@@ -33,6 +35,8 @@ use crate::stateless_transaction_validator::StatelessTransactionValidator;
3335
use crate::utils::{external_tx_to_account_tx, get_tx_hash};
3436

3537
const MEMPOOL_INVOCATIONS_QUEUE_SIZE: usize = 32;
38+
const SIERRA_TO_CASM_COMPILATION_CONFIG: SierraToCasmCompilationConfig =
39+
SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: usize::MAX };
3640

3741
#[fixture]
3842
fn mempool() -> Mempool {
@@ -121,14 +125,57 @@ fn test_compile_contract_class_compiled_class_hash_missmatch() {
121125
tx.compiled_class_hash = supplied_hash;
122126
let declare_tx = RPCDeclareTransaction::V3(tx);
123127

124-
let result = compile_contract_class(&declare_tx);
128+
let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG);
125129
assert_matches!(
126130
result.unwrap_err(),
127131
GatewayError::CompiledClassHashMismatch { supplied, hash_result }
128132
if supplied == supplied_hash && hash_result == expected_hash_result
129133
);
130134
}
131135

136+
#[rstest]
137+
#[case::bytecode_size(
138+
SierraToCasmCompilationConfig { max_bytecode_size: 1, max_raw_class_size: usize::MAX},
139+
GatewayError::CasmBytecodeSizeTooLarge { bytecode_size: 4800, max_bytecode_size: 1 }
140+
)]
141+
#[case::raw_class_size(
142+
SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: 1},
143+
GatewayError::CasmContractClassObjectSizeTooLarge {
144+
contract_class_object_size: 111037, max_contract_class_object_size: 1
145+
}
146+
)]
147+
fn test_compile_contract_class_size_validation(
148+
#[case] sierra_to_casm_compilation_config: SierraToCasmCompilationConfig,
149+
#[case] expected_error: GatewayError,
150+
) {
151+
let declare_tx = match declare_tx() {
152+
RPCTransaction::Declare(declare_tx) => declare_tx,
153+
_ => panic!("Invalid transaction type"),
154+
};
155+
156+
let result = compile_contract_class(&declare_tx, sierra_to_casm_compilation_config);
157+
if let GatewayError::CasmBytecodeSizeTooLarge {
158+
bytecode_size: expected_bytecode_size, ..
159+
} = expected_error
160+
{
161+
assert_matches!(
162+
result.unwrap_err(),
163+
GatewayError::CasmBytecodeSizeTooLarge { bytecode_size, .. }
164+
if bytecode_size == expected_bytecode_size
165+
)
166+
} else if let GatewayError::CasmContractClassObjectSizeTooLarge {
167+
contract_class_object_size: expected_contract_class_object_size,
168+
..
169+
} = expected_error
170+
{
171+
assert_matches!(
172+
result.unwrap_err(),
173+
GatewayError::CasmContractClassObjectSizeTooLarge { contract_class_object_size, .. }
174+
if contract_class_object_size == expected_contract_class_object_size
175+
)
176+
}
177+
}
178+
132179
#[test]
133180
fn test_compile_contract_class_bad_sierra() {
134181
let mut tx = assert_matches!(
@@ -139,7 +186,7 @@ fn test_compile_contract_class_bad_sierra() {
139186
tx.contract_class.sierra_program = tx.contract_class.sierra_program[..100].to_vec();
140187
let declare_tx = RPCDeclareTransaction::V3(tx);
141188

142-
let result = compile_contract_class(&declare_tx);
189+
let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG);
143190
assert_matches!(
144191
result.unwrap_err(),
145192
GatewayError::CompilationError(CompilationUtilError::AllowedLibfuncsError(
@@ -157,7 +204,8 @@ fn test_compile_contract_class() {
157204
let RPCDeclareTransaction::V3(declare_tx_v3) = &declare_tx;
158205
let contract_class = &declare_tx_v3.contract_class;
159206

160-
let class_info = compile_contract_class(&declare_tx).unwrap();
207+
let class_info =
208+
compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG).unwrap();
161209
assert_matches!(class_info.contract_class(), ContractClass::V1(_));
162210
assert_eq!(class_info.sierra_program_length(), contract_class.sierra_program.len());
163211
assert_eq!(class_info.abi_length(), contract_class.abi.len());
@@ -169,7 +217,9 @@ async fn to_bytes(res: Response) -> Bytes {
169217

170218
fn calculate_hash(external_tx: &RPCTransaction) -> TransactionHash {
171219
let optional_class_info = match &external_tx {
172-
RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx).unwrap()),
220+
RPCTransaction::Declare(declare_tx) => {
221+
Some(compile_contract_class(declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG).unwrap())
222+
}
173223
_ => None,
174224
};
175225

‎crates/gateway/src/stateful_transaction_validator_test.rs

+11-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use starknet_api::transaction::TransactionHash;
99

1010
use crate::config::StatefulTransactionValidatorConfig;
1111
use crate::errors::{StatefulTransactionValidatorError, StatefulTransactionValidatorResult};
12-
use crate::gateway::compile_contract_class;
12+
use crate::gateway::{compile_contract_class, SierraToCasmCompilationConfig};
1313
use crate::starknet_api_test_utils::{
1414
declare_tx, deploy_account_tx, invoke_tx, VALID_L1_GAS_MAX_AMOUNT,
1515
VALID_L1_GAS_MAX_PRICE_PER_UNIT,
@@ -80,7 +80,16 @@ fn test_stateful_tx_validator(
8080
},
8181
};
8282
let optional_class_info = match &external_tx {
83-
RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx).unwrap()),
83+
RPCTransaction::Declare(declare_tx) => Some(
84+
compile_contract_class(
85+
declare_tx,
86+
SierraToCasmCompilationConfig {
87+
max_bytecode_size: usize::MAX,
88+
max_raw_class_size: usize::MAX,
89+
},
90+
)
91+
.unwrap(),
92+
),
8493
_ => None,
8594
};
8695

0 commit comments

Comments
 (0)