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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! 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, FunctionBuilder, FunctionId, FunctionKind, 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`.
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,
    }
}

/// Get the `__tls_base`.
pub fn get_tls_base(module: &Module) -> Option<GlobalId> {
    let candidates = module
        .exports
        .iter()
        .filter(|ex| ex.name == "__tls_base")
        .filter_map(|ex| match ex.item {
            walrus::ExportItem::Global(id) => Some(id),
            _ => None,
        })
        .filter(|id| {
            let global = module.globals.get(*id);

            global.ty == ValType::I32
        })
        .collect::<Vec<_>>();

    match candidates.len() {
        1 => Some(candidates[0]),
        _ => 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);
}

pub fn get_or_insert_start_builder(module: &mut Module) -> &mut FunctionBuilder {
    let prev_start = {
        match module.start {
            Some(start) => match module.funcs.get_mut(start).kind {
                FunctionKind::Import(_) => Err(Some(start)),
                FunctionKind::Local(_) => Ok(start),
                FunctionKind::Uninitialized(_) => todo!(),
            },
            None => Err(None),
        }
    };

    let id = match prev_start {
        Ok(id) => id,
        Err(prev_start) => {
            let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]);

            if let Some(prev_start) = prev_start {
                builder.func_body().call(prev_start);
            }

            let id = builder.finish(Vec::new(), &mut module.funcs);
            module.start = Some(id);
            id
        }
    };

    module
        .funcs
        .get_mut(id)
        .kind
        .unwrap_local_mut()
        .builder_mut()
}