Skip to content

Commit ab37864

Browse files
ArniStarkwareArniStarkware
committed
feat: post compilation size limit validation
1 parent e51ec09 commit ab37864

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
@@ -19,6 +19,19 @@ use crate::compiler_version::{VersionId, VersionIdError};
1919
/// Errors directed towards the end-user, as a result of gateway requests.
2020
#[derive(Debug, Error)]
2121
pub enum GatewayError {
22+
#[error(
23+
"Cannot declare Casm contract class with bytecode size of {bytecode_size}; max allowed \
24+
size: {max_bytecode_size}."
25+
)]
26+
CasmBytecodeSizeTooLarge { bytecode_size: usize, max_bytecode_size: usize },
27+
#[error(
28+
"Cannot declare Casm contract class with size of {contract_class_object_size}; max \
29+
allowed size: {max_contract_class_object_size}."
30+
)]
31+
CasmContractClassObjectSizeTooLarge {
32+
contract_class_object_size: usize,
33+
max_contract_class_object_size: usize,
34+
},
2235
#[error(transparent)]
2336
CompilationError(#[from] CompilationUtilError),
2437
#[error(

‎crates/gateway/src/gateway.rs

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

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

@@ -140,10 +146,20 @@ fn process_tx(
140146
})
141147
}
142148

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

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

‎crates/gateway/src/gateway_test.rs

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

2424
use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig};
2525
use crate::errors::GatewayError;
26-
use crate::gateway::{add_tx, compile_contract_class, AppState, SharedMempoolClient};
26+
use crate::gateway::{
27+
add_tx, compile_contract_class, AppState, SharedMempoolClient, SierraToCasmCompilationConfig,
28+
};
2729
use crate::state_reader_test_utils::{
2830
local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account,
2931
TestStateReaderFactory,
@@ -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
@@ -13,7 +13,7 @@ use test_utils::starknet_api_test_utils::{
1313

1414
use crate::config::StatefulTransactionValidatorConfig;
1515
use crate::errors::{StatefulTransactionValidatorError, StatefulTransactionValidatorResult};
16-
use crate::gateway::compile_contract_class;
16+
use crate::gateway::{compile_contract_class, SierraToCasmCompilationConfig};
1717
use crate::state_reader_test_utils::{
1818
local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account,
1919
TestStateReaderFactory,
@@ -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)