1#![no_std]
2#[macro_use]
3extern crate alloc;
4use alloc::borrow::ToOwned;
5use alloc::boxed::Box;
6use alloc::collections::{BTreeMap, BTreeSet};
7use alloc::vec;
8use alloc::vec::Vec;
9
10use anyhow::Context;
11use bimap::BiBTreeMap;
12use fcopy::{obf_fn, DontObf, Obfuscate};
13use waffle::{
14 op_traits::op_outputs, util::new_sig, Block, Export, Func, FuncDecl, FunctionBody, Import,
15 ImportKind, Memory, Module, Operator, Signature, SignatureData, Type, Value, ValueDef,
16};
17
18pub mod bulk_memory_lowering;
19pub mod fcopy;
20pub mod tutils;
21pub mod wasi_compat;
22
23pub fn add_op(f: &mut FunctionBody, args: &[Value], rets: &[Type], op: Operator) -> Value {
24 let args = f.arg_pool.from_iter(args.iter().map(|a| *a));
25 let rets = f.type_pool.from_iter(rets.iter().map(|a| *a));
26 return f.add_value(ValueDef::Operator(op, args, rets));
27}
28pub trait Builder {
29 type Result;
30 fn build(
31 &mut self,
32 mo: &mut Module,
33 func: &mut FunctionBody,
34 k: Block,
35 ) -> anyhow::Result<(Self::Result, Block)>;
36}
37
38pub fn results_ref_2(f: &mut FunctionBody, c: Value) -> Vec<Value> {
39 let c = f.resolve_and_update_alias(c);
40 let b = f.value_blocks[c];
41 let mut v = vec![];
42 let s = match f.values[c] {
43 ValueDef::Operator(_, _1, _2) => f.type_pool[_2].to_owned(),
44 _ => return vec![c],
45 };
46 if s.len() == 1 {
47 return vec![c];
48 }
49 for (s, i) in s.iter().map(|a| *a).enumerate() {
50 let w = f.add_value(ValueDef::PickOutput(c, s as u32, i));
51 f.append_to_block(b, w);
52 v.push(w);
53 }
54
55 return v;
56}
57pub enum Expr {
58 Leaf(Value),
59 Bind(Operator, Vec<Expr>),
60 Mount(Box<dyn Builder<Result = waffle::Value>>),
61}
62impl Builder for Expr {
63 type Result = waffle::Value;
64
65 fn build(
66 &mut self,
67 mo: &mut waffle::Module,
68 func: &mut waffle::FunctionBody,
69 mut k: waffle::Block,
70 ) -> anyhow::Result<(Self::Result, waffle::Block)> {
71 match self {
72 Expr::Leaf(a) => Ok((*a, k)),
73 Expr::Bind(a, c) => {
74 let mut r = vec![];
75 for d in c.iter_mut() {
76 let (e, f) = d.build(mo, func, k)?;
77 k = f;
78 r.push(e);
79 }
80 let o = add_op(func, &r, &op_outputs(&mo, None, &a)?, a.clone());
81 func.append_to_block(k, o);
82 return Ok((o, k));
83 }
84 Expr::Mount(m) => m.build(mo, func, k),
85 }
86 }
87}
88pub fn find<'a>(m: &'a Module, a: usize, size: usize, mem: Memory) -> Option<&'a [u8]> {
89 return m.memories[mem]
90 .segments
91 .iter()
92 .filter(|s| a >= s.offset)
93 .find_map(|s| {
94 let b = a - s.offset;
95 let (_, t) = s.data.split_at_checked(b)?;
96 if t.len() < size {
97 return None;
98 }
99 return Some(&t[..size]);
100 });
101}
102pub fn find_val<'a>(
103 f: &FunctionBody,
104 a: Value,
105 size: Value,
106 m: &'a Module,
107 mem: Memory,
108) -> Option<&'a [u8]> {
109 let a = match f.values[a] {
110 ValueDef::Operator(a, _, _) => match a {
111 Operator::I32Const { value } => value as usize,
112 Operator::I64Const { value } => value as usize,
113 _ => return None,
114 },
115 _ => return None,
116 };
117 let size = match f.values[size] {
118 ValueDef::Operator(a, _, _) => match a {
119 Operator::I32Const { value } => value as usize,
120 Operator::I64Const { value } => value as usize,
121 _ => return None,
122 },
123 _ => return None,
124 };
125 find(m, a, size, mem)
126}
127pub trait Handler {
128 fn modify(
129 &mut self,
130 m: &mut Module,
131 x: &mut Operator,
132 f: &mut FunctionBody,
133 args: &mut [Value],
134 k: &mut Block,
135 ) -> anyhow::Result<()>;
136}
137pub struct Handle<H, O> {
138 pub handler: H,
139 pub obf: O,
140}
141impl<H: Handler, O: Obfuscate> Obfuscate for Handle<H, O> {
142 fn obf(
143 &mut self,
144 mut o: Operator,
145 f: &mut FunctionBody,
146 mut b: Block,
147 args: &[Value],
148 types: &[Type],
149 module: &mut Module,
150 ) -> anyhow::Result<(Value, Block)> {
151 let mut args = args.to_vec();
152 self.handler.modify(module, &mut o, f, &mut args, &mut b)?;
153 return self.obf.obf(o, f, b, &args, types, module);
154 }
155}
156pub struct Stamp {
157 pub func: Func,
158}
159impl Obfuscate for Stamp {
160 fn obf(
161 &mut self,
162 o: Operator,
163 f: &mut FunctionBody,
164 b: Block,
165 args: &[Value],
166 types: &[Type],
167 module: &mut Module,
168 ) -> anyhow::Result<(Value, Block)> {
169 if let Operator::Call { function_index } = &o {
170 if let Some((module_name, name)) = module.imports.iter().find_map(|i| {
171 if i.kind == ImportKind::Func(*function_index) {
172 Some((i.module.clone(), i.name.clone()))
173 } else {
174 None
175 }
176 }) {
177 let mem = module
178 .exports
179 .iter()
180 .find_map(|x| {
181 if x.name == "memory" {
182 match &x.kind {
183 waffle::ExportKind::Memory(m) => Some(*m),
184 _ => None,
185 }
186 } else {
187 None
188 }
189 })
190 .context("in getting the main memory")?;
191 if module_name == "stamp" && name == "mark_export" {
192 if let Some(b) = find_val(f, args[0], args[1], module, mem)
193 .and_then(|a| core::str::from_utf8(a).ok())
194 {
195 let b = b.to_owned();
196 module.exports.push(Export {
197 name: b,
198 kind: waffle::ExportKind::Func(self.func),
199 })
200 }
201 }
202 if module_name == "stamp/import" {
203 if let Some(bm) = find_val(f, args[0], args[1], module, mem)
204 .and_then(|a| core::str::from_utf8(a).ok())
205 {
206 let bm = bm.to_owned();
207 if let Some(bn) = find_val(f, args[2], args[3], module, mem)
208 .and_then(|a| core::str::from_utf8(a).ok())
209 {
210 let bn = bn.to_owned();
211 let args = &args[4..];
212 let itys = args
213 .iter()
214 .flat_map(|a| f.values[*a].tys(&f.type_pool).iter())
215 .cloned()
216 .collect::<Vec<_>>();
217 let rets = types.to_owned();
218 let s = new_sig(
219 module,
220 SignatureData::Func {
221 params: itys,
222 returns: rets,
223 },
224 );
225 let fx = module
226 .funcs
227 .push(waffle::FuncDecl::Import(s, format!("{bm}.{bn}")));
228 module.imports.push(Import {
229 module: bm,
230 name: bn,
231 kind: ImportKind::Func(fx),
232 });
233 return self.obf(
234 Operator::Call { function_index: fx },
235 f,
236 b,
237 args,
238 types,
239 module,
240 );
241 }
242 }
243 }
244 }
245 }
246 return DontObf {}.obf(o, f, b, args, types, module);
247 }
248}
249pub fn stamp(m: &mut Module) -> anyhow::Result<()> {
250 for f in m.funcs.iter().collect::<Vec<_>>() {
251 obf_fn(m, f, &mut Stamp { func: f })?;
252 }
253 Ok(())
254}
255pub fn collect_jmpfns2(
256 m: &mut Module,
257 mut filter: impl FnMut(Func, Block) -> bool,
258) -> BiBTreeMap<(Func, Block), Func> {
259 let a = collect_jmpfns3(m, move|_,a,b|filter(a,b));
260 let mut res = BiBTreeMap::new();
261 for (a,b) in a .into_iter(){
262 res.insert(a, b);
263 };
264 return res;
265}
266pub fn collect_jmpfns3(
267 m: &mut Module,
268 mut filter: impl FnMut(&mut Module, Func, Block) -> bool,
269) -> BTreeMap<(Func,Block),Func> {
270 let mut fns = m
271 .funcs
272 .entries()
273 .filter(|(_, b)| matches!(b, FuncDecl::Body(_, _, _)))
274 .map(|a| a.0)
275 .collect::<BTreeSet<_>>();
276 let mut res = BTreeMap::new();
277 for f in fns {
278 let n = m.funcs[f].name().to_owned();
279 let body = m.funcs[f].body().unwrap().clone();
280 let mut map = BTreeMap::new();
281 for k in body.blocks.iter() {
282 if k == body.entry {
283 map.insert(k, f);
284 continue;
285 }
286 if !filter(m, f, k) {
287 continue;
288 }
289 let mut b2 = body.clone();
290 b2.entry = k;
291 let s = new_sig(
292 m,
293 SignatureData::Func {
294 params: b2.blocks[k].params.iter().map(|a| a.0).collect(),
295 returns: b2.rets.clone(),
296 },
297 );
298 let g = m.funcs.push(FuncDecl::Body(s, format!("{n}~{k}"), b2));
299 map.insert(k, g);
300 }
301 for k in body.blocks.iter() {
302 for (a, m) in map.iter() {
303 if let Some(v) = map.get(&k).cloned() {
304 res.insert((*m, k), v);
305 }
306 }
307 }
308 }
309 return res;
310}
311pub fn collect_jmpfns(m: &mut Module) -> BiBTreeMap<(Func, Block), Func> {
312 return collect_jmpfns2(m, |_, _| true);
313}
314