wasm_metadata/
rewrite.rs

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

wasm_metadata/
rewrite.rs
