/src/wasm-tools/crates/wasm-metadata/src/rewrite.rs

Source
use crate::add_metadata::AddMetadataField;
use crate::{AddMetadata, ComponentNames, ModuleNames, Producers};
use anyhow::Result;
use std::mem;
use wasm_encoder::ComponentSection as _;
use wasm_encoder::{ComponentSectionId, Encode, Section};
use wasmparser::{KnownCustom, Parser, Payload::*};

pub(crate) fn rewrite_wasm(
    metadata: &AddMetadata,
    add_producers: &Producers,
    input: &[u8],
) -> Result<Vec<u8>> {
    let mut producers_found = false;
    let mut names_found = false;
    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 {
            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;
            }
            _ => {}
        }

        // 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())?;
                        match &metadata.name {
                            AddMetadataField::Keep => {}
                            AddMetadataField::Set(name) => {
                                names.merge(&ModuleNames::from_name(&Some(name.clone())));
                            }
                            AddMetadataField::Clear => {
                                names = ModuleNames::empty();
                            }
                        };

                        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())?;
                        match &metadata.name {
                            AddMetadataField::Keep => {}
                            AddMetadataField::Set(name) => {
                                names.merge(&ComponentNames::from_name(&Some(name.clone())));
                            }
                            AddMetadataField::Clear => {
                                names = ComponentNames::empty();
                            }
                        };
                        names.section()?.as_custom().append_to(&mut output);
                        continue;
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "authors" => {
                        if metadata.authors.is_keep() {
                            let authors = crate::Authors::parse_custom_section(c)?;
                            authors.append_to(&mut output);
                            continue;
                        } else if metadata.authors.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "description" => {
                        if metadata.description.is_keep() {
                            let description = crate::Description::parse_custom_section(c)?;
                            description.append_to(&mut output);
                            continue;
                        } else if metadata.description.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "licenses" => {
                        if metadata.licenses.is_keep() {
                            let licenses = crate::Licenses::parse_custom_section(c)?;
                            licenses.append_to(&mut output);
                            continue;
                        } else if metadata.licenses.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "source" => {
                        if metadata.source.is_keep() {
                            let source = crate::Source::parse_custom_section(c)?;
                            source.append_to(&mut output);
                            continue;
                        } else if metadata.source.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "homepage" => {
                        if metadata.homepage.is_keep() {
                            let homepage = crate::Homepage::parse_custom_section(c)?;
                            homepage.append_to(&mut output);
                            continue;
                        } else if metadata.homepage.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "revision" => {
                        if metadata.revision.is_keep() {
                            let revision = crate::Revision::parse_custom_section(c)?;
                            revision.append_to(&mut output);
                            continue;
                        } else if metadata.revision.is_clear() {
                            continue;
                        }
                    }
                    #[cfg(feature = "oci")]
                    KnownCustom::Unknown if c.name() == "version" => {
                        if metadata.version.is_keep() {
                            let version = crate::Version::parse_custom_section(c)?;
                            version.append_to(&mut output);
                            continue;
                        } else if metadata.version.is_clear() {
                            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 {
        if let AddMetadataField::Set(name) = &metadata.name {
            if output.starts_with(&wasm_encoder::Component::HEADER) {
                let names = ComponentNames::from_name(&Some(name.clone()));
                names.section()?.append_to_component(&mut output);
            } else {
                let names = ModuleNames::from_name(&Some(name.clone()));
                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:
        producers.section().append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(authors) = &metadata.authors {
        authors.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(description) = &metadata.description {
        description.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(licenses) = &metadata.licenses {
        licenses.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(source) = &metadata.source {
        source.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(homepage) = &metadata.homepage {
        homepage.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(revision) = &metadata.revision {
        revision.append_to(&mut output);
    }
    #[cfg(feature = "oci")]
    if let AddMetadataField::Set(version) = &metadata.version {
        version.append_to(&mut output);
    }
    Ok(output)
}

Coverage Report

Created: 2025-12-04 07:01

Line	Count	Source
1		use crate::add_metadata::AddMetadataField;
2		use crate::{AddMetadata, ComponentNames, ModuleNames, Producers};
3		use anyhow::Result;
4		use std::mem;
5		use wasm_encoder::ComponentSection as _;
6		use wasm_encoder::{ComponentSectionId, Encode, Section};
7		use wasmparser::{KnownCustom, Parser, Payload::*};
8
9	4.07k	pub(crate) fn rewrite_wasm(
10	4.07k	metadata: &AddMetadata,
11	4.07k	add_producers: &Producers,
12	4.07k	input: &[u8],
13	4.07k	) -> Result<Vec<u8>> {
14	4.07k	let mut producers_found = false;
15	4.07k	let mut names_found = false;
16	4.07k	let mut stack = Vec::new();
17	4.07k	let mut output = Vec::new();
18	40.4k	for payload in Parser::new(0).parse_all(&input) {
19	40.4k	let payload = payload?;
20
21		// Track nesting depth, so that we don't mess with inner producer sections:
22	40.4k	match payload {
23	4.07k	Version { encoding, .. } => {
24	4.07k	output.extend_from_slice(match encoding {
25	0	wasmparser::Encoding::Component => &wasm_encoder::Component::HEADER,
26	4.07k	wasmparser::Encoding::Module => &wasm_encoder::Module::HEADER,
27		});
28		}
29		ModuleSection { .. } \| ComponentSection { .. } => {
30	0	stack.push(mem::take(&mut output));
31	0	continue;
32		}
33		End { .. } => {
34	4.07k	let mut parent = match stack.pop() {
35	0	Some(c) => c,
36	4.07k	None => break,
37		};
38	0	if output.starts_with(&wasm_encoder::Component::HEADER) {
39	0	parent.push(ComponentSectionId::Component as u8);
40	0	output.encode(&mut parent);
41	0	} else {
42	0	parent.push(ComponentSectionId::CoreModule as u8);
43	0	output.encode(&mut parent);
44	0	}
45	0	output = parent;
46		}
47	32.2k	_ => {}
48		}
49
50		// Only rewrite the outermost custom sections
51	36.3k	if let CustomSection(c) = &payload {
52	0	if stack.len() == 0 {
53	0	match c.as_known() {
54		KnownCustom::Producers(_) => {
55	0	producers_found = true;
56	0	let mut producers = Producers::from_bytes(c.data(), c.data_offset())?;
57		// Add to the section according to the command line flags:
58	0	producers.merge(&add_producers);
59		// Encode into output:
60	0	producers.section().append_to(&mut output);
61	0	continue;
62		}
63		KnownCustom::Name(_) => {
64	0	names_found = true;
65	0	let mut names = ModuleNames::from_bytes(c.data(), c.data_offset())?;
66	0	match &metadata.name {
67	0	AddMetadataField::Keep => {}
68	0	AddMetadataField::Set(name) => {
69	0	names.merge(&ModuleNames::from_name(&Some(name.clone())));
70	0	}
71	0	AddMetadataField::Clear => {
72	0	names = ModuleNames::empty();
73	0	}
74		};
75
76	0	names.section()?.as_custom().append_to(&mut output);
77	0	continue;
78		}
79		KnownCustom::ComponentName(_) => {
80	0	names_found = true;
81	0	let mut names = ComponentNames::from_bytes(c.data(), c.data_offset())?;
82	0	match &metadata.name {
83	0	AddMetadataField::Keep => {}
84	0	AddMetadataField::Set(name) => {
85	0	names.merge(&ComponentNames::from_name(&Some(name.clone())));
86	0	}
87	0	AddMetadataField::Clear => {
88	0	names = ComponentNames::empty();
89	0	}
90		};
91	0	names.section()?.as_custom().append_to(&mut output);
92	0	continue;
93		}
94		#[cfg(feature = "oci")]
95		KnownCustom::Unknown if c.name() == "authors" => {
96		if metadata.authors.is_keep() {
97		let authors = crate::Authors::parse_custom_section(c)?;
98		authors.append_to(&mut output);
99		continue;
100		} else if metadata.authors.is_clear() {
101		continue;
102		}
103		}
104		#[cfg(feature = "oci")]
105		KnownCustom::Unknown if c.name() == "description" => {
106		if metadata.description.is_keep() {
107		let description = crate::Description::parse_custom_section(c)?;
108		description.append_to(&mut output);
109		continue;
110		} else if metadata.description.is_clear() {
111		continue;
112		}
113		}
114		#[cfg(feature = "oci")]
115		KnownCustom::Unknown if c.name() == "licenses" => {
116		if metadata.licenses.is_keep() {
117		let licenses = crate::Licenses::parse_custom_section(c)?;
118		licenses.append_to(&mut output);
119		continue;
120		} else if metadata.licenses.is_clear() {
121		continue;
122		}
123		}
124		#[cfg(feature = "oci")]
125		KnownCustom::Unknown if c.name() == "source" => {
126		if metadata.source.is_keep() {
127		let source = crate::Source::parse_custom_section(c)?;
128		source.append_to(&mut output);
129		continue;
130		} else if metadata.source.is_clear() {
131		continue;
132		}
133		}
134		#[cfg(feature = "oci")]
135		KnownCustom::Unknown if c.name() == "homepage" => {
136		if metadata.homepage.is_keep() {
137		let homepage = crate::Homepage::parse_custom_section(c)?;
138		homepage.append_to(&mut output);
139		continue;
140		} else if metadata.homepage.is_clear() {
141		continue;
142		}
143		}
144		#[cfg(feature = "oci")]
145		KnownCustom::Unknown if c.name() == "revision" => {
146		if metadata.revision.is_keep() {
147		let revision = crate::Revision::parse_custom_section(c)?;
148		revision.append_to(&mut output);
149		continue;
150		} else if metadata.revision.is_clear() {
151		continue;
152		}
153		}
154		#[cfg(feature = "oci")]
155		KnownCustom::Unknown if c.name() == "version" => {
156		if metadata.version.is_keep() {
157		let version = crate::Version::parse_custom_section(c)?;
158		version.append_to(&mut output);
159		continue;
160		} else if metadata.version.is_clear() {
161		continue;
162		}
163		}
164	0	_ => {}
165		}
166	0	}
167	36.3k	}
168		// All other sections get passed through unmodified:
169	36.3k	if let Some((id, range)) = payload.as_section() {
170	21.8k	wasm_encoder::RawSection {
171	21.8k	id,
172	21.8k	data: &input[range],
173	21.8k	}
174	21.8k	.append_to(&mut output);
175	21.8k	}
176		}
177	4.07k	if !names_found {
178	4.07k	if let AddMetadataField::Set(name) = &metadata.name {
179	0	if output.starts_with(&wasm_encoder::Component::HEADER) {
180	0	let names = ComponentNames::from_name(&Some(name.clone()));
181	0	names.section()?.append_to_component(&mut output);
182		} else {
183	0	let names = ModuleNames::from_name(&Some(name.clone()));
184	0	names.section()?.append_to(&mut output)
185		}
186	4.07k	}
187	0	}
188	4.07k	if !producers_found && !add_producers.is_empty() {
189	4.07k	let mut producers = Producers::empty();
190	4.07k	// Add to the section according to the command line flags:
191	4.07k	producers.merge(add_producers);
192	4.07k	// Encode into output:
193	4.07k	producers.section().append_to(&mut output);
194	4.07k	}
195		#[cfg(feature = "oci")]
196		if let AddMetadataField::Set(authors) = &metadata.authors {
197		authors.append_to(&mut output);
198		}
199		#[cfg(feature = "oci")]
200		if let AddMetadataField::Set(description) = &metadata.description {
201		description.append_to(&mut output);
202		}
203		#[cfg(feature = "oci")]
204		if let AddMetadataField::Set(licenses) = &metadata.licenses {
205		licenses.append_to(&mut output);
206		}
207		#[cfg(feature = "oci")]
208		if let AddMetadataField::Set(source) = &metadata.source {
209		source.append_to(&mut output);
210		}
211		#[cfg(feature = "oci")]
212		if let AddMetadataField::Set(homepage) = &metadata.homepage {
213		homepage.append_to(&mut output);
214		}
215		#[cfg(feature = "oci")]
216		if let AddMetadataField::Set(revision) = &metadata.revision {
217		revision.append_to(&mut output);
218		}
219		#[cfg(feature = "oci")]
220		if let AddMetadataField::Set(version) = &metadata.version {
221		version.append_to(&mut output);
222		}
223	4.07k	Ok(output)
224	4.07k	}