Add a custom section hook to ModuleEnvironment

This commit adds a hook to the `ModuleEnvironment` trait to learn when a
custom section in a wasm file is read. This hook can in theory be used
to parse and handle custom sections as they appear in the wasm file
without having to re-iterate over the wasm file after cranelift has
already parsed the wasm file.

The `translate_module` function is now less strict in that it doesn't
require sections to be in a particular order, but it's figured that the
wasm file is already validated elsewhere to verify the section order.

cranelift/wasm/src/environ/spec.rs

@@ -467,4 +467,14 @@ pub trait ModuleEnvironment<'data> {
         offset: usize,
         data: &'data [u8],
     ) -> WasmResult<()>;
+
+    /// Indicates that a custom section has been found in the wasm file
+    fn custom_section(
+        &mut self,
+        name: &'data str,
+        data: &'data [u8],
+    ) -> WasmResult<()> {
+        drop((name, data));
+        Ok(())
+    }
 }

cranelift/wasm/src/module_translator.rs

@@ -18,134 +18,79 @@ pub fn translate_module<'data>(
     let _tt = timing::wasm_translate_module();
     let mut reader = ModuleReader::new(data)?;
 
-    reader.skip_custom_sections()?;
-    if reader.eof() {
-        return Ok(());
-    }
-    let mut section = reader.read()?;
-
-    if let SectionCode::Type = section.code {
+    while !reader.eof() {
+        let section = reader.read()?;
+        match section.code {
+            SectionCode::Type => {
                 let types = section.get_type_section_reader()?;
                 parse_type_section(types, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Import = section.code {
+            SectionCode::Import => {
                 let imports = section.get_import_section_reader()?;
                 parse_import_section(imports, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Function = section.code {
+            SectionCode::Function => {
                 let functions = section.get_function_section_reader()?;
                 parse_function_section(functions, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Table = section.code {
+            SectionCode::Table => {
                 let tables = section.get_table_section_reader()?;
                 parse_table_section(tables, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Memory = section.code {
+            SectionCode::Memory => {
                 let memories = section.get_memory_section_reader()?;
                 parse_memory_section(memories, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Global = section.code {
+            SectionCode::Global => {
                 let globals = section.get_global_section_reader()?;
                 parse_global_section(globals, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Export = section.code {
+            SectionCode::Export => {
                 let exports = section.get_export_section_reader()?;
                 parse_export_section(exports, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Start = section.code {
+            SectionCode::Start => {
                 let start = section.get_start_section_content()?;
                 parse_start_section(start, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Element = section.code {
+            SectionCode::Element => {
                 let elements = section.get_element_section_reader()?;
                 parse_element_section(elements, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Code = section.code {
+            SectionCode::Code => {
                 let code = section.get_code_section_reader()?;
                 parse_code_section(code, environ)?;
-
-        reader.skip_custom_sections()?;
-        if reader.eof() {
-            return Ok(());
-        }
-        section = reader.read()?;
             }
 
-    if let SectionCode::Data = section.code {
+            SectionCode::Data => {
                 let data = section.get_data_section_reader()?;
                 parse_data_section(data, environ)?;
             }
 
-    reader.skip_custom_sections()?;
-    if !reader.eof() {
+            SectionCode::DataCount => {
                 return Err(WasmError::InvalidWebAssembly {
-            message: "sections must occur at most once and in the prescribed order",
+                    message: "don't know how to handle the data count section yet",
                     offset: reader.current_position(),
                 });
             }
 
+            SectionCode::Custom { name, kind: _ } => {
+                let mut reader = section.get_binary_reader();
+                let len = reader.bytes_remaining();
+                let payload = reader.read_bytes(len)?;
+                environ.custom_section(name, payload)?;
+            }
+        }
+    }
+
     Ok(())
 }