Add the concept of a "known custom section" to wasmparser (#1544)

* Add the concept of a "known custom section" to wasmparser

There are a number of places throughout wasm-tools which process custom
sections, and as surfaced in #1540 they're not necessarily all handled
uniformly due to the flexible nature of custom sections. This commit aims to
improve the parsing of custom sections in the future by having a
`wasmparser::KnownCustom` enum which statically enumerates all custom
sections that are known to wasmparser and can be parsed. This
additionally provides a nice convenience for creating the typed parsers
as well.

* Add a known "linking" section, add it to `dump`, and test it

* Add reloc to known custom sections

Author: alexcrichton

crates/wasm-encoder/src/core/dump.rs

@@ -336,7 +336,7 @@ impl Encode for CoreDumpValue {
 mod tests {
     use super::*;
     use crate::Module;
-    use wasmparser::{BinaryReader, FromReader, Parser, Payload};
+    use wasmparser::{KnownCustom, Parser, Payload};
 
     // Create new core dump section and test whether it is properly encoded and
     // parsed back out by wasmparser
@@ -361,10 +361,10 @@ mod tests {
         match payload {
             Payload::CustomSection(section) => {
                 assert_eq!(section.name(), "core");
-                let core = wasmparser::CoreDumpSection::from_reader(&mut BinaryReader::new(
-                    section.data(),
-                ))
-                .expect("data is readable into a core dump section");
+                let core = match section.as_known() {
+                    KnownCustom::CoreDump(s) => s,
+                    _ => panic!("not coredump"),
+                };
                 assert_eq!(core.name, "test.wasm");
             }
             _ => panic!("unexpected payload"),
@@ -394,10 +394,10 @@ mod tests {
         match payload {
             Payload::CustomSection(section) => {
                 assert_eq!(section.name(), "coremodules");
-                let modules = wasmparser::CoreDumpModulesSection::from_reader(
-                    &mut BinaryReader::new(section.data()),
-                )
-                .expect("data is readable into a core dump modules section");
+                let modules = match section.as_known() {
+                    KnownCustom::CoreDumpModules(s) => s,
+                    _ => panic!("not coremodules"),
+                };
                 assert_eq!(modules.modules[0], "test_module");
             }
             _ => panic!("unexpected payload"),
@@ -429,10 +429,10 @@ mod tests {
         match payload {
             Payload::CustomSection(section) => {
                 assert_eq!(section.name(), "coreinstances");
-                let coreinstances = wasmparser::CoreDumpInstancesSection::from_reader(
-                    &mut BinaryReader::new(section.data()),
-                )
-                .expect("data is readable into a core dump instances section");
+                let coreinstances = match section.as_known() {
+                    KnownCustom::CoreDumpInstances(s) => s,
+                    _ => panic!("not coreinstances"),
+                };
                 assert_eq!(coreinstances.instances.len(), 1);
                 let instance = coreinstances
                     .instances
@@ -475,10 +475,10 @@ mod tests {
         match payload {
             Payload::CustomSection(section) => {
                 assert_eq!(section.name(), "corestack");
-                let corestack = wasmparser::CoreDumpStackSection::from_reader(
-                    &mut BinaryReader::new(section.data()),
-                )
-                .expect("data is readable into a core dump stack section");
+                let corestack = match section.as_known() {
+                    KnownCustom::CoreDumpStack(s) => s,
+                    _ => panic!("not a corestack section"),
+                };
                 assert_eq!(corestack.name, "main");
                 assert_eq!(corestack.frames.len(), 1);
                 let frame = corestack
@@ -570,7 +570,7 @@ mod tests {
             // 0x0, module_idx
             0x0, 0,
             // memories count, memories
-            1, 42, 
+            1, 42,
             // globals count, globals
             1, 17
         ]);
@@ -595,7 +595,7 @@ mod tests {
             encoded,
             vec![
                 // section length
-                27, 
+                27,
                 // length of name.
                 9,
                 // section name (corestack)

crates/wasm-encoder/src/core/producers.rs

@@ -121,7 +121,7 @@ mod test {
     #[test]
     fn roundtrip_example() {
         use crate::{Module, ProducersField, ProducersSection};
-        use wasmparser::{Parser, Payload, ProducersSectionReader};
+        use wasmparser::{KnownCustom, Parser, Payload};
 
         // Create a new producers section.
         let mut field = ProducersField::new();
@@ -151,9 +151,10 @@ mod test {
         match payload {
             Payload::CustomSection(c) => {
                 assert_eq!(c.name(), "producers");
-                let mut section = ProducersSectionReader::new(c.data(), c.data_offset())
-                    .expect("readable as a producers section")
-                    .into_iter();
+                let mut section = match c.as_known() {
+                    KnownCustom::Producers(s) => s.into_iter(),
+                    _ => panic!("unknown custom section"),
+                };
                 let field = section
                     .next()
                     .expect("section has an element")

crates/wasm-metadata/src/lib.rs

@@ -9,7 +9,8 @@ use std::mem;
 use std::ops::Range;
 use wasm_encoder::{ComponentSection as _, ComponentSectionId, Encode, Section};
 use wasmparser::{
-    ComponentNameSectionReader, NameSectionReader, Parser, Payload::*, ProducersSectionReader,
+    ComponentNameSectionReader, KnownCustom, NameSectionReader, Parser, Payload::*,
+    ProducersSectionReader,
 };
 
 /// A representation of a WebAssembly producers section.
@@ -50,9 +51,11 @@ impl Producers {
             match payload {
                 ModuleSection { .. } | ComponentSection { .. } => depth += 1,
                 End { .. } => depth -= 1,
-                CustomSection(c) if c.name() == "producers" && depth == 0 => {
-                    let producers = Self::from_bytes(c.data(), c.data_offset())?;
-                    return Ok(Some(producers));
+                CustomSection(c) if depth == 0 => {
+                    if let KnownCustom::Producers(_) = c.as_known() {
+                        let producers = Self::from_bytes(c.data(), c.data_offset())?;
+                        return Ok(Some(producers));
+                    }
                 }
                 _ => {}
             }
@@ -290,57 +293,60 @@ fn rewrite_wasm(
             _ => {}
         }
 
-        // Process the wasm sections:
-        match payload {
-            // Only rewrite the outermost producers section:
-            CustomSection(c) if c.name() == "producers" && stack.len() == 0 => {
-                producers_found = true;
-                let mut producers = Producers::from_bytes(c.data(), c.data_offset())?;
-                // Add to the section according to the command line flags:
-                producers.merge(&add_producers);
-                // Encode into output:
-                producers.section().append_to(&mut output);
-            }
-
-            CustomSection(c) if c.name() == "name" && stack.len() == 0 => {
-                names_found = true;
-                let mut names = ModuleNames::from_bytes(c.data(), c.data_offset())?;
-                names.merge(&ModuleNames::from_name(add_name));
-
-                names.section()?.as_custom().append_to(&mut output);
-            }
-
-            CustomSection(c) if c.name() == "component-name" && stack.len() == 0 => {
-                names_found = true;
-                let mut names = ComponentNames::from_bytes(c.data(), c.data_offset())?;
-                names.merge(&ComponentNames::from_name(add_name));
-                names.section()?.as_custom().append_to(&mut output);
-            }
-
-            CustomSection(c) if c.name() == "registry-metadata" && stack.len() == 0 => {
-                // Pass section through if a new registry metadata isn't provided, otherwise ignore and overwrite with new
-                if add_registry_metadata.is_none() {
-                    let registry: RegistryMetadata = RegistryMetadata::from_bytes(&c.data(), 0)?;
-
-                    let registry_metadata = wasm_encoder::CustomSection {
-                        name: Cow::Borrowed("registry-metadata"),
-                        data: Cow::Owned(serde_json::to_vec(&registry)?),
-                    };
-                    registry_metadata.append_to(&mut output);
-                }
-            }
+        // Only rewrite the outermost custom sections
+        if let CustomSection(c) = &payload {
+            if stack.len() == 0 {
+                match c.as_known() {
+                    KnownCustom::Producers(_) => {
+                        producers_found = true;
+                        let mut producers = Producers::from_bytes(c.data(), c.data_offset())?;
+                        // Add to the section according to the command line flags:
+                        producers.merge(&add_producers);
+                        // Encode into output:
+                        producers.section().append_to(&mut output);
+                        continue;
+                    }
+                    KnownCustom::Name(_) => {
+                        names_found = true;
+                        let mut names = ModuleNames::from_bytes(c.data(), c.data_offset())?;
+                        names.merge(&ModuleNames::from_name(add_name));
 
-            // All other sections get passed through unmodified:
-            _ => {
-                if let Some((id, range)) = payload.as_section() {
-                    wasm_encoder::RawSection {
-                        id,
-                        data: &input[range],
+                        names.section()?.as_custom().append_to(&mut output);
+                        continue;
+                    }
+                    KnownCustom::ComponentName(_) => {
+                        names_found = true;
+                        let mut names = ComponentNames::from_bytes(c.data(), c.data_offset())?;
+                        names.merge(&ComponentNames::from_name(add_name));
+                        names.section()?.as_custom().append_to(&mut output);
+                        continue;
                     }
-                    .append_to(&mut output);
+                    KnownCustom::Unknown if c.name() == "registry-metadata" => {
+                        // Pass section through if a new registry metadata isn't provided, otherwise ignore and overwrite with new
+                        if add_registry_metadata.is_none() {
+                            let registry: RegistryMetadata =
+                                RegistryMetadata::from_bytes(&c.data(), 0)?;
+
+                            let registry_metadata = wasm_encoder::CustomSection {
+                                name: Cow::Borrowed("registry-metadata"),
+                                data: Cow::Owned(serde_json::to_vec(&registry)?),
+                            };
+                            registry_metadata.append_to(&mut output);
+                            continue;
+                        }
+                    }
+                    _ => {}
                 }
             }
         }
+        // All other sections get passed through unmodified:
+        if let Some((id, range)) = payload.as_section() {
+            wasm_encoder::RawSection {
+                id,
+                data: &input[range],
+            }
+            .append_to(&mut output);
+        }
     }
     if !names_found && add_name.is_some() {
         if output.starts_with(&wasm_encoder::Component::HEADER) {
@@ -434,39 +440,42 @@ impl Metadata {
                         metadata.last_mut().unwrap().push_child(finished);
                     }
                 }
-                CustomSection(c) if c.name() == "name" => {
-                    let names = ModuleNames::from_bytes(c.data(), c.data_offset())?;
-                    if let Some(name) = names.get_name() {
+                CustomSection(c) => match c.as_known() {
+                    KnownCustom::Name(_) => {
+                        let names = ModuleNames::from_bytes(c.data(), c.data_offset())?;
+                        if let Some(name) = names.get_name() {
+                            metadata
+                                .last_mut()
+                                .expect("non-empty metadata stack")
+                                .set_name(&name);
+                        }
+                    }
+                    KnownCustom::ComponentName(_) => {
+                        let names = ComponentNames::from_bytes(c.data(), c.data_offset())?;
+                        if let Some(name) = names.get_name() {
+                            metadata
+                                .last_mut()
+                                .expect("non-empty metadata stack")
+                                .set_name(name);
+                        }
+                    }
+                    KnownCustom::Producers(_) => {
+                        let producers = Producers::from_bytes(c.data(), c.data_offset())?;
                         metadata
                             .last_mut()
                             .expect("non-empty metadata stack")
-                            .set_name(&name);
+                            .set_producers(producers);
                     }
-                }
-                CustomSection(c) if c.name() == "component-name" => {
-                    let names = ComponentNames::from_bytes(c.data(), c.data_offset())?;
-                    if let Some(name) = names.get_name() {
+                    KnownCustom::Unknown if c.name() == "registry-metadata" => {
+                        let registry: RegistryMetadata =
+                            RegistryMetadata::from_bytes(&c.data(), 0)?;
                         metadata
                             .last_mut()
                             .expect("non-empty metadata stack")
-                            .set_name(name);
+                            .set_registry_metadata(registry);
                     }
-                }
-                CustomSection(c) if c.name() == "producers" => {
-                    let producers = Producers::from_bytes(c.data(), c.data_offset())?;
-                    metadata
-                        .last_mut()
-                        .expect("non-empty metadata stack")
-                        .set_producers(producers);
-                }
-                CustomSection(c) if c.name() == "registry-metadata" => {
-                    let registry: RegistryMetadata = RegistryMetadata::from_bytes(&c.data(), 0)?;
-                    metadata
-                        .last_mut()
-                        .expect("non-empty metadata stack")
-                        .set_registry_metadata(registry);
-                }
-
+                    _ => {}
+                },
                 _ => {}
             }
         }