feat: post compilation size limit validation

ArniStarkware · ArniStarkware · commit b2ca15c0d9c2 · 2024-07-07T11:17:34.000+03:00
diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
@@ -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 =
@@ -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 {
diff --git a/crates/gateway/src/compilation_test.rs b/crates/gateway/src/compilation_test.rs
@@ -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!(
@@ -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!(
@@ -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(
@@ -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());
diff --git 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
@@ -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;
@@ -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,
     };
 
diff --git a/crates/gateway/src/gateway_test.rs b/crates/gateway/src/gateway_test.rs
@@ -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::{
@@ -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,
     };
 
diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs
@@ -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::{
@@ -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,
     };
 
