feat: RUN-1036: Validate initial wasm memory size for Wasm64 (#1534)

In this PR we add an additional validation step for the initial size for
Wasm64 memories. The size of initial wasm heap memory should not be more
than the protocol allows.

Author: alexandru-uta

rs/embedders/src/wasm_utils/validation.rs

@@ -27,6 +27,8 @@ use crate::{
 };
 use wasmparser::{CompositeInnerType, ExternalKind, FuncType, Operator, TypeRef, ValType};
 
+const WASM_PAGE_SIZE: u32 = wasmtime_environ::Memory::DEFAULT_PAGE_SIZE;
+
 /// Symbols that are reserved and cannot be exported by canisters.
 #[doc(hidden)] // pub for usage in tests
 pub const RESERVED_SYMBOLS: [&str; 6] = [
@@ -1048,6 +1050,29 @@ fn validate_function_section(
     Ok(())
 }
 
+// Checks that the initial size of the wasm (heap) memory is not larger than
+// the allowed maximum size. This is only needed for Wasm64, because in Wasm32 this
+// is checked by Wasmtime.
+fn validate_initial_wasm_memory_size(
+    module: &Module,
+    max_wasm_memory_size_in_bytes: NumBytes,
+) -> Result<(), WasmValidationError> {
+    for memory in &module.memories {
+        if memory.memory64 {
+            let declared_size_in_wasm_pages = memory.initial;
+            let allowed_size_in_wasm_pages =
+                max_wasm_memory_size_in_bytes.get() / WASM_PAGE_SIZE as u64;
+            if declared_size_in_wasm_pages > allowed_size_in_wasm_pages {
+                return Err(WasmValidationError::InitialWasm64MemoryTooLarge {
+                    declared_size: declared_size_in_wasm_pages,
+                    allowed_size: allowed_size_in_wasm_pages,
+                });
+            }
+        }
+    }
+    Ok(())
+}
+
 // Extracts the name of the custom section.
 // Possible options:
 //      * icp:public <name>
@@ -1511,6 +1536,7 @@ pub(super) fn validate_wasm_binary<'a>(
     validate_data_section(&module)?;
     validate_global_section(&module, config.max_globals)?;
     validate_function_section(&module, config.max_functions)?;
+    validate_initial_wasm_memory_size(&module, config.max_wasm_memory_size)?;
     let (largest_function_instruction_count, max_complexity) = validate_code_section(&module)?;
     let wasm_metadata = validate_custom_section(&module, config)?;
     Ok((

rs/embedders/tests/validation.rs

@@ -1132,3 +1132,36 @@ fn wasm_with_multiple_code_sections_is_invalid() {
         ))
     )
 }
+
+#[test]
+fn test_wasm64_initial_wasm_memory_size_validation() {
+    use crate::WasmValidationError::InitialWasm64MemoryTooLarge;
+    use ic_config::embedders::FeatureFlags;
+    use ic_config::flag_status::FlagStatus;
+
+    let embedders_config = EmbeddersConfig {
+        feature_flags: FeatureFlags {
+            wasm64: FlagStatus::Enabled,
+            ..Default::default()
+        },
+        ..Default::default()
+    };
+    let allowed_wasm_memory_size_in_pages =
+        embedders_config.max_wasm_memory_size.get() / WASM_PAGE_SIZE as u64;
+    let declared_wasm_memory_size_in_pages = allowed_wasm_memory_size_in_pages + 10;
+    let wasm = wat2wasm(&format!(
+        r#"(module
+            (memory i64 {} {})
+        )"#,
+        declared_wasm_memory_size_in_pages, declared_wasm_memory_size_in_pages
+    ))
+    .unwrap();
+
+    assert_eq!(
+        validate_wasm_binary(&wasm, &embedders_config),
+        Err(InitialWasm64MemoryTooLarge {
+            declared_size: declared_wasm_memory_size_in_pages,
+            allowed_size: allowed_wasm_memory_size_in_pages
+        })
+    );
+}

rs/embedders/tests/wasmtime_embedder.rs

@@ -21,6 +21,8 @@ use ic_types::{
     Cycles, NumBytes, NumInstructions,
 };
 
+const WASM_PAGE_SIZE: u32 = wasmtime_environ::Memory::DEFAULT_PAGE_SIZE;
+
 #[cfg(target_os = "linux")]
 use ic_types::PrincipalId;
 
@@ -2790,23 +2792,52 @@ fn wasm64_cycles_burn128() {
 
 #[test]
 fn large_wasm64_memory_allocation_test() {
-    let wat = r#"
-    (module
-        (func $test (export "canister_update test"))
-        (memory i64 0 16777216)
-    )"#;
+    // This test checks if maximum memory size
+    // is capped to the maximum allowed memory size in 64 bit mode.
 
     let mut config = ic_config::embedders::Config::default();
     config.feature_flags.wasm64 = FlagStatus::Enabled;
+    let max_heap_size_in_pages = config.max_wasm_memory_size.get() / WASM_PAGE_SIZE as u64;
+    let wat = format!(
+        r#"
+    (module
+        (import "ic0" "msg_reply" (func $msg_reply))
+        (import "ic0" "msg_reply_data_append" (func $msg_reply_data_append (param $src i64) (param $size i64)))
+        (func $test (export "canister_update test")
+            ;; store the result of memory.grow at heap address 0
+            (i64.store (i64.const 0) (memory.grow (i64.const 1)))
+            ;; return the result of memory.grow
+            (call $msg_reply_data_append (i64.const 0) (i64.const 1))
+            (call $msg_reply)
+        )
+        ;; declare a memory with initial size max_heap and another max large value
+        (memory i64 {} {})
+    )"#,
+        max_heap_size_in_pages,
+        max_heap_size_in_pages * 100
+    );
+
     let mut instance = WasmtimeInstanceBuilder::new()
         .with_config(config)
-        .with_wat(wat)
+        .with_api_type(ic_system_api::ApiType::update(
+            UNIX_EPOCH,
+            vec![],
+            Cycles::zero(),
+            user_test_id(24).get(),
+            call_context_test_id(13),
+        ))
+        .with_wat(&wat)
         .build();
 
-    match instance.run(FuncRef::Method(WasmMethod::Update("test".to_string()))) {
-        Ok(_) => {}
-        Err(e) => panic!("Unexpected error: {:?}", e),
-    }
+    let result = instance.run(FuncRef::Method(WasmMethod::Update("test".to_string())));
+    let wasm_res = instance
+        .store_data_mut()
+        .system_api_mut()
+        .unwrap()
+        .take_execution_result(result.as_ref().err());
+
+    // The reply is actually the encoding of -1 (the memory grow failed).
+    assert_eq!(wasm_res, Ok(Some(WasmResult::Reply(vec![255]))));
 }
 
 #[test]

rs/types/wasm_types/src/errors.rs

@@ -99,6 +99,11 @@ pub enum WasmValidationError {
     CodeSectionTooLarge { size: u32, allowed: u32 },
     /// The total module size is too large.
     ModuleTooLarge { size: u64, allowed: u64 },
+    // The initial Wasm64 heap memory is too large.
+    InitialWasm64MemoryTooLarge {
+        declared_size: u64,
+        allowed_size: u64,
+    },
 }
 
 impl std::fmt::Display for WasmValidationError {
@@ -204,6 +209,14 @@ impl std::fmt::Display for WasmValidationError {
                 "Wasm module size of {size} exceeds the maximum \
                     allowed size of {allowed}.",
             ),
+            Self::InitialWasm64MemoryTooLarge {
+                declared_size,
+                allowed_size,
+            } => write!(
+                f,
+                "Wasm module declares an initial Wasm64 heap memory size of {declared_size} pages \
+                    which exceeds the maximum allowed size of {allowed_size} pages.",
+            ),
         }
     }
 }
@@ -270,6 +283,10 @@ impl AsErrorHelp for WasmValidationError {
                     .to_string(),
                 doc_link: doc_ref("wasm-module-too-large"),
             },
+            WasmValidationError::InitialWasm64MemoryTooLarge { .. } => ErrorHelp::UserError {
+                suggestion: "Try reducing the initial Wasm64 heap memory size.".to_string(),
+                doc_link: doc_ref("wasm-module-initial-wasm64-memory-too-large"),
+            },
         }
     }
 }