From caac45522331bc35f167c768a6467aeddc17cadc Mon Sep 17 00:00:00 2001
From: Jay Geng <jay@stellar.org>
Date: Fri, 19 May 2023 18:25:56 -0400
Subject: [PATCH] Capture and surface parsed custom sections in modules.

---
 crates/wasmi/src/module/builder.rs |  9 +++++++++
 crates/wasmi/src/module/mod.rs     | 12 ++++++++++++
 crates/wasmi/src/module/parser.rs  |  9 ++++++++-
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/crates/wasmi/src/module/builder.rs b/crates/wasmi/src/module/builder.rs
index 35dcc461c7..2c8851e4d1 100644
--- a/crates/wasmi/src/module/builder.rs
+++ b/crates/wasmi/src/module/builder.rs
@@ -2,6 +2,7 @@ use super::{
     export::ExternIdx,
     import::FuncTypeIdx,
     ConstExpr,
+    CustomSection,
     DataSegment,
     ElementSegment,
     ExternTypeIdx,
@@ -39,6 +40,7 @@ pub struct ModuleBuilder<'engine> {
     pub compiled_funcs: Vec<CompiledFunc>,
     pub element_segments: Vec<ElementSegment>,
     pub data_segments: Vec<DataSegment>,
+    pub(super) custom_sections: Vec<CustomSection>,
 }
 
 /// The import names of the [`Module`] imports.
@@ -139,6 +141,7 @@ impl<'engine> ModuleBuilder<'engine> {
             compiled_funcs: Vec::new(),
             element_segments: Vec::new(),
             data_segments: Vec::new(),
+            custom_sections: Vec::new(),
         }
     }
 
@@ -390,6 +393,12 @@ impl<'engine> ModuleBuilder<'engine> {
         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) -> Module {
         Module::from_builder(self)
diff --git a/crates/wasmi/src/module/mod.rs b/crates/wasmi/src/module/mod.rs
index 1d28ac4989..a2b145a809 100644
--- a/crates/wasmi/src/module/mod.rs
+++ b/crates/wasmi/src/module/mod.rs
@@ -49,6 +49,12 @@ use crate::{
 use alloc::{boxed::Box, collections::BTreeMap, sync::Arc};
 use core::{iter, slice::Iter as SliceIter};
 
+#[derive(Debug)]
+pub struct CustomSection {
+    pub name: Box<str>,
+    pub data: Box<[u8]>,
+}
+
 /// A parsed and validated WebAssembly module.
 #[derive(Debug)]
 pub struct Module {
@@ -65,6 +71,7 @@ pub struct Module {
     compiled_funcs: Box<[CompiledFunc]>,
     element_segments: Box<[ElementSegment]>,
     data_segments: Box<[DataSegment]>,
+    custom_sections: Box<[CustomSection]>,
 }
 
 /// The index of the default Wasm linear memory.
@@ -168,6 +175,7 @@ impl Module {
             compiled_funcs: builder.compiled_funcs.into(),
             element_segments: builder.element_segments.into(),
             data_segments: builder.data_segments.into(),
+            custom_sections: builder.custom_sections.into(),
         }
     }
 
@@ -265,6 +273,10 @@ impl Module {
         ModuleExportsIter::new(self)
     }
 
+    pub fn custom_sections(&self) -> &[CustomSection] {
+        &self.custom_sections
+    }
+
     /// Looks up an export in this [`Module`] by its `name`.
     ///
     /// Returns `None` if no export with the name was found.
diff --git a/crates/wasmi/src/module/parser.rs b/crates/wasmi/src/module/parser.rs
index d617abbb2b..30359550e5 100644
--- a/crates/wasmi/src/module/parser.rs
+++ b/crates/wasmi/src/module/parser.rs
@@ -26,6 +26,7 @@ use core::{
 };
 use wasmparser::{
     Chunk,
+    CustomSectionReader,
     DataSectionReader,
     ElementSectionReader,
     Encoding,
@@ -178,7 +179,7 @@ impl<'engine> ModuleParser<'engine> {
             Payload::ElementSection(section) => self.process_element(section),
             Payload::DataCountSection { count, range } => self.process_data_count(count, range),
             Payload::DataSection(section) => self.process_data(section),
-            Payload::CustomSection { .. } => Ok(()),
+            Payload::CustomSection(section) => self.process_custom_section(section),
             Payload::CodeSectionStart { count, range, .. } => self.process_code_start(count, range),
             Payload::CodeSectionEntry(func_body) => self.process_code_entry(func_body),
             Payload::UnknownSection { id, range, .. } => self.process_unknown(id, range),
@@ -460,6 +461,12 @@ impl<'engine> ModuleParser<'engine> {
         Ok(())
     }
 
+    fn process_custom_section(&mut self, section: CustomSectionReader) -> Result<(), ModuleError> {
+        self.builder
+            .push_custom_section(section.name(), section.data());
+        Ok(())
+    }
+
     /// Process module code section start.
     ///
     /// # Note