Skip to content

Commit

Permalink
feat: post compilation size limit validation
Browse files Browse the repository at this point in the history
  • Loading branch information
ArniStarkware committed Jul 7, 2024
1 parent fb92713 commit b2ca15c
Show file tree
Hide file tree
Showing 6 changed files with 142 additions and 10 deletions.
48 changes: 47 additions & 1 deletion crates/gateway/src/compilation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,22 @@ use crate::utils::is_subsequence;
#[path = "compilation_test.rs"]
mod compilation_test;

// 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,
}

// TODO(Define a function for `compile_contract_class` - which ignores the `config` parameter).

/// 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<ClassInfo> {
pub fn compile_contract_class(
declare_tx: &RPCDeclareTransaction,
config: SierraToCasmCompilationConfig,
) -> GatewayResult<ClassInfo> {
let RPCDeclareTransaction::V3(tx) = declare_tx;
let starknet_api_contract_class = &tx.contract_class;
let cairo_lang_contract_class =
Expand All @@ -39,6 +51,40 @@ pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResu
};
validate_casm_class(&casm_contract_class)?;

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 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(casm_contract_class.compiled_class_hash());
if hash_result != tx.compiled_class_hash {
return Err(GatewayError::CompiledClassHashMismatch {
Expand Down
56 changes: 52 additions & 4 deletions crates/gateway/src/compilation_test.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,18 @@
use assert_matches::assert_matches;
use blockifier::execution::contract_class::ContractClass;
use cairo_lang_starknet_classes::allowed_libfuncs::AllowedLibfuncsError;
use rstest::rstest;
use starknet_api::core::CompiledClassHash;
use starknet_api::rpc_transaction::{RPCDeclareTransaction, RPCTransaction};
use starknet_sierra_compile::errors::CompilationUtilError;
use test_utils::starknet_api_test_utils::declare_tx;

use crate::compilation::compile_contract_class;
use crate::compilation::{compile_contract_class, SierraToCasmCompilationConfig};
use crate::errors::GatewayError;

const SIERRA_TO_CASM_COMPILATION_CONFIG: SierraToCasmCompilationConfig =
SierraToCasmCompilationConfig { max_bytecode_size: usize::MAX, max_raw_class_size: usize::MAX };

#[test]
fn test_compile_contract_class_compiled_class_hash_missmatch() {
let mut tx = assert_matches!(
Expand All @@ -21,14 +25,57 @@ 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 }
if supplied == supplied_hash && hash_result == expected_hash_result
);
}

#[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!(
Expand All @@ -39,7 +86,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(
Expand All @@ -57,7 +104,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());
Expand Down
13 changes: 13 additions & 0 deletions crates/gateway/src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
10 changes: 8 additions & 2 deletions crates/gateway/src/gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use starknet_mempool_infra::component_runner::{ComponentRunner, ComponentStartEr
use starknet_mempool_types::communication::SharedMempoolClient;
use starknet_mempool_types::mempool_types::{Account, MempoolInput};

use crate::compilation::compile_contract_class;
use crate::compilation::{compile_contract_class, SierraToCasmCompilationConfig};
use crate::config::{GatewayConfig, GatewayNetworkConfig, RpcStateReaderConfig};
use crate::errors::{GatewayError, GatewayResult, GatewayRunError};
use crate::rpc_state_reader::RpcStateReaderFactory;
Expand Down Expand Up @@ -118,8 +118,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,
};

Expand Down
12 changes: 11 additions & 1 deletion crates/gateway/src/gateway_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use test_utils::starknet_api_test_utils::{declare_tx, deploy_account_tx, invoke_
use tokio::sync::mpsc::channel;
use tokio::task;

use crate::compilation::SierraToCasmCompilationConfig;
use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig};
use crate::gateway::{add_tx, compile_contract_class, AppState, SharedMempoolClient};
use crate::state_reader_test_utils::{
Expand Down Expand Up @@ -110,7 +111,16 @@ 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,
SierraToCasmCompilationConfig {
max_bytecode_size: usize::MAX,
max_raw_class_size: usize::MAX,
},
)
.unwrap(),
),
_ => None,
};

Expand Down
13 changes: 11 additions & 2 deletions crates/gateway/src/stateful_transaction_validator_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use test_utils::starknet_api_test_utils::{
VALID_L1_GAS_MAX_PRICE_PER_UNIT,
};

use crate::compilation::compile_contract_class;
use crate::compilation::{compile_contract_class, SierraToCasmCompilationConfig};
use crate::config::StatefulTransactionValidatorConfig;
use crate::errors::{StatefulTransactionValidatorError, StatefulTransactionValidatorResult};
use crate::state_reader_test_utils::{
Expand Down Expand Up @@ -81,7 +81,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,
};

Expand Down

0 comments on commit b2ca15c

Please sign in to comment.