Skip to content

Commit 6ec8522

Browse files
ArniStarkwareArniStarkware
committed
feat: post compilation size limit validation
1 parent 203eed8 commit 6ec8522

File tree

4 files changed

+113
-9
lines changed

4 files changed

+113
-9
lines changed

‎crates/gateway/src/errors.rs

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

+54-5
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ use tokio::task;
2020

2121
use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig};
2222
use crate::errors::GatewayError;
23-
use crate::gateway::{add_tx, compile_contract_class, AppState, SharedMempoolClient};
23+
use crate::gateway::{
24+
add_tx, compile_contract_class, AppState, SharedMempoolClient, SierraToCasmCompilationConfig,
25+
};
2426
use crate::starknet_api_test_utils::{declare_tx, deploy_account_tx, invoke_tx};
2527
use crate::state_reader_test_utils::{
2628
local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account,
@@ -31,6 +33,8 @@ use crate::stateless_transaction_validator::StatelessTransactionValidator;
3133
use crate::utils::{external_tx_to_account_tx, get_tx_hash};
3234

3335
const MEMPOOL_INVOCATIONS_QUEUE_SIZE: usize = 32;
36+
const SIERRA_TO_CASM_COMPILATION_CONFIG: SierraToCasmCompilationConfig =
37+
SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: usize::MAX };
3438

3539
#[fixture]
3640
fn mempool() -> Mempool {
@@ -105,7 +109,7 @@ async fn test_add_tx(
105109
}
106110

107111
#[test]
108-
fn test_compile_contract_class_failure() {
112+
fn test_compile_contract_class_compiled_class_hash_missmatch() {
109113
let mut declare_tx_v3 = match declare_tx() {
110114
RPCTransaction::Declare(RPCDeclareTransaction::V3(declare_tx)) => declare_tx,
111115
_ => panic!("Invalid transaction type"),
@@ -116,14 +120,57 @@ fn test_compile_contract_class_failure() {
116120
declare_tx_v3.compiled_class_hash = supplied_hash;
117121
let declare_tx = RPCDeclareTransaction::V3(declare_tx_v3);
118122

119-
let result = compile_contract_class(&declare_tx);
123+
let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG);
120124
assert_matches!(
121125
result.unwrap_err(),
122126
GatewayError::CompiledClassHashMismatch { supplied, hash_result }
123127
if supplied == supplied_hash && hash_result == expected_hash_result
124128
);
125129
}
126130

131+
#[rstest]
132+
#[case::bytecode_size(
133+
SierraToCasmCompilationConfig { max_bytecode_size: 1, max_raw_class_size: usize::MAX},
134+
GatewayError::CasmBytecodeSizeTooLarge { bytecode_size: 4800, max_bytecode_size: 1 }
135+
)]
136+
#[case::raw_class_size(
137+
SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: 1},
138+
GatewayError::CasmContractClassObjectSizeTooLarge {
139+
contract_class_object_size: 111037, max_contract_class_object_size: 1
140+
}
141+
)]
142+
fn test_compile_contract_class_size_validation(
143+
#[case] sierra_to_casm_compilation_config: SierraToCasmCompilationConfig,
144+
#[case] expected_error: GatewayError,
145+
) {
146+
let declare_tx = match declare_tx() {
147+
RPCTransaction::Declare(declare_tx) => declare_tx,
148+
_ => panic!("Invalid transaction type"),
149+
};
150+
151+
let result = compile_contract_class(&declare_tx, sierra_to_casm_compilation_config);
152+
if let GatewayError::CasmBytecodeSizeTooLarge {
153+
bytecode_size: expected_bytecode_size, ..
154+
} = expected_error
155+
{
156+
assert_matches!(
157+
result.unwrap_err(),
158+
GatewayError::CasmBytecodeSizeTooLarge { bytecode_size, .. }
159+
if bytecode_size == expected_bytecode_size
160+
)
161+
} else if let GatewayError::CasmContractClassObjectSizeTooLarge {
162+
contract_class_object_size: expected_contract_class_object_size,
163+
..
164+
} = expected_error
165+
{
166+
assert_matches!(
167+
result.unwrap_err(),
168+
GatewayError::CasmContractClassObjectSizeTooLarge { contract_class_object_size, .. }
169+
if contract_class_object_size == expected_contract_class_object_size
170+
)
171+
}
172+
}
173+
127174
#[test]
128175
fn test_compile_contract_class() {
129176
let declare_tx = match declare_tx() {
@@ -133,7 +180,7 @@ fn test_compile_contract_class() {
133180
let RPCDeclareTransaction::V3(declare_tx_v3) = &declare_tx;
134181
let contract_class = &declare_tx_v3.contract_class;
135182

136-
let result = compile_contract_class(&declare_tx);
183+
let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG);
137184
assert_matches!(
138185
result,
139186
Ok(class_info)
@@ -151,7 +198,9 @@ async fn to_bytes(res: Response) -> Bytes {
151198

152199
fn calculate_hash(external_tx: &RPCTransaction) -> TransactionHash {
153200
let optional_class_info = match &external_tx {
154-
RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx).unwrap()),
201+
RPCTransaction::Declare(declare_tx) => {
202+
Some(compile_contract_class(declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG).unwrap())
203+
}
155204
_ => None,
156205
};
157206

‎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)