Fix metadata/strip on components with core modules (#932)

This fixes the `metadata add` and `strip` loops in the `wasm-tools` CLI
for components to correctly handle nested components and nested modules.
Previously the output contents would get corrupted due to modules
showing up multiple times by accident.

Author: alexcrichton

crates/wasm-encoder/src/component.rs

@@ -42,6 +42,12 @@ const INSTANCE_SORT: u8 = 0x05;
 pub trait ComponentSection: Encode {
     /// Gets the section identifier for this section.
     fn id(&self) -> u8;
+
+    /// Appends this section to the specified destination list of bytes.
+    fn append_to_component(&self, dst: &mut Vec<u8>) {
+        dst.push(self.id());
+        self.encode(dst);
+    }
 }
 
 /// Known section identifiers of WebAssembly components.
@@ -101,13 +107,19 @@ pub struct Component {
 }
 
 impl Component {
+    /// The 8-byte header at the beginning of all components.
+    #[rustfmt::skip]
+    pub const HEADER: [u8; 8] = [
+        // Magic
+        0x00, 0x61, 0x73, 0x6D,
+        // Version
+        0x0c, 0x00, 0x01, 0x00,
+    ];
+
     /// Begin writing a new `Component`.
     pub fn new() -> Self {
         Self {
-            bytes: vec![
-                0x00, 0x61, 0x73, 0x6D, // magic (`\0asm`)
-                0x0c, 0x00, 0x01, 0x00, // version
-            ],
+            bytes: Self::HEADER.to_vec(),
         }
     }
 

crates/wasm-encoder/src/core.rs

@@ -48,6 +48,12 @@ pub(crate) const CORE_TAG_SORT: u8 = 0x04;
 pub trait Section: Encode {
     /// Gets the section identifier for this section.
     fn id(&self) -> u8;
+
+    /// Appends this section to the specified destination list of bytes.
+    fn append_to(&self, dst: &mut Vec<u8>) {
+        dst.push(self.id());
+        self.encode(dst);
+    }
 }
 
 /// Known section identifiers of WebAssembly modules.
@@ -110,16 +116,20 @@ pub struct Module {
 }
 
 impl Module {
+    /// The 8-byte header at the beginning of all core wasm modules.
+    #[rustfmt::skip]
+    pub const HEADER: [u8; 8] = [
+        // Magic
+        0x00, 0x61, 0x73, 0x6D,
+        // Version
+        0x01, 0x00, 0x00, 0x00,
+    ];
+
     /// Begin writing a new `Module`.
     #[rustfmt::skip]
     pub fn new() -> Self {
         Module {
-            bytes: vec![
-                // Magic
-                0x00, 0x61, 0x73, 0x6D,
-                // Version
-                0x01, 0x00, 0x00, 0x00,
-            ],
+            bytes: Self::HEADER.to_vec(),
         }
     }
 

crates/wasm-metadata/src/lib.rs

@@ -2,6 +2,8 @@ use anyhow::Result;
 use indexmap::{map::Entry, IndexMap};
 use serde::Serialize;
 use std::fmt;
+use std::mem;
+use wasm_encoder::{ComponentSection as _, ComponentSectionId, Encode, Section};
 use wasmparser::{
     ComponentNameSectionReader, NameSectionReader, Parser, Payload::*, ProducersSectionReader,
 };
@@ -222,125 +224,101 @@ fn rewrite_wasm(
     add_producers: &Producers,
     input: &[u8],
 ) -> Result<Vec<u8>> {
-    let mut parser = Parser::new(0).parse_all(&input);
-
-    enum Output {
-        Component(wasm_encoder::Component),
-        Module(wasm_encoder::Module),
-    }
-    impl Output {
-        fn section(
-            &mut self,
-            section: &(impl wasm_encoder::Section + wasm_encoder::ComponentSection),
-        ) {
-            match self {
-                Output::Component(c) => {
-                    c.section(section);
-                }
-                Output::Module(m) => {
-                    m.section(section);
-                }
-            }
-        }
-        fn finish(self) -> Vec<u8> {
-            match self {
-                Output::Component(c) => c.finish(),
-                Output::Module(m) => m.finish(),
-            }
-        }
-    }
-
-    let mut output = match parser
-        .next()
-        .ok_or_else(|| anyhow::anyhow!("at least a version tag on binary"))??
-    {
-        Version {
-            encoding: wasmparser::Encoding::Component,
-            ..
-        } => Output::Component(wasm_encoder::Component::new()),
-        Version {
-            encoding: wasmparser::Encoding::Module,
-            ..
-        } => Output::Module(wasm_encoder::Module::new()),
-        _ => {
-            panic!("first item from parser must be a Version tag")
-        }
-    };
-
     let mut producers_found = false;
     let mut names_found = false;
-    let mut depth = 0;
-    for payload in parser {
+    let mut stack = Vec::new();
+    let mut output = Vec::new();
+    for payload in Parser::new(0).parse_all(&input) {
         let payload = payload?;
 
         // Track nesting depth, so that we don't mess with inner producer sections:
         match payload {
-            ModuleSection { .. } | ComponentSection { .. } => depth += 1,
-            End { .. } => depth -= 1,
+            Version { encoding, .. } => {
+                output.extend_from_slice(match encoding {
+                    wasmparser::Encoding::Component => &wasm_encoder::Component::HEADER,
+                    wasmparser::Encoding::Module => &wasm_encoder::Module::HEADER,
+                });
+            }
+            ModuleSection { .. } | ComponentSection { .. } => {
+                stack.push(mem::take(&mut output));
+                continue;
+            }
+            End { .. } => {
+                let mut parent = match stack.pop() {
+                    Some(c) => c,
+                    None => break,
+                };
+                if output.starts_with(&wasm_encoder::Component::HEADER) {
+                    parent.push(ComponentSectionId::Component as u8);
+                    output.encode(&mut parent);
+                } else {
+                    parent.push(ComponentSectionId::CoreModule as u8);
+                    output.encode(&mut parent);
+                }
+                output = parent;
+            }
             _ => {}
         }
 
         // Process the wasm sections:
         match payload {
             // Only rewrite the outermost producers section:
-            CustomSection(c) if c.name() == "producers" && depth == 0 => {
+            CustomSection(c) if c.name() == "producers" && stack.len() == 0 => {
                 producers_found = true;
                 let section = ProducersSectionReader::new(c.data(), c.data_offset())?;
                 let mut producers = Producers::from_reader(section)?;
                 // Add to the section according to the command line flags:
                 producers.merge(&add_producers);
                 // Encode into output:
-                output.section(&producers.section());
+                producers.section().append_to(&mut output);
             }
 
-            CustomSection(c) if c.name() == "name" && depth == 0 => {
+            CustomSection(c) if c.name() == "name" && stack.len() == 0 => {
                 names_found = true;
                 let section = NameSectionReader::new(c.data(), c.data_offset());
                 let mut names = ModuleNames::from_reader(section)?;
                 names.merge(&ModuleNames::from_name(add_name));
 
-                output.section(&names.section()?.as_custom());
+                names.section()?.as_custom().append_to(&mut output);
             }
 
-            CustomSection(c) if c.name() == "component-name" && depth == 0 => {
+            CustomSection(c) if c.name() == "component-name" && stack.len() == 0 => {
                 names_found = true;
                 let section = ComponentNameSectionReader::new(c.data(), c.data_offset());
                 let mut names = ComponentNames::from_reader(section)?;
                 names.merge(&ComponentNames::from_name(add_name));
-                output.section(&names.section()?.as_custom());
+                names.section()?.as_custom().append_to(&mut output);
             }
 
             // All other sections get passed through unmodified:
             _ => {
                 if let Some((id, range)) = payload.as_section() {
-                    output.section(&wasm_encoder::RawSection {
+                    wasm_encoder::RawSection {
                         id,
                         data: &input[range],
-                    });
+                    }
+                    .append_to(&mut output);
                 }
             }
         }
     }
     if !names_found && add_name.is_some() {
-        match &mut output {
-            Output::Component(c) => {
-                let names = ComponentNames::from_name(add_name);
-                c.section(&names.section()?);
-            }
-            Output::Module(m) => {
-                let names = ModuleNames::from_name(add_name);
-                m.section(&names.section()?);
-            }
+        if output.starts_with(&wasm_encoder::Component::HEADER) {
+            let names = ComponentNames::from_name(add_name);
+            names.section()?.append_to_component(&mut output);
+        } else {
+            let names = ModuleNames::from_name(add_name);
+            names.section()?.append_to(&mut output)
         }
     }
     if !producers_found && !add_producers.is_empty() {
         let mut producers = Producers::empty();
         // Add to the section according to the command line flags:
         producers.merge(add_producers);
         // Encode into output:
-        output.section(&producers.section());
+        producers.section().append_to(&mut output);
     }
-    Ok(output.finish())
+    Ok(output)
 }
 
 /// A tree of the metadata found in a WebAssembly binary.

src/bin/wasm-tools/strip.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
-use wasm_encoder::RawSection;
+use std::mem;
+use wasm_encoder::{ComponentSectionId, Encode, RawSection, Section};
 use wasmparser::{Parser, Payload::*};
 
 /// Removes custom sections from an input WebAssembly file.
@@ -45,10 +46,41 @@ impl Opts {
             name != "name"
         };
 
-        let mut module = wasm_encoder::Module::new();
+        let mut output = Vec::new();
+        let mut stack = Vec::new();
 
         for payload in Parser::new(0).parse_all(&input) {
             let payload = payload?;
+
+            // Track nesting depth, so that we don't mess with inner producer sections:
+            match payload {
+                Version { encoding, .. } => {
+                    output.extend_from_slice(match encoding {
+                        wasmparser::Encoding::Component => &wasm_encoder::Component::HEADER,
+                        wasmparser::Encoding::Module => &wasm_encoder::Module::HEADER,
+                    });
+                }
+                ModuleSection { .. } | ComponentSection { .. } => {
+                    stack.push(mem::take(&mut output));
+                    continue;
+                }
+                End { .. } => {
+                    let mut parent = match stack.pop() {
+                        Some(c) => c,
+                        None => break,
+                    };
+                    if output.starts_with(&wasm_encoder::Component::HEADER) {
+                        parent.push(ComponentSectionId::Component as u8);
+                        output.encode(&mut parent);
+                    } else {
+                        parent.push(ComponentSectionId::CoreModule as u8);
+                        output.encode(&mut parent);
+                    }
+                    output = parent;
+                }
+                _ => {}
+            }
+
             match &payload {
                 CustomSection(c) => {
                     if strip_custom_section(c.name()) {
@@ -59,15 +91,16 @@ impl Opts {
                 _ => {}
             }
             if let Some((id, range)) = payload.as_section() {
-                module.section(&RawSection {
+                RawSection {
                     id,
                     data: &input[range],
-                });
+                }
+                .append_to(&mut output);
             }
         }
 
         self.io.output(wasm_tools::Output::Wasm {
-            bytes: module.as_slice(),
+            bytes: &output,
             wat: self.wat,
         })?;
         Ok(())

tests/cli/add-metadata-merge-sections.wat

@@ -0,0 +1,2 @@
+;; RUN: metadata add --language foo % | metadata add --language bar | metadata add --sdk foo=2 | metadata show
+(module)

tests/cli/add-metadata-merge-sections.wat.stdout

@@ -0,0 +1,6 @@
+module:
+    language:
+        foo
+        bar
+    sdk:
+        foo: 2

tests/cli/add-metadata-overwrite-name.wat

@@ -0,0 +1,2 @@
+;; RUN: metadata add % --name foo | metadata show
+(module $bar)

tests/cli/add-metadata-overwrite-name.wat.stdout

@@ -0,0 +1 @@
+module foo:

tests/cli/add-metadata.wat

@@ -0,0 +1,2 @@
+;; RUN: metadata add % --name foo --language bar --processed-by baz=1 --sdk my-sdk=2 | metadata show
+(module)

tests/cli/add-metadata.wat.stdout

@@ -0,0 +1,7 @@
+module foo:
+    language:
+        bar
+    processed-by:
+        baz: 1
+    sdk:
+        my-sdk: 2