1 //! The initial instrumentation pass.
2
3 use crate::info::ModuleContext;
4 use wasm_encoder::SectionId;
5 use wasm_encoder::reencode::{Reencode, RoundtripReencoder};
6
7 /// Instrument the input Wasm so that it exports its memories and globals,
8 /// allowing us to inspect their state after the module is instantiated and
9 /// initialized.
10 ///
11 /// For example, given this input module:
12 ///
13 /// ```wat
14 /// (module $A
15 /// (module $B
16 /// (memory $B_mem)
17 /// (global $B_glob (mut i32))
18 /// )
19 ///
20 /// (instance $x (instantiate $B))
21 /// (instance $y (instantiate $B))
22 ///
23 /// (memory $A_mem)
24 /// (global $A_glob (mut i32))
25 /// )
26 /// ```
27 ///
28 /// this pass will produce the following instrumented module:
29 ///
30 /// ```wat
31 /// (module $A
32 /// (module $B
33 /// (memory $B_mem)
34 /// (global $B_glob (mut i32))
35 ///
36 /// ;; Export all state.
37 /// (export "__wizer_memory_0" (memory $B_mem))
38 /// (export "__wizer_global_0" (global $B_glob))
39 /// )
40 ///
41 /// (instance $x (instantiate $B))
42 /// (instance $y (instantiate $B))
43 ///
44 /// (memory $A_mem)
45 /// (global $A_glob (mut i32))
46 ///
47 /// ;; Export of all state (including transitively re-exporting nested
48 /// ;; instantiations' state).
49 /// (export "__wizer_memory_0" (memory $A_mem))
50 /// (export "__wizer_global_0" (global $A_glob))
51 /// (export "__wizer_instance_0" (instance $x))
52 /// (export "__wizer_instance_1" (instance $y))
53 /// )
54 /// ```
55 ///
56 /// NB: we re-export nested instantiations as a whole instance export because we
57 /// can do this without disturbing existing instances' indices. If we were to
58 /// export their memories and globals individually, that would disturb the
59 /// modules locally defined memoryies' and globals' indices, which would require
60 /// rewriting the code section, which would break debug info offsets.
instrument(module: &mut ModuleContext<'_>) -> Vec<u8>61 pub(crate) fn instrument(module: &mut ModuleContext<'_>) -> Vec<u8> {
62 log::debug!("Instrumenting the input Wasm");
63
64 let mut encoder = wasm_encoder::Module::new();
65 let mut defined_global_exports = Vec::new();
66 let mut defined_memory_exports = Vec::new();
67
68 for section in module.raw_sections() {
69 match section.id {
70 // For the exports section, we need to transitively export internal
71 // state so that we can read the initialized state after we call the
72 // initialization function.
73 id if id == u8::from(SectionId::Export) => {
74 let mut exports = wasm_encoder::ExportSection::new();
75
76 // First, copy over all the original exports.
77 for export in module.exports() {
78 RoundtripReencoder
79 .parse_export(&mut exports, *export)
80 .unwrap();
81 }
82
83 // Now export all of this module's defined globals, memories,
84 // and instantiations under well-known names so we can inspect
85 // them after initialization.
86 for (i, ty, _) in module.defined_globals() {
87 if !ty.mutable {
88 continue;
89 }
90 let name = format!("__wizer_global_{i}");
91 exports.export(&name, wasm_encoder::ExportKind::Global, i);
92 defined_global_exports.push((i, name));
93 }
94 for (i, (j, _)) in module.defined_memories().enumerate() {
95 let name = format!("__wizer_memory_{i}");
96 exports.export(&name, wasm_encoder::ExportKind::Memory, j);
97 defined_memory_exports.push(name);
98 }
99
100 encoder.section(&exports);
101 }
102
103 // All other sections don't need instrumentation and can be copied
104 // over directly.
105 _other => {
106 encoder.section(section);
107 }
108 }
109 }
110
111 module.defined_global_exports = Some(defined_global_exports);
112 module.defined_memory_exports = Some(defined_memory_exports);
113
114 encoder.finish()
115 }
116