Skip to main content

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::*;
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
14/// WebAssembly features accepted by the IC replica, used to sanity-check the
15/// module ic-wasm emits (a mismatch only produces a warning).
16///
17/// The feature set mirrors the IC `wasmtime` validation config:
18/// <https://github.com/dfinity/ic/blob/master/rs/embedders/src/wasm_utils/validation.rs>
19/// (function `wasmtime_validation_config`).
20///
21/// The replica validates canister modules with wasmtime, starting from
22/// `wasmtime::Config::default()` and disabling a handful of proposals (see
23/// `wasmtime_validation_config` linked above). wasmparser's
24/// `WasmFeatures::default()` tracks the same wasmtime-default set of enabled
25/// proposals, so we mirror the replica by starting from that default and
26/// removing exactly the proposals it turns off. New proposals then inherit the
27/// shared default, instead of an exhaustive struct literal that breaks the
28/// build every time wasmparser adds a feature flag.
29pub fn make_validator_with_features() -> Validator {
30    let mut features = WasmFeatures::default();
31    features.remove(
32        // config.wasm_function_references(false)
33        WasmFeatures::FUNCTION_REFERENCES
34        // config.wasm_gc(false)
35            | WasmFeatures::GC
36        // config.wasm_multi_memory(false) (disabled during validation)
37            | WasmFeatures::MULTI_MEMORY
38        // config.wasm_relaxed_simd(false) (disabled for determinism)
39            | WasmFeatures::RELAXED_SIMD
40        // config.wasm_extended_const(false)
41            | 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            // System API cost taken from https://github.com/dfinity/ic/blob/master/rs/embedders/src/wasmtime_embedder/system_api_complexity.rs
107            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    // Cost taken from https://github.com/dfinity/ic/blob/master/rs/embedders/src/wasm_utils/instrumentation.rs
163    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, // missing ReturnCall/Indirect
176        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, // missing TruncSat?
187            _ => 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]  // tag for blob
342        && blob[8..12] == [0x00, 0x61, 0x73, 0x6d]
343    // Wasm magic number
344    {
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}