1use crate::info::ExportedMethodInfo;
2use crate::Error;
3use libflate::gzip;
4use std::borrow::Cow;
5use std::collections::HashMap;
6use std::io::{self, Read};
7use walrus::*;
8use wasmparser::{Validator, WasmFeatures};
9
10pub const WASM_MAGIC_BYTES: &[u8] = &[0, 97, 115, 109];
11
12pub const GZIPPED_WASM_MAGIC_BYTES: &[u8] = &[31, 139, 8];
13
14pub fn make_validator_with_features() -> Validator {
30 let mut features = WasmFeatures::default();
31 features.remove(
32 WasmFeatures::FUNCTION_REFERENCES
34 | WasmFeatures::GC
36 | WasmFeatures::MULTI_MEMORY
38 | WasmFeatures::RELAXED_SIMD
40 | WasmFeatures::EXTENDED_CONST,
42 );
43 Validator::new_with_features(features)
44}
45
46fn wasm_parser_config(keep_name_section: bool) -> ModuleConfig {
47 let mut config = walrus::ModuleConfig::new();
48 config.generate_name_section(keep_name_section);
49 config.generate_producers_section(false);
50 config
51}
52
53fn decompress(bytes: &[u8]) -> Result<Vec<u8>, std::io::Error> {
54 let mut decoder = gzip::Decoder::new(bytes)?;
55 let mut decoded_data = Vec::new();
56 decoder.read_to_end(&mut decoded_data)?;
57 Ok(decoded_data)
58}
59
60pub fn parse_wasm(bytes: &[u8], keep_name_section: bool) -> Result<Module, Error> {
61 let wasm = if bytes.starts_with(WASM_MAGIC_BYTES) {
62 Ok(Cow::Borrowed(bytes))
63 } else if bytes.starts_with(GZIPPED_WASM_MAGIC_BYTES) {
64 decompress(bytes).map(Cow::Owned)
65 } else {
66 Err(io::Error::new(
67 io::ErrorKind::InvalidInput,
68 "Input must be either gzipped or uncompressed WASM.",
69 ))
70 }
71 .map_err(Error::IO)?;
72 let config = wasm_parser_config(keep_name_section);
73 config
74 .parse(&wasm)
75 .map_err(|e| Error::WasmParse(e.to_string()))
76}
77
78pub fn parse_wasm_file(file: std::path::PathBuf, keep_name_section: bool) -> Result<Module, Error> {
79 let bytes = std::fs::read(file).map_err(Error::IO)?;
80 parse_wasm(&bytes[..], keep_name_section)
81}
82
83#[derive(Clone, Copy, PartialEq, Eq)]
84pub(crate) enum InjectionKind {
85 Static,
86 Dynamic,
87 Dynamic64,
88}
89
90pub(crate) struct FunctionCost(HashMap<FunctionId, (i64, InjectionKind)>);
91impl FunctionCost {
92 pub fn new(m: &Module) -> Self {
93 let mut res = HashMap::new();
94 for (method, func) in m.imports.iter().filter_map(|i| {
95 if let ImportKind::Function(func) = i.kind {
96 if i.module == "ic0" {
97 Some((i.name.as_str(), func))
98 } else {
99 None
100 }
101 } else {
102 None
103 }
104 }) {
105 use InjectionKind::*;
106 let cost = match method {
108 "accept_message" => (500, Static),
109 "call_cycles_add" | "call_cycles_add128" => (500, Static),
110 "call_data_append" => (500, Dynamic),
111 "call_new" => (1500, Static),
112 "call_on_cleanup" => (500, Static),
113 "call_perform" => (5000, Static),
114 "canister_cycle_balance" | "canister_cycle_balance128" => (500, Static),
115 "canister_self_copy" => (500, Dynamic),
116 "canister_self_size" => (500, Static),
117 "canister_status" | "canister_version" => (500, Static),
118 "certified_data_set" => (500, Dynamic),
119 "data_certificate_copy" => (500, Dynamic),
120 "data_certificate_present" | "data_certificate_size" => (500, Static),
121 "debug_print" => (100, Dynamic),
122 "global_timer_set" => (500, Static),
123 "is_controller" => (1000, Dynamic),
124 "msg_arg_data_copy" => (500, Dynamic),
125 "msg_arg_data_size" => (500, Static),
126 "msg_caller_copy" => (500, Dynamic),
127 "msg_caller_size" => (500, Static),
128 "msg_cycles_accept" | "msg_cycles_accept128" => (500, Static),
129 "msg_cycles_available" | "msg_cycles_available128" => (500, Static),
130 "msg_cycles_refunded" | "msg_cycles_refunded128" => (500, Static),
131 "cycles_burn128" => (100, Static),
132 "msg_method_name_copy" => (500, Dynamic),
133 "msg_method_name_size" => (500, Static),
134 "msg_reject_code" | "msg_reject_msg_size" => (500, Static),
135 "msg_reject_msg_copy" => (500, Dynamic),
136 "msg_reject" => (500, Dynamic),
137 "msg_reply_data_append" => (500, Dynamic),
138 "msg_reply" => (500, Static),
139 "performance_counter" => (200, Static),
140 "stable_grow" | "stable64_grow" => (100, Static),
141 "stable_size" | "stable64_size" => (20, Static),
142 "stable_read" => (20, Dynamic),
143 "stable_write" => (20, Dynamic),
144 "stable64_read" => (20, Dynamic64),
145 "stable64_write" => (20, Dynamic64),
146 "trap" => (500, Dynamic),
147 "time" => (500, Static),
148 _ => (20, Static),
149 };
150 res.insert(func, cost);
151 }
152 Self(res)
153 }
154 pub fn get_cost(&self, id: FunctionId) -> Option<(i64, InjectionKind)> {
155 self.0.get(&id).copied()
156 }
157}
158pub(crate) fn instr_cost(i: &ir::Instr) -> i64 {
159 use ir::*;
160 use BinaryOp::*;
161 use UnaryOp::*;
162 match i {
164 Instr::Block(..) | Instr::Loop(..) => 0,
165 Instr::Const(..) | Instr::Load(..) | Instr::Store(..) => 1,
166 Instr::GlobalGet(..) | Instr::GlobalSet(..) => 2,
167 Instr::TableGet(..) | Instr::TableSet(..) => 5,
168 Instr::TableGrow(..) | Instr::MemoryGrow(..) => 300,
169 Instr::MemorySize(..) => 20,
170 Instr::TableSize(..) => 100,
171 Instr::MemoryFill(..) | Instr::MemoryCopy(..) | Instr::MemoryInit(..) => 100,
172 Instr::TableFill(..) | Instr::TableCopy(..) | Instr::TableInit(..) => 100,
173 Instr::DataDrop(..) | Instr::ElemDrop(..) => 300,
174 Instr::Call(..) => 5,
175 Instr::CallIndirect(..) => 10, Instr::IfElse(..) | Instr::Br(..) | Instr::BrIf(..) | Instr::BrTable(..) => 2,
177 Instr::RefIsNull(..) => 5,
178 Instr::RefFunc(..) => 130,
179 Instr::Unop(Unop { op }) => match op {
180 F32Ceil | F32Floor | F32Trunc | F32Nearest | F32Sqrt => 20,
181 F64Ceil | F64Floor | F64Trunc | F64Nearest | F64Sqrt => 20,
182 F32Abs | F32Neg | F64Abs | F64Neg => 2,
183 F32ConvertSI32 | F64ConvertSI64 | F32ConvertSI64 | F64ConvertSI32 => 3,
184 F64ConvertUI32 | F32ConvertUI64 | F32ConvertUI32 | F64ConvertUI64 => 16,
185 I64TruncSF32 | I64TruncUF32 | I64TruncSF64 | I64TruncUF64 => 20,
186 I32TruncSF32 | I32TruncUF32 | I32TruncSF64 | I32TruncUF64 => 20, _ => 1,
188 },
189 Instr::Binop(Binop { op }) => match op {
190 I32DivS | I32DivU | I32RemS | I32RemU => 10,
191 I64DivS | I64DivU | I64RemS | I64RemU => 10,
192 F32Add | F32Sub | F32Mul | F32Div | F32Min | F32Max => 20,
193 F64Add | F64Sub | F64Mul | F64Div | F64Min | F64Max => 20,
194 F32Copysign | F64Copysign => 2,
195 F32Eq | F32Ne | F32Lt | F32Gt | F32Le | F32Ge => 3,
196 F64Eq | F64Ne | F64Lt | F64Gt | F64Le | F64Ge => 3,
197 _ => 1,
198 },
199 _ => 1,
200 }
201}
202
203pub(crate) fn get_ic_func_id(m: &mut Module, method: &str) -> FunctionId {
204 match m.imports.find("ic0", method) {
205 Some(id) => match m.imports.get(id).kind {
206 ImportKind::Function(func_id) => func_id,
207 _ => unreachable!(),
208 },
209 None => {
210 let ty = match method {
211 "stable_write" => m
212 .types
213 .add(&[ValType::I32, ValType::I32, ValType::I32], &[]),
214 "stable64_write" => m
215 .types
216 .add(&[ValType::I64, ValType::I64, ValType::I64], &[]),
217 "stable_read" => m
218 .types
219 .add(&[ValType::I32, ValType::I32, ValType::I32], &[]),
220 "stable64_read" => m
221 .types
222 .add(&[ValType::I64, ValType::I64, ValType::I64], &[]),
223 "stable_grow" => m.types.add(&[ValType::I32], &[ValType::I32]),
224 "stable64_grow" => m.types.add(&[ValType::I64], &[ValType::I64]),
225 "stable_size" => m.types.add(&[], &[ValType::I32]),
226 "stable64_size" => m.types.add(&[], &[ValType::I64]),
227 "call_cycles_add" => m.types.add(&[ValType::I64], &[]),
228 "call_cycles_add128" => m.types.add(&[ValType::I64, ValType::I64], &[]),
229 "cycles_burn128" => m
230 .types
231 .add(&[ValType::I64, ValType::I64, ValType::I32], &[]),
232 "call_new" => m.types.add(
233 &[
234 ValType::I32,
235 ValType::I32,
236 ValType::I32,
237 ValType::I32,
238 ValType::I32,
239 ValType::I32,
240 ValType::I32,
241 ValType::I32,
242 ],
243 &[],
244 ),
245 "debug_print" => m.types.add(&[ValType::I32, ValType::I32], &[]),
246 "trap" => m.types.add(&[ValType::I32, ValType::I32], &[]),
247 "msg_arg_data_size" => m.types.add(&[], &[ValType::I32]),
248 "msg_arg_data_copy" => m
249 .types
250 .add(&[ValType::I32, ValType::I32, ValType::I32], &[]),
251 "msg_reply_data_append" => m.types.add(&[ValType::I32, ValType::I32], &[]),
252 "msg_reply" => m.types.add(&[], &[]),
253 _ => unreachable!(),
254 };
255 m.add_import_func("ic0", method, ty).0
256 }
257 }
258}
259
260pub(crate) fn get_memory_id(m: &Module) -> MemoryId {
261 m.memories
262 .iter()
263 .next()
264 .expect("only single memory is supported")
265 .id()
266}
267
268pub(crate) fn get_export_func_id(m: &Module, method: &str) -> Option<FunctionId> {
269 let e = m.exports.iter().find(|e| e.name == method)?;
270 if let ExportItem::Function(id) = e.item {
271 Some(id)
272 } else {
273 None
274 }
275}
276pub(crate) fn get_or_create_export_func<'a>(
277 m: &'a mut Module,
278 method: &'a str,
279) -> InstrSeqBuilder<'a> {
280 let id = match get_export_func_id(m, method) {
281 Some(id) => id,
282 None => {
283 let builder = FunctionBuilder::new(&mut m.types, &[], &[]);
284 let id = builder.finish(vec![], &mut m.funcs);
285 m.exports.add(method, id);
286 id
287 }
288 };
289 get_builder(m, id)
290}
291
292pub(crate) fn get_builder(m: &mut Module, id: FunctionId) -> InstrSeqBuilder<'_> {
293 if let FunctionKind::Local(func) = &mut m.funcs.get_mut(id).kind {
294 let id = func.entry_block();
295 func.builder_mut().instr_seq(id)
296 } else {
297 unreachable!()
298 }
299}
300
301pub(crate) fn inject_top(builder: &mut InstrSeqBuilder<'_>, instrs: Vec<ir::Instr>) {
302 for instr in instrs.into_iter().rev() {
303 builder.instr_at(0, instr);
304 }
305}
306
307pub(crate) fn get_exported_methods(m: &Module) -> Vec<ExportedMethodInfo> {
308 m.exports
309 .iter()
310 .filter_map(|e| match e.item {
311 ExportItem::Function(id) => Some(ExportedMethodInfo {
312 name: e.name.clone(),
313 internal_name: get_func_name(m, id),
314 }),
315 _ => None,
316 })
317 .collect()
318}
319
320pub(crate) fn get_func_name(m: &Module, id: FunctionId) -> String {
321 m.funcs
322 .get(id)
323 .name
324 .as_ref()
325 .unwrap_or(&format!("func_{}", id.index()))
326 .to_string()
327}
328
329pub(crate) fn is_motoko_canister(m: &Module) -> bool {
330 m.customs.iter().any(|(_, s)| {
331 s.name() == "icp:private motoko:compiler" || s.name() == "icp:public motoko:compiler"
332 }) || m
333 .exports
334 .iter()
335 .any(|e| e.name == "canister_update __motoko_async_helper")
336}
337
338pub(crate) fn is_motoko_wasm_data_section(blob: &[u8]) -> Option<&[u8]> {
339 let len = blob.len() as u32;
340 if len > 100
341 && blob[0..4] == [0x11, 0x00, 0x00, 0x00] && blob[8..12] == [0x00, 0x61, 0x73, 0x6d]
343 {
345 let decoded_len = u32::from_le_bytes(blob[4..8].try_into().unwrap());
346 if decoded_len + 8 == len {
347 return Some(&blob[8..]);
348 }
349 }
350 None
351}
352
353pub(crate) fn get_motoko_wasm_data_sections(m: &Module) -> Vec<(DataId, Module)> {
354 m.data
355 .iter()
356 .filter_map(|d| {
357 let blob = is_motoko_wasm_data_section(&d.value)?;
358 let mut config = ModuleConfig::new();
359 config.generate_name_section(false);
360 config.generate_producers_section(false);
361 let m = config.parse(blob).ok()?;
362 Some((d.id(), m))
363 })
364 .collect()
365}
366
367pub(crate) fn encode_module_as_data_section(mut m: Module) -> Vec<u8> {
368 let blob = m.emit_wasm();
369 let blob_len = blob.len();
370 let mut res = Vec::with_capacity(blob_len + 8);
371 res.extend_from_slice(&[0x11, 0x00, 0x00, 0x00]);
372 let encoded_len = (blob_len as u32).to_le_bytes();
373 res.extend_from_slice(&encoded_len);
374 res.extend_from_slice(&blob);
375 res
376}