1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

//! A tiny crate of utilities for working with implicit Wasm codegen conventions
//! (often established by LLVM and lld).
//!
//! Examples conventions include:
//!
//! * The shadow stack pointer
//! * The canonical linear memory that contains the shadow stack

use anyhow::{anyhow, bail, Result};
use walrus::{
    ir::Value, ElementId, FunctionId, GlobalId, GlobalKind, InitExpr, MemoryId, Module, ValType,
};

/// Get a Wasm module's canonical linear memory.
pub fn get_memory(module: &Module) -> Result<MemoryId> {
    let mut memories = module.memories.iter().map(|m| m.id());
    let memory = memories.next();
    if memories.next().is_some() {
        bail!(
            "expected a single memory, found multiple; multiple memories \
             currently not supported"
        );
    }
    memory.ok_or_else(|| {
        anyhow!(
            "module does not have a memory; must have a memory \
             to transform return pointers into Wasm multi-value"
        )
    })
}

/// Get the `__shadow_stack_pointer`.
///
/// It must have been previously added to the module's exports via
/// `export_shadow_stack_pointer`.
pub fn get_shadow_stack_pointer(module: &Module) -> Option<GlobalId> {
    let candidates = module
        .globals
        .iter()
        .filter(|g| g.ty == ValType::I32)
        .filter(|g| g.mutable)
        // The stack pointer is guaranteed to not be initialized to 0, and it's
        // guaranteed to have an i32 initializer, so find globals which are
        // locally defined, are an i32, and have a nonzero initializer
        .filter(|g| match g.kind {
            GlobalKind::Local(InitExpr::Value(Value::I32(n))) => n != 0,
            _ => false,
        })
        .collect::<Vec<_>>();

    match candidates.len() {
        0 => None,
        // TODO: have an actual check here.
        1 => Some(candidates[0].id()),
        _ => None,
    }
}

pub struct FunctionTableEntry {
    pub element: ElementId,
    pub idx: usize,
    pub func: Option<FunctionId>,
}

/// Looks up a function table entry by index in the main function table.
pub fn get_function_table_entry(module: &Module, idx: u32) -> Result<FunctionTableEntry> {
    let table = module
        .tables
        .main_function_table()?
        .ok_or_else(|| anyhow!("no function table found in module"))?;
    let table = module.tables.get(table);
    for &segment in table.elem_segments.iter() {
        let segment = module.elements.get(segment);
        let offset = match &segment.kind {
            walrus::ElementKind::Active {
                offset: InitExpr::Value(Value::I32(n)),
                ..
            } => *n as u32,
            _ => continue,
        };
        let idx = (idx - offset) as usize;
        match segment.members.get(idx) {
            Some(slot) => {
                return Ok(FunctionTableEntry {
                    element: segment.id(),
                    idx,
                    func: slot.clone(),
                })
            }
            None => continue,
        }
    }
    bail!("failed to find `{}` in function table", idx);
}