feat(core,pallets): add data segments amount limit checks (#3938)

core/src/code/errors.rs

@@ -92,6 +92,14 @@ pub enum DataSectionError {
     /// Data segment end address is out of static memory.
     #[display(fmt = "Data segment {_0:#x}..={_1:#x} is out of static memory 0x0..{_2:#x}")]
     EndAddressOutOfStaticMemory(u32, u32, u64),
+    /// Data segment amount exceeds the limit.
+    #[display(fmt = "Data segment amount limit exceeded: limit={limit}, actual={actual}")]
+    DataSegmentsAmountLimit {
+        /// Limit of data segments.
+        limit: u32,
+        /// Actual amount of data segments.
+        actual: u32,
+    },
 }
 
 /// Export error in WASM module.

core/src/code/mod.rs

@@ -62,6 +62,8 @@ pub struct TryNewCodeConfig {
     pub version: u32,
     /// Stack height limit
     pub stack_height: Option<u32>,
+    /// Limit of data section amount
+    pub data_segments_amount_limit: Option<u32>,
     /// Export `STACK_HEIGHT_EXPORT_NAME` global
     pub export_stack_height: bool,
     /// Check exports (wasm contains init or handle exports)
@@ -85,6 +87,7 @@ impl Default for TryNewCodeConfig {
         Self {
             version: 1,
             stack_height: None,
+            data_segments_amount_limit: None,
             export_stack_height: false,
             check_exports: true,
             check_imports: true,
@@ -135,7 +138,12 @@ impl Code {
 
         // Not changing steps
         if config.check_data_section {
-            utils::check_data_section(&module, static_pages, stack_end)?;
+            utils::check_data_section(
+                &module,
+                static_pages,
+                stack_end,
+                config.data_segments_amount_limit,
+            )?;
         }
         if config.check_mut_global_exports {
             utils::check_mut_global_exports(&module)?;
@@ -254,6 +262,7 @@ impl Code {
         version: u32,
         get_gas_rules: GetRulesFn,
         stack_height: Option<u32>,
+        data_segments_amount_limit: Option<u32>,
     ) -> Result<Self, CodeError>
     where
         R: Rules,
@@ -265,6 +274,7 @@ impl Code {
             TryNewCodeConfig {
                 version,
                 stack_height,
+                data_segments_amount_limit,
                 ..Default::default()
             },
         )
@@ -412,15 +422,24 @@ mod tests {
         };
     }
 
-    fn try_new_code_from_wat(wat: &str, stack_height: Option<u32>) -> Result<Code, CodeError> {
+    fn try_new_code_from_wat_with_params(
+        wat: &str,
+        stack_height: Option<u32>,
+        data_segments_amount_limit: Option<u32>,
+    ) -> Result<Code, CodeError> {
         Code::try_new(
             wat2wasm(wat),
             1,
             |_| CustomConstantCostRules::default(),
             stack_height,
+            data_segments_amount_limit,
         )
     }
 
+    fn try_new_code_from_wat(wat: &str, stack_height: Option<u32>) -> Result<Code, CodeError> {
+        try_new_code_from_wat_with_params(wat, stack_height, None)
+    }
+
     #[test]
     fn reject_unknown_exports() {
         let wat = r#"
@@ -703,6 +722,7 @@ mod tests {
             1,
             |_| CustomConstantCostRules::default(),
             None,
+            None,
         );
 
         assert_code_err!(res, CodeError::Validation(_));
@@ -747,4 +767,35 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn data_segments_amount_limit() {
+        const DATA_SEGMENTS_AMOUNT_LIMIT: u32 = 1024;
+
+        let segment = r#"(data (i32.const 0x0) "gear")"#;
+
+        let wat = format!(
+            r#"
+            (module
+                (import "env" "memory" (memory 1))
+                (func $init)
+                (export "init" (func $init))
+                {}
+            )
+        "#,
+            segment.repeat(1025)
+        );
+
+        assert_code_err!(
+            try_new_code_from_wat_with_params(
+                wat.as_str(),
+                None,
+                DATA_SEGMENTS_AMOUNT_LIMIT.into()
+            ),
+            CodeError::DataSection(DataSectionError::DataSegmentsAmountLimit {
+                limit: DATA_SEGMENTS_AMOUNT_LIMIT,
+                actual: 1025
+            })
+        );
+    }
 }

core/src/code/utils.rs

@@ -264,12 +264,24 @@ pub fn check_data_section(
     module: &Module,
     static_pages: WasmPagesAmount,
     stack_end: Option<WasmPage>,
+    data_section_amount_limit: Option<u32>,
 ) -> Result<(), CodeError> {
     let Some(data_section) = module.data_section() else {
         // No data section - nothing to check.
         return Ok(());
     };
 
+    // Check that data segments amount does not exceed the limit.
+    if let Some(data_segments_amount_limit) = data_section_amount_limit {
+        let number_of_data_segments = data_section.entries().len() as u32;
+        if number_of_data_segments > data_segments_amount_limit {
+            Err(DataSectionError::DataSegmentsAmountLimit {
+                limit: data_segments_amount_limit,
+                actual: number_of_data_segments,
+            })?;
+        }
+    }
+
     for data_segment in data_section.entries() {
         let data_segment_offset = data_segment
             .offset()

core/src/program.rs

@@ -171,7 +171,14 @@ mod tests {
 
         let binary: Vec<u8> = parse_wat(wat);
 
-        let code = Code::try_new(binary, 1, |_| CustomConstantCostRules::default(), None).unwrap();
+        let code = Code::try_new(
+            binary,
+            1,
+            |_| CustomConstantCostRules::default(),
+            None,
+            None,
+        )
+        .unwrap();
         let (code, _) = code.into_parts();
         let program = Program::new(ProgramId::from(1), Default::default(), code);
 

gsdk/src/metadata/generated.rs

@@ -2431,6 +2431,7 @@ pub mod runtime_types {
                     pub call_depth: ::core::primitive::u32,
                     pub payload_len: ::core::primitive::u32,
                     pub code_len: ::core::primitive::u32,
+                    pub data_segments_amount: ::core::primitive::u32,
                 }
                 #[derive(Debug, crate::gp::Decode, crate::gp::DecodeAsType, crate::gp::Encode)]
                 pub struct MemoryWeights {

gtest/src/manager.rs

@@ -1156,6 +1156,7 @@ impl JournalHandler for ExtManager {
                         schedule.instruction_weights.version,
                         |module| schedule.rules(module),
                         schedule.limits.stack_height,
+                        schedule.limits.data_segments_amount.into(),
                     )
                     .expect("Program can't be constructed with provided code");
 

gtest/src/program.rs

@@ -356,6 +356,7 @@ impl ProgramBuilder {
             schedule.instruction_weights.version,
             |module| schedule.rules(module),
             schedule.limits.stack_height,
+            schedule.limits.data_segments_amount.into(),
         )
         .expect("Failed to create Program from code");
 

pallets/gear/src/benchmarking/utils.rs

@@ -93,6 +93,7 @@ where
                 schedule.instruction_weights.version,
                 |module| schedule.rules(module),
                 schedule.limits.stack_height,
+                schedule.limits.data_segments_amount.into(),
             )
             .map_err(|_| "Code failed to load")?;
 

pallets/gear/src/lib.rs

@@ -1135,6 +1135,7 @@ pub mod pallet {
                 schedule.instruction_weights.version,
                 |module| schedule.rules(module),
                 schedule.limits.stack_height,
+                schedule.limits.data_segments_amount.into(),
             )?;
 
             let code_and_id = CodeAndId::from_parts_unchecked(code, code_id);
@@ -1158,6 +1159,7 @@ pub mod pallet {
                 schedule.instruction_weights.version,
                 |module| schedule.rules(module),
                 schedule.limits.stack_height,
+                schedule.limits.data_segments_amount.into(),
             )
             .map_err(|e| {
                 log::debug!("Code checking or instrumentation failed: {e}");

pallets/gear/src/schedule.rs

@@ -68,6 +68,11 @@ pub const STACK_HEIGHT_LIMIT: u32 = 36_743;
 #[cfg(feature = "fuzz")]
 pub const FUZZER_STACK_HEIGHT_LIMIT: u32 = 65_000;
 
+/// Maximum number of data segments in a wasm module.
+/// It has been determined that the maximum number of data segments in a wasm module
+/// does not exceed 1024 by a large margin.
+pub const DATA_SEGMENTS_AMOUNT_LIMIT: u32 = 1024;
+
 /// Definition of the cost schedule and other parameterization for the wasm vm.
 ///
 /// Its [`Default`] implementation is the designated way to initialize this type. It uses
@@ -202,6 +207,9 @@ pub struct Limits {
     /// version of the code. Therefore `instantiate_with_code` can fail even when supplying
     /// a wasm binary below this maximum size.
     pub code_len: u32,
+
+    /// The maximum number of wasm data segments allowed for a program.
+    pub data_segments_amount: u32,
 }
 
 impl Limits {
@@ -769,6 +777,7 @@ impl Default for Limits {
             stack_height: Some(STACK_HEIGHT_LIMIT),
             #[cfg(feature = "fuzz")]
             stack_height: Some(FUZZER_STACK_HEIGHT_LIMIT),
+            data_segments_amount: DATA_SEGMENTS_AMOUNT_LIMIT,
             globals: 256,
             locals: 1024,
             parameters: 128,

pallets/gear/src/tests.rs

@@ -224,6 +224,7 @@ fn state_rpc_calls_trigger_reinstrumentation() {
             0, // invalid version
             |module| schedule.rules(module),
             schedule.limits.stack_height,
+            schedule.limits.data_segments_amount.into(),
         )
         .expect("Failed to create dummy code");
 
@@ -4970,6 +4971,7 @@ fn test_code_submission_pass() {
             schedule.instruction_weights.version,
             |module| schedule.rules(module),
             schedule.limits.stack_height,
+            schedule.limits.data_segments_amount.into(),
         )
         .expect("Error creating Code");
         assert_eq!(saved_code.unwrap().code(), code.code());
@@ -6465,6 +6467,7 @@ fn test_create_program_works() {
             schedule.instruction_weights.version,
             |module| schedule.rules(module),
             schedule.limits.stack_height,
+            schedule.limits.data_segments_amount.into(),
         )
         .expect("Code failed to load");
 
@@ -10004,6 +10007,7 @@ fn test_mad_big_prog_instrumentation() {
             schedule.instruction_weights.version,
             |module| schedule.rules(module),
             schedule.limits.stack_height,
+            schedule.limits.data_segments_amount.into(),
         );
         // In any case of the defined weights on the platform, instrumentation of the valid
         // huge wasm mustn't fail
@@ -13069,7 +13073,8 @@ fn wrong_entry_type() {
                 ProgramCodeKind::Custom(wat).to_bytes(),
                 1,
                 |_| CustomConstantCostRules::default(),
-                None
+                None,
+                None,
             ),
             Err(CodeError::Export(ExportError::InvalidExportFnSignature(0)))
         ));

utils/calc-stack-height/src/main.rs

@@ -90,6 +90,7 @@ fn main() -> anyhow::Result<()> {
             schedule.instruction_weights.version,
             |module| schedule.rules(module),
             Some(mid),
+            schedule.limits.data_segments_amount.into(),
         )
         .map_err(|e| anyhow::anyhow!("{e}"))?;
 

utils/wasm-builder/src/code_validator.rs

@@ -319,7 +319,13 @@ impl CodeValidator {
     /// Validates wasm code in the same way as
     /// `pallet_gear::pallet::Pallet::upload_program(...)`.
     pub fn validate_program(self) -> anyhow::Result<()> {
-        match Code::try_new(self.code, 1, |_| CustomConstantCostRules::default(), None) {
+        match Code::try_new(
+            self.code,
+            1,
+            |_| CustomConstantCostRules::default(),
+            None,
+            None,
+        ) {
             Err(code_error) => Err(CodeErrorWithContext::from((self.module, code_error)))?,
             _ => Ok(()),
         }

utils/wasm-gen/src/tests.rs

@@ -59,7 +59,7 @@ proptest! {
         let original_code = generate_gear_program_code(&mut u, configs_bundle)
             .expect("failed generating wasm");
 
-        let code_res = Code::try_new(original_code, 1, |_| CustomConstantCostRules::default(), None);
+        let code_res = Code::try_new(original_code, 1, |_| CustomConstantCostRules::default(), None, None);
         assert!(code_res.is_ok());
     }
 
@@ -767,8 +767,14 @@ fn execute_wasm_with_custom_configs(
 
     let code =
         generate_gear_program_code(unstructured, gear_config).expect("failed wasm generation");
-    let code = Code::try_new(code, 1, |_| CustomConstantCostRules::new(0, 0, 0), None)
-        .expect("Failed to create Code");
+    let code = Code::try_new(
+        code,
+        1,
+        |_| CustomConstantCostRules::new(0, 0, 0),
+        None,
+        None,
+    )
+    .expect("Failed to create Code");
 
     let code_id = CodeId::generate(code.original_code());
     let program_id = ProgramId::generate_from_user(code_id, b"");

utils/wasm-instrument/src/gas_metering/schedule.rs

@@ -76,6 +76,7 @@ pub struct Limits {
     pub call_depth: u32,
     pub payload_len: u32,
     pub code_len: u32,
+    pub data_segments_amount: u32,
 }
 
 impl Default for Limits {
@@ -92,6 +93,7 @@ impl Default for Limits {
             call_depth: 32,
             payload_len: 8388608,
             code_len: 524288,
+            data_segments_amount: 1024,
         }
     }
 }