From ab37864b514a3a521afc8cc1bb485b775e100b43 Mon Sep 17 00:00:00 2001 From: Arni Hod Date: Wed, 26 Jun 2024 11:43:46 +0300 Subject: [PATCH] feat: post compilation size limit validation --- crates/gateway/src/errors.rs | 13 ++++ crates/gateway/src/gateway.rs | 37 +++++++++++- crates/gateway/src/gateway_test.rs | 60 +++++++++++++++++-- .../stateful_transaction_validator_test.rs | 13 +++- 4 files changed, 114 insertions(+), 9 deletions(-) diff --git a/crates/gateway/src/errors.rs b/crates/gateway/src/errors.rs index a1cc259e2..ce5fb2962 100644 --- a/crates/gateway/src/errors.rs +++ b/crates/gateway/src/errors.rs @@ -19,6 +19,19 @@ use crate::compiler_version::{VersionId, VersionIdError}; /// Errors directed towards the end-user, as a result of gateway requests. #[derive(Debug, Error)] pub enum GatewayError { + #[error( + "Cannot declare Casm contract class with bytecode size of {bytecode_size}; max allowed \ + size: {max_bytecode_size}." + )] + CasmBytecodeSizeTooLarge { bytecode_size: usize, max_bytecode_size: usize }, + #[error( + "Cannot declare Casm contract class with size of {contract_class_object_size}; max \ + allowed size: {max_contract_class_object_size}." + )] + CasmContractClassObjectSizeTooLarge { + contract_class_object_size: usize, + max_contract_class_object_size: usize, + }, #[error(transparent)] CompilationError(#[from] CompilationUtilError), #[error( diff --git a/crates/gateway/src/gateway.rs b/crates/gateway/src/gateway.rs index 527766bdb..d37452280 100644 --- a/crates/gateway/src/gateway.rs +++ b/crates/gateway/src/gateway.rs @@ -124,8 +124,14 @@ fn process_tx( stateless_tx_validator.validate(&tx)?; // Compile Sierra to Casm. + let sierra_to_casm_compilation_config = SierraToCasmCompilationConfig { + max_bytecode_size: stateless_tx_validator.config.max_bytecode_size, + max_raw_class_size: stateless_tx_validator.config.max_raw_class_size, + }; let optional_class_info = match &tx { - RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx)?), + RPCTransaction::Declare(declare_tx) => { + Some(compile_contract_class(declare_tx, sierra_to_casm_compilation_config)?) + } _ => None, }; @@ -140,10 +146,20 @@ fn process_tx( }) } +// TODO(Arni): Move the gateway compilation util to a dedicated file. +// TODO(Arni): Find a better place for this config. +pub struct SierraToCasmCompilationConfig { + pub max_bytecode_size: usize, + pub max_raw_class_size: usize, +} + /// Formats the contract class for compilation, compiles it, and returns the compiled contract class /// wrapped in a [`ClassInfo`]. /// Assumes the contract class is of a Sierra program which is compiled to Casm. -pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResult { +pub fn compile_contract_class( + declare_tx: &RPCDeclareTransaction, + config: SierraToCasmCompilationConfig, +) -> GatewayResult { let RPCDeclareTransaction::V3(tx) = declare_tx; let starknet_api_contract_class = &tx.contract_class; let cairo_lang_contract_class = @@ -160,6 +176,23 @@ pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResu } }; + let bytecode_size = casm_contract_class.bytecode.len(); + if bytecode_size > config.max_bytecode_size { + return Err(GatewayError::CasmBytecodeSizeTooLarge { + bytecode_size, + max_bytecode_size: config.max_bytecode_size, + }); + } + let contract_class_object_size = serde_json::to_string(&casm_contract_class) + .expect("Unexpected error serializing Casm contract class.") + .len(); + if contract_class_object_size > config.max_raw_class_size { + return Err(GatewayError::CasmContractClassObjectSizeTooLarge { + contract_class_object_size, + max_contract_class_object_size: config.max_raw_class_size, + }); + } + let hash_result = CompiledClassHash(felt_to_stark_felt(&casm_contract_class.compiled_class_hash())); if hash_result != tx.compiled_class_hash { diff --git a/crates/gateway/src/gateway_test.rs b/crates/gateway/src/gateway_test.rs index 6308054f4..2251446da 100644 --- a/crates/gateway/src/gateway_test.rs +++ b/crates/gateway/src/gateway_test.rs @@ -23,7 +23,9 @@ use tokio::task; use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig}; use crate::errors::GatewayError; -use crate::gateway::{add_tx, compile_contract_class, AppState, SharedMempoolClient}; +use crate::gateway::{ + add_tx, compile_contract_class, AppState, SharedMempoolClient, SierraToCasmCompilationConfig, +}; use crate::state_reader_test_utils::{ local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account, TestStateReaderFactory, @@ -33,6 +35,8 @@ use crate::stateless_transaction_validator::StatelessTransactionValidator; use crate::utils::{external_tx_to_account_tx, get_tx_hash}; const MEMPOOL_INVOCATIONS_QUEUE_SIZE: usize = 32; +const SIERRA_TO_CASM_COMPILATION_CONFIG: SierraToCasmCompilationConfig = + SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: usize::MAX }; #[fixture] fn mempool() -> Mempool { @@ -121,7 +125,7 @@ fn test_compile_contract_class_compiled_class_hash_missmatch() { tx.compiled_class_hash = supplied_hash; let declare_tx = RPCDeclareTransaction::V3(tx); - let result = compile_contract_class(&declare_tx); + let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG); assert_matches!( result.unwrap_err(), GatewayError::CompiledClassHashMismatch { supplied, hash_result } @@ -129,6 +133,49 @@ fn test_compile_contract_class_compiled_class_hash_missmatch() { ); } +#[rstest] +#[case::bytecode_size( + SierraToCasmCompilationConfig { max_bytecode_size: 1, max_raw_class_size: usize::MAX}, + GatewayError::CasmBytecodeSizeTooLarge { bytecode_size: 4800, max_bytecode_size: 1 } +)] +#[case::raw_class_size( + SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: 1}, + GatewayError::CasmContractClassObjectSizeTooLarge { + contract_class_object_size: 111037, max_contract_class_object_size: 1 + } +)] +fn test_compile_contract_class_size_validation( + #[case] sierra_to_casm_compilation_config: SierraToCasmCompilationConfig, + #[case] expected_error: GatewayError, +) { + let declare_tx = match declare_tx() { + RPCTransaction::Declare(declare_tx) => declare_tx, + _ => panic!("Invalid transaction type"), + }; + + let result = compile_contract_class(&declare_tx, sierra_to_casm_compilation_config); + if let GatewayError::CasmBytecodeSizeTooLarge { + bytecode_size: expected_bytecode_size, .. + } = expected_error + { + assert_matches!( + result.unwrap_err(), + GatewayError::CasmBytecodeSizeTooLarge { bytecode_size, .. } + if bytecode_size == expected_bytecode_size + ) + } else if let GatewayError::CasmContractClassObjectSizeTooLarge { + contract_class_object_size: expected_contract_class_object_size, + .. + } = expected_error + { + assert_matches!( + result.unwrap_err(), + GatewayError::CasmContractClassObjectSizeTooLarge { contract_class_object_size, .. } + if contract_class_object_size == expected_contract_class_object_size + ) + } +} + #[test] fn test_compile_contract_class_bad_sierra() { let mut tx = assert_matches!( @@ -139,7 +186,7 @@ fn test_compile_contract_class_bad_sierra() { tx.contract_class.sierra_program = tx.contract_class.sierra_program[..100].to_vec(); let declare_tx = RPCDeclareTransaction::V3(tx); - let result = compile_contract_class(&declare_tx); + let result = compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG); assert_matches!( result.unwrap_err(), GatewayError::CompilationError(CompilationUtilError::AllowedLibfuncsError( @@ -157,7 +204,8 @@ fn test_compile_contract_class() { let RPCDeclareTransaction::V3(declare_tx_v3) = &declare_tx; let contract_class = &declare_tx_v3.contract_class; - let class_info = compile_contract_class(&declare_tx).unwrap(); + let class_info = + compile_contract_class(&declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG).unwrap(); assert_matches!(class_info.contract_class(), ContractClass::V1(_)); assert_eq!(class_info.sierra_program_length(), contract_class.sierra_program.len()); assert_eq!(class_info.abi_length(), contract_class.abi.len()); @@ -169,7 +217,9 @@ async fn to_bytes(res: Response) -> Bytes { fn calculate_hash(external_tx: &RPCTransaction) -> TransactionHash { let optional_class_info = match &external_tx { - RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx).unwrap()), + RPCTransaction::Declare(declare_tx) => { + Some(compile_contract_class(declare_tx, SIERRA_TO_CASM_COMPILATION_CONFIG).unwrap()) + } _ => None, }; diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs index d3821782f..cf29a825a 100644 --- a/crates/gateway/src/stateful_transaction_validator_test.rs +++ b/crates/gateway/src/stateful_transaction_validator_test.rs @@ -13,7 +13,7 @@ use test_utils::starknet_api_test_utils::{ use crate::config::StatefulTransactionValidatorConfig; use crate::errors::{StatefulTransactionValidatorError, StatefulTransactionValidatorResult}; -use crate::gateway::compile_contract_class; +use crate::gateway::{compile_contract_class, SierraToCasmCompilationConfig}; use crate::state_reader_test_utils::{ local_test_state_reader_factory, local_test_state_reader_factory_for_deploy_account, TestStateReaderFactory, @@ -80,7 +80,16 @@ fn test_stateful_tx_validator( }, }; let optional_class_info = match &external_tx { - RPCTransaction::Declare(declare_tx) => Some(compile_contract_class(declare_tx).unwrap()), + RPCTransaction::Declare(declare_tx) => Some( + compile_contract_class( + declare_tx, + SierraToCasmCompilationConfig { + max_bytecode_size: usize::MAX, + max_raw_class_size: usize::MAX, + }, + ) + .unwrap(), + ), _ => None, };