fixup! Capture and surface parsed custom sections in modules.

crates/wasmi/src/module/builder.rs

@@ -32,6 +32,7 @@ use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, vec::Vec};
 pub struct ModuleBuilder {
     pub header: ModuleHeader,
     pub data_segments: Vec<DataSegment>,
+    pub custom_sections: Vec<CustomSection>,
 }
 
 /// A builder for a WebAssembly [`Module`] header.
@@ -50,7 +51,6 @@ pub struct ModuleHeaderBuilder {
     pub compiled_funcs: Vec<CompiledFunc>,
     pub compiled_funcs_idx: BTreeMap<CompiledFunc, FuncIdx>,
     pub element_segments: Vec<ElementSegment>,
-    pub(super) custom_sections: Vec<CustomSection>,
 }
 
 impl ModuleHeaderBuilder {
@@ -70,7 +70,6 @@ impl ModuleHeaderBuilder {
             compiled_funcs: Vec::new(),
             compiled_funcs_idx: BTreeMap::new(),
             element_segments: Vec::new(),
-            custom_sections: Vec::new(),
         }
     }
 
@@ -91,7 +90,6 @@ impl ModuleHeaderBuilder {
                 compiled_funcs: self.compiled_funcs.into(),
                 compiled_funcs_idx: self.compiled_funcs_idx,
                 element_segments: self.element_segments.into(),
-                custom_sections: self.custom_sections.into()
             }),
         }
     }
@@ -138,6 +136,7 @@ impl ModuleBuilder {
         Self {
             header,
             data_segments: Vec::new(),
+            custom_sections: Vec::new(),
         }
     }
 }
@@ -370,12 +369,6 @@ impl ModuleHeaderBuilder {
         self.element_segments = elements.into_iter().collect::<Result<Vec<_>, _>>()?;
         Ok(())
     }
-
-    pub fn push_custom_section(&mut self, name: &str, data: &[u8]) {
-        let name: Box<str> = name.into();
-        let data: Box<[u8]> = data.into();
-        self.custom_sections.push(CustomSection { name, data })
-    }
 }
 
 impl ModuleBuilder {
@@ -400,12 +393,19 @@ impl ModuleBuilder {
         Ok(())
     }
 
+    pub fn push_custom_section(&mut self, name: &str, data: &[u8]) {
+        let name: Box<str> = name.into();
+        let data: Box<[u8]> = data.into();
+        self.custom_sections.push(CustomSection { name, data })
+    }
+
     /// Finishes construction of the WebAssembly [`Module`].
     pub fn finish(self, engine: &Engine) -> Module {
         Module {
             engine: engine.clone(),
             header: self.header,
             data_segments: self.data_segments.into(),
+            custom_sections: self.custom_sections.into(),
         }
     }
 }

crates/wasmi/src/module/mod.rs

@@ -56,6 +56,7 @@ pub struct Module {
     engine: Engine,
     header: ModuleHeader,
     data_segments: Box<[DataSegment]>,
+    custom_sections: Box<[CustomSection]>,
 }
 
 /// A parsed and validated WebAssembly module header.
@@ -79,7 +80,6 @@ struct ModuleHeaderInner {
     compiled_funcs: Box<[CompiledFunc]>,
     compiled_funcs_idx: BTreeMap<CompiledFunc, FuncIdx>,
     element_segments: Box<[ElementSegment]>,
-    custom_sections: Box<[CustomSection]>,
 }
 
 impl ModuleHeader {
@@ -370,7 +370,7 @@ impl Module {
     }
 
     pub fn custom_sections(&self) -> &[CustomSection] {
-        &self.header.inner.custom_sections
+        &self.custom_sections
     }
 
     /// Looks up an export in this [`Module`] by its `name`.

crates/wasmi/src/module/parser.rs

@@ -130,9 +130,10 @@ impl ModuleParser {
     ) -> Result<Module, Error> {
         let mut buffer = Vec::new();
         let header = Self::parse_header(&mut self, &mut stream, &mut buffer)?;
-        let builder =
+        let mut builder =
             Self::parse_code(&mut self, validation_mode, &mut stream, &mut buffer, header)?;
-        let module = Self::parse_data(&mut self, &mut stream, &mut buffer, builder)?;
+        builder = Self::parse_data(&mut self, &mut stream, &mut buffer, builder)?;
+        let module = Self::parse_custom_section(&mut self, &mut stream, &mut buffer, builder)?;
         Ok(module)
     }
 
@@ -202,9 +203,7 @@ impl ModuleParser {
                         }
                         Payload::DataSection(_) => break,
                         Payload::End(_) => break,
-                        Payload::CustomSection(section)  => {
-                            self.process_custom_section(section, &mut header)
-                        },
+                        Payload::CustomSection { .. } => break,
                         Payload::UnknownSection { id, range, .. } => {
                             self.process_unknown(id, range)
                         }
@@ -253,7 +252,6 @@ impl ModuleParser {
                             let bytes = &buffer[start..consumed];
                             self.process_code_entry(func_body, validation_mode, bytes, &header)?;
                         }
-                        Payload::CustomSection { .. } => {}
                         Payload::UnknownSection { id, range, .. } => {
                             self.process_unknown(id, range)?
                         }
@@ -272,7 +270,7 @@ impl ModuleParser {
         stream: &mut impl Read,
         buffer: &mut Vec<u8>,
         mut builder: ModuleBuilder,
-    ) -> Result<Module, Error> {
+    ) -> Result<ModuleBuilder, Error> {
         loop {
             match self.parser.parse(&buffer[..], self.eof)? {
                 Chunk::NeedMoreData(hint) => {
@@ -283,12 +281,44 @@ impl ModuleParser {
                         Payload::DataSection(section) => {
                             self.process_data(section, &mut builder)?;
                         }
+                        Payload::End { .. } => break,
+                        Payload::CustomSection { .. } => break,
+                        Payload::UnknownSection { id, range, .. } => {
+                            self.process_unknown(id, range)?
+                        }
+                        unexpected => {
+                            unreachable!("encountered unexpected Wasm section: {unexpected:?}")
+                        }
+                    }
+                    // Cut away the parts from the intermediate buffer that have already been parsed.
+                    buffer.drain(..consumed);
+                }
+            }
+        }
+        Ok(builder)
+    }
+
+    fn parse_custom_section(
+        &mut self,
+        stream: &mut impl Read,
+        buffer: &mut Vec<u8>,
+        mut builder: ModuleBuilder,
+    ) -> Result<Module, Error> {
+        loop {
+            match self.parser.parse(&buffer[..], self.eof)? {
+                Chunk::NeedMoreData(hint) => {
+                    self.eof = Self::pull_bytes(buffer, hint, stream)?;
+                }
+                Chunk::Parsed { consumed, payload } => {
+                    match payload {
                         Payload::End(offset) => {
                             self.process_end(offset)?;
                             buffer.drain(..consumed);
                             break;
                         }
-                        Payload::CustomSection { .. } => {}
+                        Payload::CustomSection(section)  => {
+                            self.process_custom_section(section, &mut builder)?
+                        }
                         Payload::UnknownSection { id, range, .. } => {
                             self.process_unknown(id, range)?
                         }
@@ -301,7 +331,7 @@ impl ModuleParser {
                 }
             }
         }
-        Ok(builder.finish(&self.engine))
+        Ok(builder.finish(&self.engine))        
     }
 
     /// Pulls more bytes from the `stream` in order to produce Wasm payload.
@@ -605,9 +635,9 @@ impl ModuleParser {
     fn process_custom_section(
         &mut self, 
         section: CustomSectionReader, 
-        header: &mut ModuleHeaderBuilder
+        builder: &mut ModuleBuilder
     ) -> Result<(), Error> {
-        header.push_custom_section(section.name(), section.data());
+        builder.push_custom_section(section.name(), section.data());
         Ok(())
     }