From 31c3c663b10a1d2d1aa476f175551450e9771b9b Mon Sep 17 00:00:00 2001 From: Jay Geng Date: Wed, 3 Jan 2024 14:35:25 -0500 Subject: [PATCH] fixup! Capture and surface parsed custom sections in modules. --- crates/wasmi/src/module/builder.rs | 18 +++++------ crates/wasmi/src/module/mod.rs | 4 +-- crates/wasmi/src/module/parser.rs | 52 +++++++++++++++++++++++------- 3 files changed, 52 insertions(+), 22 deletions(-) diff --git a/crates/wasmi/src/module/builder.rs b/crates/wasmi/src/module/builder.rs index 50efd0c854..4db215cbee 100644 --- a/crates/wasmi/src/module/builder.rs +++ b/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, + pub custom_sections: Vec, } /// A builder for a WebAssembly [`Module`] header. @@ -50,7 +51,6 @@ pub struct ModuleHeaderBuilder { pub compiled_funcs: Vec, pub compiled_funcs_idx: BTreeMap, pub element_segments: Vec, - pub(super) custom_sections: Vec, } 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::, _>>()?; Ok(()) } - - pub fn push_custom_section(&mut self, name: &str, data: &[u8]) { - let name: Box = 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 = 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(), } } } diff --git a/crates/wasmi/src/module/mod.rs b/crates/wasmi/src/module/mod.rs index d83a8a9eeb..13d94967e7 100644 --- a/crates/wasmi/src/module/mod.rs +++ b/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, 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`. diff --git a/crates/wasmi/src/module/parser.rs b/crates/wasmi/src/module/parser.rs index 0b44eeaba4..6107565e88 100644 --- a/crates/wasmi/src/module/parser.rs +++ b/crates/wasmi/src/module/parser.rs @@ -130,9 +130,10 @@ impl ModuleParser { ) -> Result { 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, mut builder: ModuleBuilder, - ) -> Result { + ) -> Result { 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, + mut builder: ModuleBuilder, + ) -> Result { + 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(()) }