Skip to content

Commit

Permalink
fixup! fixup! Capture and surface parsed custom sections in modules.
Browse files Browse the repository at this point in the history
  • Loading branch information
@jayz22
jayz22 committed Jan 4, 2024
1 parent 31c3c66 commit de10022
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 45 deletions.
6 changes: 2 additions & 4 deletions crates/wasmi/src/module/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -393,10 +393,8 @@ 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 })
pub fn push_custom_sections(&mut self, custom_sections: Vec<CustomSection>) {
self.custom_sections = custom_sections;
}

/// Finishes construction of the WebAssembly [`Module`].
Expand Down
64 changes: 23 additions & 41 deletions crates/wasmi/src/module/parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use super::{
ModuleBuilder,
ModuleHeader,
Read,
CustomSection
};
use crate::{engine::CompiledFunc, Engine, Error, FuncType, MemoryType, TableType};
use alloc::{boxed::Box, vec::Vec};
Expand Down Expand Up @@ -129,12 +130,13 @@ impl ModuleParser {
mut stream: impl Read,
) -> Result<Module, Error> {
let mut buffer = Vec::new();
let header = Self::parse_header(&mut self, &mut stream, &mut buffer)?;
let mut custom_sections: Vec<CustomSection> = vec![];
let header = Self::parse_header(&mut self, &mut stream, &mut buffer, &mut custom_sections)?;
let mut builder =
Self::parse_code(&mut self, validation_mode, &mut stream, &mut buffer, header)?;
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)
Self::parse_code(&mut self, validation_mode, &mut stream, &mut buffer, header, &mut custom_sections)?;
builder = Self::parse_data(&mut self, &mut stream, &mut buffer, builder, &mut custom_sections)?;
builder.push_custom_sections(custom_sections);
Ok(builder.finish(&self.engine))
}

/// Parse the Wasm module header.
Expand All @@ -151,6 +153,7 @@ impl ModuleParser {
&mut self,
stream: &mut impl Read,
buffer: &mut Vec<u8>,
custom_sections: &mut Vec<CustomSection>
) -> Result<ModuleHeader, Error> {
let mut header = ModuleHeaderBuilder::new(&self.engine);
loop {
Expand Down Expand Up @@ -203,7 +206,9 @@ impl ModuleParser {
}
Payload::DataSection(_) => break,
Payload::End(_) => break,
Payload::CustomSection { .. } => break,
Payload::CustomSection(section) => {
self.process_custom_section(section, custom_sections)
}
Payload::UnknownSection { id, range, .. } => {
self.process_unknown(id, range)
}
Expand Down Expand Up @@ -234,6 +239,7 @@ impl ModuleParser {
stream: &mut impl Read,
buffer: &mut Vec<u8>,
header: ModuleHeader,
custom_sections: &mut Vec<CustomSection>
) -> Result<ModuleBuilder, Error> {
loop {
match self.parser.parse(&buffer[..], self.eof)? {
Expand All @@ -252,6 +258,9 @@ impl ModuleParser {
let bytes = &buffer[start..consumed];
self.process_code_entry(func_body, validation_mode, bytes, &header)?;
}
Payload::CustomSection(section) => {
self.process_custom_section(section, custom_sections)?
}
Payload::UnknownSection { id, range, .. } => {
self.process_unknown(id, range)?
}
Expand All @@ -270,6 +279,7 @@ impl ModuleParser {
stream: &mut impl Read,
buffer: &mut Vec<u8>,
mut builder: ModuleBuilder,
custom_sections: &mut Vec<CustomSection>
) -> Result<ModuleBuilder, Error> {
loop {
match self.parser.parse(&buffer[..], self.eof)? {
Expand All @@ -281,43 +291,13 @@ 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(section) => {
self.process_custom_section(section, &mut builder)?
Payload::CustomSection(section) => {
self.process_custom_section(section, custom_sections)?
}
Payload::UnknownSection { id, range, .. } => {
self.process_unknown(id, range)?
Expand All @@ -331,7 +311,7 @@ impl ModuleParser {
}
}
}
Ok(builder.finish(&self.engine))
Ok(builder)
}

/// Pulls more bytes from the `stream` in order to produce Wasm payload.
Expand Down Expand Up @@ -635,9 +615,11 @@ impl ModuleParser {
fn process_custom_section(
&mut self,
section: CustomSectionReader,
builder: &mut ModuleBuilder
custom_sections: &mut Vec<CustomSection>
) -> Result<(), Error> {
builder.push_custom_section(section.name(), section.data());
let name: Box<str> = section.name().into();
let data: Box<[u8]> = section.data().into();
custom_sections.push(CustomSection { name, data });
Ok(())
}

Expand Down

0 comments on commit de10022

Please sign in to comment.