feat: post compilation size limit validation

starkware-libs · Jul 7, 2024 · 6b2e9f1 · 6b2e9f1
1 parent 9622e9a
commit 6b2e9f1
Show file tree

Hide file tree

Showing 9 changed files with 194 additions and 12 deletions.
diff --git a/config/default_config.json b/config/default_config.json
@@ -9,6 +9,16 @@
     "privacy": "Public",
     "value": true
   },
+  "gateway_config.gateway_compiler_config.max_bytecode_size": {
+    "description": "The maximum bytecode size allowed for a contract.",
+    "privacy": "Public",
+    "value": 0
+  },
+  "gateway_config.gateway_compiler_config.max_raw_class_size": {
+    "description": "The maximum raw class size allowed for a contract.",
+    "privacy": "Public",
+    "value": 0
+  },
   "gateway_config.network_config.ip": {
     "description": "The gateway server ip.",
     "privacy": "Public",
@@ -49,11 +59,21 @@
     "privacy": "Public",
     "value": 0
   },
+  "gateway_config.stateless_tx_validator_config.max_bytecode_size": {
+    "description": "The maximum bytecode size allowed for a contract.",
+    "privacy": "Public",
+    "value": 0
+  },
   "gateway_config.stateless_tx_validator_config.max_calldata_length": {
     "description": "Validates that a transaction has calldata length less than or equal to this value.",
     "privacy": "Public",
     "value": 0
   },
+  "gateway_config.stateless_tx_validator_config.max_raw_class_size": {
+    "description": "The maximum raw class size allowed for a contract.",
+    "privacy": "Public",
+    "value": 0
+  },
   "gateway_config.stateless_tx_validator_config.max_sierra_version.major": {
     "description": "The major version of the configuration.",
     "privacy": "Public",

diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
@@ -19,6 +19,15 @@ 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).
+#[derive(Clone)]
 pub struct GatewayCompiler {
     pub config: GatewayCompilerConfig,
 }
@@ -47,6 +56,7 @@ impl GatewayCompiler {
             }
         };
         self.validate_casm_class(&casm_contract_class)?;
+        self.validate_casm_class_size(&casm_contract_class)?;
 
         let hash_result = CompiledClassHash(casm_contract_class.compiled_class_hash());
         if hash_result != tx.compiled_class_hash {
@@ -85,6 +95,47 @@ impl GatewayCompiler {
         }
         Ok(())
     }
+
+    fn validate_casm_class_size(
+        &self,
+        casm_contract_class: &CasmContractClass,
+    ) -> Result<(), GatewayError> {
+        let bytecode_size = casm_contract_class.bytecode.len();
+        if bytecode_size > self.config.max_bytecode_size {
+            return Err(GatewayError::CasmBytecodeSizeTooLarge {
+                bytecode_size,
+                max_bytecode_size: self.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 > self.config.max_raw_class_size {
+            return Err(GatewayError::CasmContractClassObjectSizeTooLarge {
+                contract_class_object_size,
+                max_contract_class_object_size: self.config.max_raw_class_size,
+            });
+        }
+
+        let bytecode_size = casm_contract_class.bytecode.len();
+        if bytecode_size > self.config.max_bytecode_size {
+            return Err(GatewayError::CasmBytecodeSizeTooLarge {
+                bytecode_size,
+                max_bytecode_size: self.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 > self.config.max_raw_class_size {
+            return Err(GatewayError::CasmContractClassObjectSizeTooLarge {
+                contract_class_object_size,
+                max_contract_class_object_size: self.config.max_raw_class_size,
+            });
+        }
+
+        Ok(())
+    }
 }
 
 // TODO(Arni): Add to a config.

diff --git a/crates/gateway/src/compilation_config.rs b/crates/gateway/src/compilation_config.rs
@@ -1,15 +1,31 @@
 use std::collections::BTreeMap;
 
-use papyrus_config::dumping::SerializeConfig;
-use papyrus_config::{ParamPath, SerializedParam};
+use papyrus_config::dumping::{ser_param, SerializeConfig};
+use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
 use serde::{Deserialize, Serialize};
 use validator::Validate;
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize, Validate, PartialEq)]
-pub struct GatewayCompilerConfig {}
+pub struct GatewayCompilerConfig {
+    pub max_bytecode_size: usize,
+    pub max_raw_class_size: usize,
+}
 
 impl SerializeConfig for GatewayCompilerConfig {
     fn dump(&self) -> BTreeMap<ParamPath, SerializedParam> {
-        BTreeMap::new()
+        BTreeMap::from_iter([
+            ser_param(
+                "max_bytecode_size",
+                &self.max_bytecode_size,
+                "The maximum bytecode size allowed for a contract.",
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "max_raw_class_size",
+                &self.max_raw_class_size,
+                "The maximum raw class size allowed for a contract.",
+                ParamPrivacyInput::Public,
+            ),
+        ])
     }
 }
diff --git a/crates/gateway/src/compilation_test.rs b/crates/gateway/src/compilation_test.rs
@@ -1,12 +1,14 @@
 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::GatewayCompiler;
+use crate::compilation_config::GatewayCompilerConfig;
 use crate::errors::GatewayError;
 
 #[test]
@@ -21,14 +23,61 @@ fn test_compile_contract_class_compiled_class_hash_missmatch() {
     tx.compiled_class_hash = supplied_hash;
     let declare_tx = RPCDeclareTransaction::V3(tx);
 
-    let result = GatewayCompiler { config: Default::default() }.compile_contract_class(&declare_tx);
+    let result = GatewayCompiler {
+        config: GatewayCompilerConfig { max_bytecode_size: 4800, max_raw_class_size: 111037 },
+    }
+    .compile_contract_class(&declare_tx);
     assert_matches!(
         result.unwrap_err(),
         GatewayError::CompiledClassHashMismatch { supplied, hash_result }
         if supplied == supplied_hash && hash_result == expected_hash_result
     );
 }
 
+#[rstest]
+#[case::bytecode_size(
+    GatewayCompilerConfig { max_bytecode_size: 1, max_raw_class_size: usize::MAX},
+    GatewayError::CasmBytecodeSizeTooLarge { bytecode_size: 4800, max_bytecode_size: 1 }
+)]
+#[case::raw_class_size(
+    GatewayCompilerConfig { 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: GatewayCompilerConfig,
+    #[case] expected_error: GatewayError,
+) {
+    let declare_tx = match declare_tx() {
+        RPCTransaction::Declare(declare_tx) => declare_tx,
+        _ => panic!("Invalid transaction type"),
+    };
+
+    let gateway_compiler = GatewayCompiler { config: sierra_to_casm_compilation_config };
+    let result = gateway_compiler.compile_contract_class(&declare_tx);
+    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!(
@@ -57,8 +106,11 @@ fn test_compile_contract_class() {
     let RPCDeclareTransaction::V3(declare_tx_v3) = &declare_tx;
     let contract_class = &declare_tx_v3.contract_class;
 
-    let class_info =
-        GatewayCompiler { config: Default::default() }.compile_contract_class(&declare_tx).unwrap();
+    let class_info = GatewayCompiler {
+        config: GatewayCompilerConfig { max_bytecode_size: 4800, max_raw_class_size: 111037 },
+    }
+    .compile_contract_class(&declare_tx)
+    .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/config.rs b/crates/gateway/src/config.rs
@@ -125,6 +125,18 @@ impl SerializeConfig for StatelessTransactionValidatorConfig {
                  value.",
                 ParamPrivacyInput::Public,
             ),
+            ser_param(
+                "max_bytecode_size",
+                &self.max_bytecode_size,
+                "The maximum bytecode size allowed for a contract.",
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "max_raw_class_size",
+                &self.max_raw_class_size,
+                "The maximum raw class size allowed for a contract.",
+                ParamPrivacyInput::Public,
+            ),
         ]);
         vec![
             members,

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
@@ -35,6 +35,7 @@ pub struct AppState {
     pub stateless_tx_validator: StatelessTransactionValidator,
     pub stateful_tx_validator: Arc<StatefulTransactionValidator>,
     pub state_reader_factory: Arc<dyn StateReaderFactory>,
+    pub gateway_compiler: GatewayCompiler,
     pub mempool_client: SharedMempoolClient,
 }
 
@@ -52,6 +53,7 @@ impl Gateway {
                 config: config.stateful_tx_validator_config.clone(),
             }),
             state_reader_factory,
+            gateway_compiler: GatewayCompiler { config: config.gateway_compiler_config.clone() },
             mempool_client,
         };
         Gateway { config, app_state }
@@ -90,6 +92,7 @@ async fn add_tx(
             app_state.stateless_tx_validator,
             app_state.stateful_tx_validator.as_ref(),
             app_state.state_reader_factory.as_ref(),
+            app_state.gateway_compiler,
             tx,
         )
     })
@@ -110,6 +113,7 @@ fn process_tx(
     stateless_tx_validator: StatelessTransactionValidator,
     stateful_tx_validator: &StatefulTransactionValidator,
     state_reader_factory: &dyn StateReaderFactory,
+    gateway_compiler: GatewayCompiler,
     tx: RPCTransaction,
 ) -> GatewayResult<MempoolInput> {
     // TODO(Arni, 1/5/2024): Perform congestion control.
@@ -120,7 +124,6 @@ fn process_tx(
     // Compile Sierra to Casm.
     let optional_class_info = match &tx {
         RPCTransaction::Declare(declare_tx) => {
-            let gateway_compiler = GatewayCompiler { config: Default::default() };
             Some(gateway_compiler.compile_contract_class(declare_tx)?)
         }
         _ => 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_config::GatewayCompilerConfig;
 use crate::config::{StatefulTransactionValidatorConfig, StatelessTransactionValidatorConfig};
 use crate::gateway::{add_tx, AppState, GatewayCompiler, SharedMempoolClient};
 use crate::state_reader_test_utils::{
@@ -52,6 +53,9 @@ pub fn app_state(
         stateful_tx_validator: Arc::new(StatefulTransactionValidator {
             config: StatefulTransactionValidatorConfig::create_for_testing(),
         }),
+        gateway_compiler: GatewayCompiler {
+            config: GatewayCompilerConfig { max_bytecode_size: 10000, max_raw_class_size: 1000000 },
+        },
         state_reader_factory: Arc::new(state_reader_factory),
         mempool_client,
     }
@@ -111,9 +115,14 @@ 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(
-            GatewayCompiler { config: Default::default() }
-                .compile_contract_class(declare_tx)
-                .unwrap(),
+            GatewayCompiler {
+                config: GatewayCompilerConfig {
+                    max_bytecode_size: 4800,
+                    max_raw_class_size: 111037,
+                },
+            }
+            .compile_contract_class(declare_tx)
+            .unwrap(),
         ),
         _ => None,
     };

diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs
@@ -13,6 +13,7 @@ use test_utils::starknet_api_test_utils::{
 };
 
 use crate::compilation::GatewayCompiler;
+use crate::compilation_config::GatewayCompilerConfig;
 use crate::config::StatefulTransactionValidatorConfig;
 use crate::errors::{StatefulTransactionValidatorError, StatefulTransactionValidatorResult};
 use crate::state_reader_test_utils::{
@@ -82,7 +83,12 @@ fn test_stateful_tx_validator(
     };
     let optional_class_info = match &external_tx {
         RPCTransaction::Declare(declare_tx) => {
-            let gateway_compiler = GatewayCompiler { config: Default::default() };
+            let gateway_compiler = GatewayCompiler {
+                config: GatewayCompilerConfig {
+                    max_bytecode_size: 4800,
+                    max_raw_class_size: 111037,
+                },
+            };
             Some(gateway_compiler.compile_contract_class(declare_tx).unwrap())
         }
         _ => None,