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