Merge branch 'abk/run-647-larger-wasms' into 'master'

RUN-647: Allow large Wasms with smaller code sections

This change allows the runtime to accept larger Wasm modules. In order
to keep compilation times fast, we still impose a 10MB limit on the size
of the code sections. 

See merge request dfinity-lab/public/ic!12528

Author: adambratschikaye

rs/embedders/src/wasm_utils.rs

@@ -175,6 +175,7 @@ impl SystemApiFunc {
     }
 }
 
+#[derive(Debug)]
 pub struct InstrumentationOutput {
     /// All exported methods that are relevant to the IC.
     /// Methods relevant to the IC are:

rs/embedders/src/wasm_utils/decoding.rs

@@ -3,7 +3,7 @@ use std::io::Read;
 use std::sync::Arc;
 
 /// Maximum size of a WebAssembly module.
-pub const MAX_WASM_MODULE_SIZE_BYTES: usize = 10 * 1024 * 1024;
+pub const MAX_WASM_MODULE_SIZE_BYTES: usize = 30 * 1024 * 1024;
 
 fn make_module_too_large_error() -> WasmValidationError {
     WasmValidationError::DecodingError(format!(

rs/embedders/src/wasm_utils/validation.rs

@@ -37,6 +37,7 @@ pub const RESERVED_SYMBOLS: [&str; 6] = [
 
 const WASM_FUNCTION_COMPLEXITY_LIMIT: usize = 15_000;
 const WASM_FUNCTION_SIZE_LIMIT: usize = 1_000_000;
+const MAX_CODE_SECTION_SIZE_IN_BYTES: u32 = 10 * 1024 * 1024;
 
 // Represents the expected function signature for any System APIs the Internet
 // Computer provides or any special exported user functions.
@@ -1148,6 +1149,30 @@ fn can_compile(wasm: &BinaryEncodedWasm) -> Result<(), WasmValidationError> {
     })
 }
 
+fn check_code_section_size(wasm: &BinaryEncodedWasm) -> Result<(), WasmValidationError> {
+    let parser = wasmparser::Parser::new(0);
+    let payloads = parser.parse_all(wasm.as_slice());
+    for payload in payloads {
+        if let wasmparser::Payload::CodeSectionStart {
+            count: _,
+            range: _,
+            size,
+        } = payload.map_err(|e| {
+            WasmValidationError::DecodingError(format!("Error finding code section: {}", e))
+        })? {
+            if size > MAX_CODE_SECTION_SIZE_IN_BYTES {
+                return Err(WasmValidationError::CodeSectionTooLarge {
+                    size,
+                    allowed: MAX_CODE_SECTION_SIZE_IN_BYTES,
+                });
+            } else {
+                return Ok(());
+            }
+        }
+    }
+    Ok(())
+}
+
 /// Validates a Wasm binary against the requirements of the interface spec
 /// defined in https://sdk.dfinity.org/docs/interface-spec/index.html.
 ///
@@ -1167,6 +1192,7 @@ pub(super) fn validate_wasm_binary<'a>(
     wasm: &'a BinaryEncodedWasm,
     config: &EmbeddersConfig,
 ) -> Result<(WasmValidationDetails, Module<'a>), WasmValidationError> {
+    check_code_section_size(wasm)?;
     can_compile(wasm)?;
     let module = Module::parse(wasm.as_slice(), false)
         .map_err(|err| WasmValidationError::DecodingError(format!("{}", err)))?;

rs/embedders/tests/compressed/zeros.gz

@@ -17,6 +17,9 @@ use ic_replicated_state::canister_state::execution_state::{
 };
 use ic_types::{NumBytes, NumInstructions};
 use maplit::btreemap;
+use wasmtime_environ::WASM_PAGE_SIZE;
+
+const KB: u32 = 1024;
 
 fn wat2wasm(wat: &str) -> Result<BinaryEncodedWasm, wat::Error> {
     wat::parse_str(wat).map(BinaryEncodedWasm::new)
@@ -958,3 +961,144 @@ fn complex_function_rejected() {
         })
     )
 }
+
+/// Creates a was with roughly the given sizes for the code and data sections
+/// (may be off by a few bytes).
+fn wasm_with_fixed_sizes(code_section_size: u32, data_section_size: u32) -> BinaryEncodedWasm {
+    // Initial memory needs to be large enough to fit the data
+    let memory_size = data_section_size / WASM_PAGE_SIZE + 1;
+    let mut wat = "(module (func".to_string();
+    // Each (block) is 3 bytes: 2 bytes for "block" and 1 for "end"
+    for _ in 0..code_section_size / 3 {
+        wat.push_str("(block)");
+    }
+    wat.push(')');
+    wat.push_str(&format!("(memory {})", memory_size));
+    wat.push_str(&format!(
+        "(data (i32.const 0) \"{}\")",
+        "a".repeat(data_section_size as usize),
+    ));
+    wat.push(')');
+    wat2wasm(&wat).unwrap()
+}
+
+#[test]
+fn large_code_section_rejected() {
+    let wasm = wasm_with_fixed_sizes(10 * KB * KB + 10, 0);
+    let embedder = WasmtimeEmbedder::new(EmbeddersConfig::default(), no_op_logger());
+    let result = validate_and_instrument_for_testing(&embedder, &wasm);
+    assert_matches!(
+        result,
+        Err(HypervisorError::InvalidWasm(
+            WasmValidationError::CodeSectionTooLarge { .. },
+        ))
+    )
+}
+
+#[test]
+fn large_wasm_with_small_code_accepted() {
+    let wasm = wasm_with_fixed_sizes(KB, 20 * KB * KB);
+    let embedder = WasmtimeEmbedder::new(EmbeddersConfig::default(), no_op_logger());
+    let result = validate_and_instrument_for_testing(&embedder, &wasm);
+    assert_matches!(result, Ok(_))
+}
+
+/// We are trusting the code section size reported in the header when
+/// determining if the code section is too long.  A Wasm which has been
+/// manipulated to report an incorrectly small size in the header should be
+/// rejected should later be rejected when we try to validate it with Wasmtime.
+#[test]
+fn incorrect_wasm_code_size_is_invalid() {
+    use wasmparser::{Parser, Payload};
+
+    let wasm = wasm_with_fixed_sizes(10 * KB * KB + 10, 0);
+
+    let parser = Parser::new(0);
+    let payloads = parser.parse_all(wasm.as_slice());
+    let mut manipulated_wasm = vec![];
+    for payload in payloads {
+        if let Payload::CodeSectionStart { range, .. } = payload.unwrap() {
+            // The section header contains the byte 10 as the section id
+            // followed by a variable length encoded u32 for the size.
+            //
+            // Note that the `size` field doesn't include the encoding of the
+            // function count (which is 1), so it differs from the length of the
+            // `range` (which does include the count encoding).
+            //
+            // The code section should have size 0xa0000f, which is 0x8f808005
+            // as a variable-length u32.
+            assert_eq!(range.end - range.start, 0xa0000f);
+            assert_eq!(
+                wasm.as_slice()[range.start - 5..range.start],
+                [0xa /*Code section id*/, 0x8f, 0x80, 0x80, 0x05]
+            );
+            // Copy everything up to and including the code section id.
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[..range.start - 4]);
+            // Push 0x7f = 127 for the code section size.
+            manipulated_wasm.push(0x7f);
+            // Copy everything after the code section size unchanged.
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[range.start..]);
+            break;
+        }
+    }
+
+    let manipulated_wasm = BinaryEncodedWasm::new(manipulated_wasm);
+    let embedder = WasmtimeEmbedder::new(EmbeddersConfig::default(), no_op_logger());
+    let result = validate_and_instrument_for_testing(&embedder, &manipulated_wasm);
+    assert_matches!(
+        result,
+        Err(HypervisorError::InvalidWasm(
+            WasmValidationError::WasmtimeValidation(_),
+        ))
+    )
+}
+
+/// We're assuming there is at most one code section in the Wasm. The spec
+/// doesn't allow multiple code sections, so if there are multiple code sections
+/// the module should fail validation.
+#[test]
+fn wasm_with_multiple_code_sections_is_invalid() {
+    use wasmparser::{Parser, Payload};
+
+    let wasm = wasm_with_fixed_sizes(10, 0);
+
+    let parser = Parser::new(0);
+    let payloads = parser.parse_all(wasm.as_slice());
+    let mut manipulated_wasm = vec![];
+    for payload in payloads {
+        if let Payload::CodeSectionStart { range, .. } = payload.unwrap() {
+            // The section header contains the byte 10 as the section id
+            // followed by a variable length encoded u32 for the size.
+            //
+            // Note that the `size` field doesn't include the encoding of the
+            // function count (which is 1), so it differs from the length of the
+            // `range` (which does include the count encoding).
+            //
+            // The code section should have size 0xa, which is unchanged as a
+            // variable-length u32.
+            assert_eq!(range.end - range.start, 0xd);
+            assert_eq!(
+                wasm.as_slice()[range.start - 2..range.start],
+                [0xa /*Code section id*/, 0x0d]
+            );
+            // Copy everything before the code section
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[..range.start - 2]);
+            // Copy the code section twice
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[range.start - 2..range.end]);
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[range.start - 2..range.end]);
+            // Copy everything after the code section size unchanged.
+            manipulated_wasm.extend_from_slice(&wasm.as_slice()[range.start..]);
+            break;
+        }
+    }
+
+    let manipulated_wasm = BinaryEncodedWasm::new(manipulated_wasm);
+    let embedder = WasmtimeEmbedder::new(EmbeddersConfig::default(), no_op_logger());
+    let result = validate_and_instrument_for_testing(&embedder, &manipulated_wasm);
+    assert_matches!(
+        result,
+        Err(HypervisorError::InvalidWasm(
+            WasmValidationError::WasmtimeValidation(_),
+        ))
+    )
+}

rs/embedders/tests/validation.rs

@@ -56,6 +56,8 @@ pub enum WasmValidationError {
         size: usize,
         allowed: usize,
     },
+    /// The code section is too large.
+    CodeSectionTooLarge { size: u32, allowed: u32 },
 }
 
 impl std::fmt::Display for WasmValidationError {
@@ -118,6 +120,11 @@ impl std::fmt::Display for WasmValidationError {
                 "Wasm module contains a function at index {} of size {} that exceeds the maximum allowed size of {}",
                 index, size, allowed,
             ),
+            Self::CodeSectionTooLarge{size, allowed} => write!(
+                f,
+                "Wasm model code section size of {} exceeds the maximum allowed size of {}",
+                size, allowed,
+            ),
         }
     }
 }

rs/types/wasm_types/src/errors.rs

@@ -19,6 +19,13 @@ const WASM_HASH_LENGTH: usize = 32;
 #[derive(Clone)]
 pub struct BinaryEncodedWasm(Arc<Vec<u8>>);
 
+impl std::fmt::Debug for BinaryEncodedWasm {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Ignore the actual binary contents when debug formatting.
+        f.debug_tuple("BinaryEncodedWasm").finish()
+    }
+}
+
 impl BinaryEncodedWasm {
     pub fn new(wasm: Vec<u8>) -> Self {
         Self::new_shared(Arc::new(wasm))