1use alloc::borrow::ToOwned;
2use alloc::collections::{BTreeMap, BTreeSet};
3use alloc::format;
4use alloc::string::{String, ToString};
5use alloc::vec;
6use alloc::vec::Vec;
7use core::iter::once;
8use core::mem::{replace, take};
9use portal_pc_waffle::copying::fcopy::{obf_mod, DontObf, Obfuscate};
10
11use anyhow::Context;
12use portal_pc_waffle::{
13 entity::EntityRef, util::new_sig, ExportKind, Func, FuncDecl, FunctionBody, HeapType, Import,
14 ImportKind, Module, Operator, SignatureData, StorageType, TableData, Type, ValueDef,
15 WithNullable,
16};
17use crate::canon::canon;
24use crate::util::add_op;
25pub fn patch_ty(t: &mut Type) {
26 if let Type::Heap(WithNullable {
27 value: HeapType::ExternRef,
28 nullable,
29 }) = t.clone()
30 {
31 *t = Type::I32
32 }
33}
34pub struct LowerTables {}
35impl Obfuscate for LowerTables {
36 fn obf(
37 &mut self,
38 o: Operator,
39 f: &mut portal_pc_waffle::FunctionBody,
40 b: portal_pc_waffle::Block,
41 args: &[portal_pc_waffle::Value],
42 types: &[Type],
43 module: &mut Module,
44 ) -> anyhow::Result<(portal_pc_waffle::Value, portal_pc_waffle::Block)> {
45 match o {
46 Operator::TableGet { table_index }
47 | Operator::TableSet { table_index }
48 | Operator::TableSize { table_index }
49 | Operator::TableGrow { table_index } => {
50 match module.tables[table_index].ty.clone() {
51 Type::I32 | Type::I64 | Type::F32 | Type::F64 => {
52 let w = add_op(
54 f,
55 &[],
56 &[Type::I32],
57 Operator::I32Const {
58 value: table_index.index() as u32,
59 },
60 );
61 f.append_to_block(b, w);
62 let id = format!(
63 "pit-patch-rt/@{}/{}",
64 module.tables[table_index].ty.to_string(),
65 o.to_string().split_once("<").unwrap().0
66 );
67 let mut a = module.exports.iter();
68 let a = loop {
69 let Some(b) = a.next() else {
70 anyhow::bail!("pit patch rt not found")
71 };
72 if b.name != id {
73 continue;
74 }
75 let ExportKind::Func(a) = &b.kind else {
76 continue;
77 };
78 break *a;
79 };
80 DontObf {}.obf(
81 Operator::Call { function_index: a },
82 f,
83 b,
84 &once(w).chain(args.iter().cloned()).collect::<Vec<_>>(),
85 types,
86 module,
87 )
88 }
89 _ => DontObf {}.obf(o, f, b, args, types, module),
90 }
91 }
92 _ => DontObf {}.obf(o, f, b, args, types, module),
93 }
94 }
95}
96pub fn import_fn(m: &mut Module, mo: String, n: String, s: SignatureData) -> Func {
97 for i in m.imports.iter() {
98 if i.module == mo && i.name == n {
99 if let ImportKind::Func(f) = &i.kind {
100 return *f;
101 }
102 }
103 }
104 let s = new_sig(m, s);
105 let f = m
106 .funcs
107 .push(portal_pc_waffle::FuncDecl::Import(s, format!("{mo}.{n}")));
108 m.imports.push(Import {
109 module: mo,
110 name: n,
111 kind: ImportKind::Func(f),
112 });
113 return f;
114}
115pub struct Cfg {
116 pub unexportable_i32_tables: bool,
117}
118pub fn canon_all(m: &mut Module, cfg: &Cfg, root: &str) -> anyhow::Result<()> {
119 let i = crate::get_interfaces(m)?;
120 let interfaces = i
121 .into_iter()
122 .collect::<BTreeSet<_>>()
123 .into_iter()
124 .collect::<Vec<_>>();
125 for j in interfaces.iter() {
126 canon(m, &j.rid_str(), root)?;
127 }
128 if !cfg.unexportable_i32_tables {
129 obf_mod(m, &mut LowerTables {})?;
130 }
131 Ok(())
132}
133pub fn instantiate(m: &mut Module, cfg: &Cfg) -> anyhow::Result<()> {
134 let root = "pit_patch_internal_instantiate";
135 let i = crate::get_interfaces(m)?;
136 let interfaces = i
137 .into_iter()
138 .collect::<BTreeSet<_>>()
139 .into_iter()
140 .collect::<Vec<_>>();
141 for j in interfaces.iter() {
142 canon(m, &j.rid_str(), root)?;
143 }
144 for s in m.signatures.values_mut() {
145 match s {
146 SignatureData::Func {
147 params, returns, ..
148 } => {
149 for p in params.iter_mut().chain(returns.iter_mut()) {
150 patch_ty(p)
151 }
152 }
153 SignatureData::Struct { fields, .. } => {
154 for f in fields.iter_mut() {
155 if let StorageType::Val(v) = &mut f.value {
156 patch_ty(v);
157 }
158 }
159 }
160 SignatureData::Array { ty, .. } => {
161 if let StorageType::Val(v) = &mut ty.value {
162 patch_ty(v);
163 }
164 }
165 SignatureData::None => {}
166 _ => todo!(),
167 }
168 }
169 m.try_take_per_func_body(|m, b| {
170 let x = b.type_pool.from_iter(once(Type::I32));
171 for p in b
172 .blocks
173 .values_mut()
174 .flat_map(|a| a.params.iter_mut().map(|a| &mut a.0))
175 .chain(b.rets.iter_mut())
176 {
177 patch_ty(p)
178 }
179 for v in b.values.iter().collect::<Vec<_>>() {
180 if let ValueDef::Operator(o, _1, tys) = &mut b.values[v] {
181 match o.clone() {
182 Operator::RefNull { ty } => {
183 if matches!(
184 ty,
185 Type::Heap(WithNullable {
186 value: HeapType::ExternRef,
187 nullable
188 })
189 ) {
190 *o = Operator::I32Const { value: 0 }
191 }
192 }
193 Operator::RefIsNull => {
194 if matches!(
195 b.type_pool[*tys][0],
196 Type::Heap(WithNullable {
197 value: HeapType::ExternRef,
198 nullable
199 })
200 ) {
201 *o = Operator::I32Eqz
202 }
203 }
204 _ => {}
205 }
206 }
207 }
208 for t in b.type_pool.storage.iter_mut() {
209 patch_ty(t)
210 }
211 Ok::<_, anyhow::Error>(())
212 })?;
213 for t in m.tables.values_mut() {
215 if matches!(
216 t.ty,
217 Type::Heap(WithNullable {
218 value: HeapType::ExternRef,
219 nullable
220 })
221 ) {
222 t.ty = Type::I32;
223 }
224 }
225 let rl = interfaces.len();
234 for i in take(&mut m.imports) {
235 if let Some(rid) = i.module.strip_prefix("pit/") {
236 let ridx = interfaces
237 .iter()
238 .enumerate()
239 .find_map(|x| {
240 if x.1.rid_str() == rid {
241 Some(x.0)
242 } else {
243 None
244 }
245 })
246 .context("in getting the index")?;
247
248 if i.name.ends_with(&format!("~{root}")) {
249 if let ImportKind::Func(f) = i.kind {
250 let fs = m.funcs[f].sig();
251 let fname = m.funcs[f].name().to_owned();
252 let mut b = FunctionBody::new(&m, fs);
253 let k = b.entry;
254 let mut values: Vec<portal_pc_waffle::Value> =
255 b.blocks[k].params.iter().map(|a| a.1).collect();
256 let tys = b.blocks[k].params.iter().map(|a| a.0).collect::<Vec<_>>();
257
258 let v = {
273 let ridx = b.add_op(
274 k,
275 Operator::I32Const { value: ridx as u32 },
276 &[],
277 &[Type::I32],
278 );
279 let rl = b.add_op(
280 k,
281 Operator::I32Const { value: rl as u32 },
282 &[],
283 &[Type::I32],
284 );
285 let v = b.add_op(k, Operator::I32Mul, &[rl, values[0]], &[Type::I32]);
286 b.add_op(k, Operator::I32Add, &[v, ridx], &[Type::I32])
287 };
288 values = vec![v];
289 b.set_terminator(k, portal_pc_waffle::Terminator::Return { values });
290 m.funcs[f] = FuncDecl::Body(fs, fname, b);
291 continue;
292 }
293 } else {
294 let ek = format!("{}/~{root}{}", i.module, i.name);
295 let mut ex = m.exports.iter();
296 let ex = loop {
297 let Some(x) = ex.next() else {
298 anyhow::bail!("export not found")
299 };
300 if x.name != ek {
301 continue;
302 };
303 let ExportKind::Func(ef) = &x.kind else {
304 continue;
305 };
306 break *ef;
307 };
308 if let ImportKind::Func(f) = i.kind {
309 let fs = m.funcs[f].sig();
310 let fname = m.funcs[f].name().to_owned();
311 let mut b = FunctionBody::new(&m, fs);
312 let k = b.entry;
313 let mut args = b.blocks[b.entry]
314 .params
315 .iter()
316 .map(|a| a.1)
317 .collect::<Vec<_>>();
318 let v = match args[0] {
319 v => {
320 let ridx = b.add_op(
321 k,
322 Operator::I32Const { value: ridx as u32 },
323 &[],
324 &[Type::I32],
325 );
326 let rl = b.add_op(
327 k,
328 Operator::I32Const { value: rl as u32 },
329 &[],
330 &[Type::I32],
331 );
332 b.add_op(k, Operator::I32DivU, &[v, rl], &[Type::I32])
334 }
335 };
336 args[0] = v;
337 b.set_terminator(
338 k,
339 portal_pc_waffle::Terminator::ReturnCall { func: ex, args },
340 );
341 m.funcs[f] = FuncDecl::Body(fs, fname, b);
342 continue;
343 }
344 }
345 }
346 if i.module == "pit" {
347 let n = &i.name;
348 let fs = interfaces
349 .iter()
350 .filter_map(|i| {
351 let mut ex = m.exports.iter();
352 let ex = loop {
353 let Some(x) = ex.next() else {
354 return None;
355 };
356 if x.name != format!("pit/{}/{root}.{n}", i.rid_str()) {
357 continue;
358 };
359 let ExportKind::Func(ef) = &x.kind else {
360 continue;
361 };
362 break *ef;
363 };
364 return Some(ex);
365 })
366 .collect::<Vec<_>>();
367 let t = m.tables.push(TableData {
368 ty: Type::Heap(WithNullable {
369 nullable: true,
370 value: portal_pc_waffle::HeapType::FuncRef,
371 }),
372 initial: fs.len() as u64,
373 max: Some(fs.len() as u64),
374 func_elements: Some(fs.clone()),
375 table64: false,
376 });
377 if let ImportKind::Func(f) = i.kind {
378 let fsi = m.funcs[f].sig();
379 let fname = m.funcs[f].name().to_owned();
380 let mut b = FunctionBody::new(&m, fsi);
381 let k = b.entry;
382 let mut args = b.blocks[b.entry]
383 .params
384 .iter()
385 .map(|a| a.1)
386 .collect::<Vec<_>>();
387 let (a, c) = match b.blocks[k].params[0].1 {
414 v => {
415 let rl = b.add_op(
417 k,
418 Operator::I32Const { value: rl as u32 },
419 &[],
420 &[Type::I32],
421 );
422 (
424 b.add_op(k, Operator::I32DivU, &[v, rl], &[Type::I32]),
425 b.add_op(k, Operator::I32RemU, &[v, rl], &[Type::I32]),
426 )
427 }
428 };
429 args[0] = a;
430 args.push(c);
431 b.set_terminator(
432 k,
433 portal_pc_waffle::Terminator::ReturnCallIndirect {
434 sig: fsi,
435 table: t,
436 args,
437 },
438 );
439 m.funcs[f] = FuncDecl::Body(fsi, fname, b);
440 continue;
441 }
442 }
443 if i.module == "system" && i.name == "stub" {
444 if let ImportKind::Func(f) = i.kind {
445 let fsi = m.funcs[f].sig();
446 let fname = m.funcs[f].name().to_owned();
447 let mut b = FunctionBody::new(&m, fsi);
448 let k = b.entry;
449 let s = b.add_op(k, Operator::I32Const { value: 1 }, &[], &[Type::I32]);
450 b.set_terminator(k, portal_pc_waffle::Terminator::Return { values: vec![s] });
451 m.funcs[f] = FuncDecl::Body(fsi, fname, b);
452 continue;
453 }
454 }
455 m.imports.push(i);
456 }
457 if !cfg.unexportable_i32_tables {
458 obf_mod(m, &mut LowerTables {})?;
459 }
460 Ok(())
461}