ic_wasm/
utils.rs

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

ic_wasm/
utils.rs
