From 19acd4582aecfa01900a8f806e52fadc8e59e230 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Wed, 5 Oct 2022 08:16:25 -0700 Subject: [PATCH] Add a library function to view `Payload` as a section This makes the `wasm-tools strip` function easier to write and additionally enables this sort of functionality externally as `wit-bindgen` wants to strip custom sections as well. --- crates/wasmparser/src/parser.rs | 136 +++++++++++++++++++++++++------- src/bin/wasm-tools/strip.rs | 72 +++-------------- 2 files changed, 120 insertions(+), 88 deletions(-) diff --git a/crates/wasmparser/src/parser.rs b/crates/wasmparser/src/parser.rs index f23e7f894c..a29cb57e47 100644 --- a/crates/wasmparser/src/parser.rs +++ b/crates/wasmparser/src/parser.rs @@ -5,8 +5,8 @@ use crate::{ ComponentInstanceSectionReader, ComponentStartSectionReader, ComponentTypeSectionReader, CustomSectionReader, DataSectionReader, ElementSectionReader, ExportSectionReader, FunctionBody, FunctionSectionReader, GlobalSectionReader, ImportSectionReader, - InstanceSectionReader, MemorySectionReader, Result, TableSectionReader, TagSectionReader, - TypeSectionReader, + InstanceSectionReader, MemorySectionReader, Result, SectionReader, TableSectionReader, + TagSectionReader, TypeSectionReader, }; use std::convert::TryInto; use std::fmt; @@ -281,6 +281,33 @@ pub enum Payload<'a> { End(usize), } +const CUSTOM_SECTION: u8 = 0; +const TYPE_SECTION: u8 = 1; +const IMPORT_SECTION: u8 = 2; +const FUNCTION_SECTION: u8 = 3; +const TABLE_SECTION: u8 = 4; +const MEMORY_SECTION: u8 = 5; +const GLOBAL_SECTION: u8 = 6; +const EXPORT_SECTION: u8 = 7; +const START_SECTION: u8 = 8; +const ELEMENT_SECTION: u8 = 9; +const CODE_SECTION: u8 = 10; +const DATA_SECTION: u8 = 11; +const DATA_COUNT_SECTION: u8 = 12; +const TAG_SECTION: u8 = 13; + +const COMPONENT_MODULE_SECTION: u8 = 1; +const COMPONENT_CORE_INSTANCE_SECTION: u8 = 2; +const COMPONENT_CORE_TYPE_SECTION: u8 = 3; +const COMPONENT_SECTION: u8 = 4; +const COMPONENT_INSTANCE_SECTION: u8 = 5; +const COMPONENT_ALIAS_SECTION: u8 = 6; +const COMPONENT_TYPE_SECTION: u8 = 7; +const COMPONENT_CANONICAL_SECTION: u8 = 8; +const COMPONENT_START_SECTION: u8 = 9; +const COMPONENT_IMPORT_SECTION: u8 = 10; +const COMPONENT_EXPORT_SECTION: u8 = 11; + impl Parser { /// Creates a new parser. /// @@ -541,35 +568,35 @@ impl Parser { (_, 0) => section(reader, len, CustomSectionReader::new, CustomSection), // Module sections - (Encoding::Module, 1) => { + (Encoding::Module, TYPE_SECTION) => { section(reader, len, TypeSectionReader::new, TypeSection) } - (Encoding::Module, 2) => { + (Encoding::Module, IMPORT_SECTION) => { section(reader, len, ImportSectionReader::new, ImportSection) } - (Encoding::Module, 3) => { + (Encoding::Module, FUNCTION_SECTION) => { section(reader, len, FunctionSectionReader::new, FunctionSection) } - (Encoding::Module, 4) => { + (Encoding::Module, TABLE_SECTION) => { section(reader, len, TableSectionReader::new, TableSection) } - (Encoding::Module, 5) => { + (Encoding::Module, MEMORY_SECTION) => { section(reader, len, MemorySectionReader::new, MemorySection) } - (Encoding::Module, 6) => { + (Encoding::Module, GLOBAL_SECTION) => { section(reader, len, GlobalSectionReader::new, GlobalSection) } - (Encoding::Module, 7) => { + (Encoding::Module, EXPORT_SECTION) => { section(reader, len, ExportSectionReader::new, ExportSection) } - (Encoding::Module, 8) => { + (Encoding::Module, START_SECTION) => { let (func, range) = single_u32(reader, len, "start")?; Ok(StartSection { func, range }) } - (Encoding::Module, 9) => { + (Encoding::Module, ELEMENT_SECTION) => { section(reader, len, ElementSectionReader::new, ElementSection) } - (Encoding::Module, 10) => { + (Encoding::Module, CODE_SECTION) => { let start = reader.original_position(); let count = delimited(reader, &mut len, |r| r.read_var_u32())?; let range = start..reader.original_position() + len as usize; @@ -583,20 +610,20 @@ impl Parser { size: len, }) } - (Encoding::Module, 11) => { + (Encoding::Module, DATA_SECTION) => { section(reader, len, DataSectionReader::new, DataSection) } - (Encoding::Module, 12) => { + (Encoding::Module, DATA_COUNT_SECTION) => { let (count, range) = single_u32(reader, len, "data count")?; Ok(DataCountSection { count, range }) } - (Encoding::Module, 13) => { + (Encoding::Module, TAG_SECTION) => { section(reader, len, TagSectionReader::new, TagSection) } // Component sections - (Encoding::Component, 1 /* module */) - | (Encoding::Component, 4 /* component */) => { + (Encoding::Component, COMPONENT_MODULE_SECTION) + | (Encoding::Component, COMPONENT_SECTION) => { if len as usize > MAX_WASM_MODULE_SIZE { bail!( len_pos, @@ -618,50 +645,49 @@ impl Parser { _ => unreachable!(), }) } - (Encoding::Component, 2) => { + (Encoding::Component, COMPONENT_CORE_INSTANCE_SECTION) => { section(reader, len, InstanceSectionReader::new, InstanceSection) } - (Encoding::Component, 3) => { + (Encoding::Component, COMPONENT_CORE_TYPE_SECTION) => { section(reader, len, CoreTypeSectionReader::new, CoreTypeSection) } - // Section 5 handled above - (Encoding::Component, 5) => section( + (Encoding::Component, COMPONENT_INSTANCE_SECTION) => section( reader, len, ComponentInstanceSectionReader::new, ComponentInstanceSection, ), - (Encoding::Component, 6) => section( + (Encoding::Component, COMPONENT_ALIAS_SECTION) => section( reader, len, ComponentAliasSectionReader::new, ComponentAliasSection, ), - (Encoding::Component, 7) => section( + (Encoding::Component, COMPONENT_TYPE_SECTION) => section( reader, len, ComponentTypeSectionReader::new, ComponentTypeSection, ), - (Encoding::Component, 8) => section( + (Encoding::Component, COMPONENT_CANONICAL_SECTION) => section( reader, len, ComponentCanonicalSectionReader::new, ComponentCanonicalSection, ), - (Encoding::Component, 9) => section( + (Encoding::Component, COMPONENT_START_SECTION) => section( reader, len, ComponentStartSectionReader::new, ComponentStartSection, ), - (Encoding::Component, 10) => section( + (Encoding::Component, COMPONENT_IMPORT_SECTION) => section( reader, len, ComponentImportSectionReader::new, ComponentImportSection, ), - (Encoding::Component, 11) => section( + (Encoding::Component, COMPONENT_EXPORT_SECTION) => section( reader, len, ComponentExportSectionReader::new, @@ -930,6 +956,62 @@ impl Default for Parser { } } +impl Payload<'_> { + /// If this `Payload` represents a section in the original wasm module then + /// the section's id and range within the original wasm binary are returned. + /// + /// Not all payloads refer to entire sections, such as the `Version` and + /// `CodeSectionEntry` variants. These variants will return `None` from this + /// function. + /// + /// Otherwise this function will return `Some` where the first element is + /// the byte identifier for the section and the second element is the range + /// of the contents of the section within the original wasm binary. + /// + /// The purpose of this method is to enable tools to easily iterate over + /// entire sections if necessary and handle sections uniformly, for example + /// dropping custom sections while preserving all other sections. + pub fn as_section(&self) -> Option<(u8, Range)> { + use Payload::*; + + match self { + Version { .. } => None, + TypeSection(s) => Some((TYPE_SECTION, s.range())), + ImportSection(s) => Some((IMPORT_SECTION, s.range())), + FunctionSection(s) => Some((FUNCTION_SECTION, s.range())), + TableSection(s) => Some((TABLE_SECTION, s.range())), + MemorySection(s) => Some((MEMORY_SECTION, s.range())), + TagSection(s) => Some((TAG_SECTION, s.range())), + GlobalSection(s) => Some((GLOBAL_SECTION, s.range())), + ExportSection(s) => Some((EXPORT_SECTION, s.range())), + ElementSection(s) => Some((ELEMENT_SECTION, s.range())), + DataSection(s) => Some((DATA_SECTION, s.range())), + StartSection { range, .. } => Some((START_SECTION, range.clone())), + DataCountSection { range, .. } => Some((DATA_COUNT_SECTION, range.clone())), + CodeSectionStart { range, .. } => Some((CODE_SECTION, range.clone())), + CodeSectionEntry(_) => None, + + ModuleSection { range, .. } => Some((COMPONENT_MODULE_SECTION, range.clone())), + InstanceSection(s) => Some((COMPONENT_CORE_INSTANCE_SECTION, s.range())), + CoreTypeSection(s) => Some((COMPONENT_CORE_TYPE_SECTION, s.range())), + ComponentSection { range, .. } => Some((COMPONENT_SECTION, range.clone())), + ComponentInstanceSection(s) => Some((COMPONENT_INSTANCE_SECTION, s.range())), + ComponentAliasSection(s) => Some((COMPONENT_ALIAS_SECTION, s.range())), + ComponentTypeSection(s) => Some((COMPONENT_TYPE_SECTION, s.range())), + ComponentCanonicalSection(s) => Some((COMPONENT_CANONICAL_SECTION, s.range())), + ComponentStartSection(s) => Some((COMPONENT_START_SECTION, s.range())), + ComponentImportSection(s) => Some((COMPONENT_IMPORT_SECTION, s.range())), + ComponentExportSection(s) => Some((COMPONENT_EXPORT_SECTION, s.range())), + + CustomSection(c) => Some((CUSTOM_SECTION, c.range())), + + UnknownSection { id, range, .. } => Some((*id, range.clone())), + + End(_) => None, + } + } +} + impl fmt::Debug for Payload<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { use Payload::*; diff --git a/src/bin/wasm-tools/strip.rs b/src/bin/wasm-tools/strip.rs index f7a4be5612..274a2c4e0e 100644 --- a/src/bin/wasm-tools/strip.rs +++ b/src/bin/wasm-tools/strip.rs @@ -1,7 +1,6 @@ use anyhow::{bail, Result}; -use std::ops::Range; -use wasm_encoder::{RawSection, SectionId}; -use wasmparser::{Encoding, Parser, Payload::*, SectionReader}; +use wasm_encoder::RawSection; +use wasmparser::{Encoding, Parser, Payload::*}; /// Removes custom sections from an input WebAssembly file. /// @@ -50,69 +49,20 @@ impl Opts { for payload in Parser::new(0).parse_all(&input) { let payload = payload?; - let mut section = |id: SectionId, range: Range| { - module.section(&RawSection { - id: id as u8, - data: &input[range], - }); - }; - match payload { - Version { - encoding: Encoding::Module, - .. - } => {} - Version { - encoding: Encoding::Component, - .. - } => { - bail!("components are not supported yet with the `strip` command"); - } - - TypeSection(s) => section(SectionId::Type, s.range()), - ImportSection(s) => section(SectionId::Import, s.range()), - FunctionSection(s) => section(SectionId::Function, s.range()), - TableSection(s) => section(SectionId::Table, s.range()), - MemorySection(s) => section(SectionId::Memory, s.range()), - TagSection(s) => section(SectionId::Tag, s.range()), - GlobalSection(s) => section(SectionId::Global, s.range()), - ExportSection(s) => section(SectionId::Export, s.range()), - ElementSection(s) => section(SectionId::Element, s.range()), - DataSection(s) => section(SectionId::Data, s.range()), - StartSection { range, .. } => section(SectionId::Start, range), - DataCountSection { range, .. } => section(SectionId::DataCount, range), - CodeSectionStart { range, .. } => section(SectionId::Code, range), - CodeSectionEntry(_) => {} - - ModuleSection { .. } - | InstanceSection(_) - | CoreTypeSection(_) - | ComponentSection { .. } - | ComponentInstanceSection(_) - | ComponentAliasSection(_) - | ComponentTypeSection(_) - | ComponentCanonicalSection(_) - | ComponentStartSection(_) - | ComponentImportSection(_) - | ComponentExportSection(_) => unimplemented!("component model"), - + match &payload { CustomSection(c) => { - if !strip_custom_section(c.name()) { - module.section(&RawSection { - id: SectionId::Custom as u8, - data: &input[c.range()], - }); + if strip_custom_section(c.name()) { + continue; } } - UnknownSection { + _ => {} + } + if let Some((id, range)) = payload.as_section() { + module.section(&RawSection { id, - contents, - range: _, - } => { - module.section(&RawSection { id, data: contents }); - } - - End(_) => {} + data: &input[range], + }); } }