1 use crate::component::ComponentContext;
2 use std::collections::hash_map::Entry;
3 use wasmparser::{
4     CanonicalFunction, ComponentAlias, ComponentExternalKind, ComponentOuterAliasKind, Encoding,
5     Instance, Parser, Payload,
6 };
7 use wasmtime::{Result, bail};
8 
9 /// Parse the given Wasm bytes into a `ComponentContext` tree.
10 ///
11 /// This will parse the input component and build up metadata that Wizer needs
12 /// to know about the component. At this time there are limitations to this
13 /// parsing phase which in theory could be lifted in the future but serve as
14 /// simplifying assumptions for now:
15 ///
16 /// * Nested components with modules are not supported.
17 /// * Imported modules or components are not supported.
18 /// * Instantiating a module twice is not supported.
19 /// * Component-level start functions are not supported.
20 ///
21 /// Some of these restrictions are likely to be loosened over time, however.
22 pub(crate) fn parse<'a>(full_wasm: &'a [u8]) -> wasmtime::Result<ComponentContext<'a>> {
23     let mut component = ComponentContext::default();
24     let parser = Parser::new(0).parse_all(full_wasm);
25     parse_into(Some(&mut component), full_wasm, parser)?;
26     Ok(component)
27 }
28 
29 fn parse_into<'a>(
30     mut cx: Option<&mut ComponentContext<'a>>,
31     full_wasm: &'a [u8],
32     mut iter: impl Iterator<Item = wasmparser::Result<Payload<'a>>>,
33 ) -> wasmtime::Result<()> {
34     let mut stack = Vec::new();
35     while let Some(payload) = iter.next() {
36         let payload = payload?;
37 
38         match &payload {
39             // Module sections get parsed with wizer's core wasm support.
40             Payload::ModuleSection { .. } => match &mut cx {
41                 Some(component) => {
42                     let info = crate::parse::parse_with(&full_wasm, &mut iter)?;
43                     component.push_module_section(info);
44                 }
45                 None => {
46                     bail!("nested components with modules not currently supported");
47                 }
48             },
49 
50             // All other sections get pushed raw as-is into the component.
51             _ => {
52                 if let Some((id, range)) = payload.as_section()
53                     && let Some(component) = &mut cx
54                 {
55                     component.push_raw_section(wasm_encoder::RawSection {
56                         id,
57                         data: &full_wasm[range],
58                     });
59                 }
60             }
61         }
62 
63         // Further validation/handling of each section, mostly about maintaining
64         // index spaces so we know what indices are used when the instrumented
65         // component is produced.
66         match payload {
67             Payload::Version {
68                 encoding: Encoding::Module,
69                 ..
70             } => {
71                 bail!("expected a component, found a core module");
72             }
73 
74             Payload::ComponentSection { .. } => {
75                 stack.push(cx.take());
76             }
77 
78             Payload::End(_) => {
79                 if stack.len() > 0 {
80                     cx = stack.pop().unwrap();
81                 }
82             }
83 
84             Payload::InstanceSection(reader) => {
85                 if let Some(component) = &mut cx {
86                     for instance in reader {
87                         let instance_index = component.inc_core_instances();
88 
89                         if let Instance::Instantiate { module_index, .. } = instance? {
90                             match component.core_instantiations.entry(module_index) {
91                                 Entry::Vacant(entry) => {
92                                     entry.insert(instance_index);
93                                 }
94                                 Entry::Occupied(_) => {
95                                     bail!("modules may be instantiated at most once")
96                                 }
97                             }
98                         }
99                     }
100                 }
101             }
102             Payload::ComponentInstanceSection(reader) => {
103                 if let Some(component) = &mut cx {
104                     for _ in reader {
105                         component.inc_instances();
106                     }
107                 }
108             }
109 
110             Payload::ComponentAliasSection(reader) => {
111                 for alias in reader {
112                     match alias? {
113                         ComponentAlias::CoreInstanceExport { kind, .. } => {
114                             if let Some(component) = &mut cx {
115                                 component.inc_core(kind);
116                             }
117                         }
118                         ComponentAlias::InstanceExport { kind, .. } => {
119                             validate_item_kind(kind, "aliases")?;
120                             if let Some(component) = &mut cx {
121                                 component.inc(kind);
122                             }
123                         }
124                         ComponentAlias::Outer { kind, .. } => match kind {
125                             ComponentOuterAliasKind::CoreType => {}
126                             ComponentOuterAliasKind::Type => {
127                                 if let Some(component) = &mut cx {
128                                     component.inc_types();
129                                 }
130                             }
131                             ComponentOuterAliasKind::CoreModule => {
132                                 bail!("wizer does not currently support module aliases");
133                             }
134                             ComponentOuterAliasKind::Component => {
135                                 bail!("wizer does not currently support component aliases");
136                             }
137                         },
138                     }
139                 }
140             }
141 
142             Payload::ComponentCanonicalSection(reader) => {
143                 for function in reader {
144                     match function? {
145                         CanonicalFunction::Lift { .. } => {
146                             if let Some(component) = &mut cx {
147                                 component.inc_funcs();
148                             }
149                         }
150                         _ => {
151                             if let Some(component) = &mut cx {
152                                 component.inc_core_funcs();
153                             }
154                         }
155                     }
156                 }
157             }
158 
159             Payload::ComponentImportSection(reader) => {
160                 for import in reader {
161                     let kind = import?.ty.kind();
162                     validate_item_kind(kind, "imports")?;
163                     if let Some(component) = &mut cx {
164                         component.inc(kind);
165                     }
166                 }
167             }
168 
169             Payload::ComponentExportSection(reader) => {
170                 for export in reader {
171                     let kind = export?.kind;
172                     validate_item_kind(kind, "exports")?;
173                     if let Some(component) = &mut cx {
174                         component.inc(kind);
175                     }
176                 }
177             }
178 
179             Payload::ComponentTypeSection(reader) => {
180                 for _ in reader {
181                     if let Some(component) = &mut cx {
182                         component.inc_types();
183                     }
184                 }
185             }
186 
187             // The `start` section for components is itself not stable, so not
188             // super urgent to handle. A simplifying assumption is made to just
189             // reject it outright for now if it shows up.
190             Payload::ComponentStartSection { .. } => {
191                 bail!("wizer does not currently support component start functions");
192             }
193 
194             _ => {}
195         }
196     }
197 
198     Ok(())
199 }
200 
201 fn validate_item_kind(kind: ComponentExternalKind, msg: &str) -> Result<()> {
202     match kind {
203         // Aliasing modules would require keeping track of where a module's
204         // original definition is. There's not much usage of this in the wild
205         // so reject it for now as a simplifying assumption.
206         ComponentExternalKind::Module => {
207             bail!("wizer does not currently support module {msg}");
208         }
209 
210         // Aliasing components deals with nested instantiations or similar,
211         // where a simplifying assumption is made to not worry about that for now.
212         ComponentExternalKind::Component => {
213             bail!("wizer does not currently support component {msg}");
214         }
215 
216         // Like start sections, support for this is just deferred to some other
217         // time, if ever.
218         ComponentExternalKind::Value => {
219             bail!("wizer does not currently support value {msg}");
220         }
221 
222         ComponentExternalKind::Func
223         | ComponentExternalKind::Type
224         | ComponentExternalKind::Instance => Ok(()),
225     }
226 }
227