portal_pc_waffle_robake/
lib.rs

1use std::{
2    cell::OnceCell,
3    collections::{BTreeMap, BTreeSet},
4};
5
6use itertools::Itertools;
7use once_map::{OnceMap, RandomState};
8use waffle::{
9    copying::fcopy::{obf_mod, DontObf, Obfuscate},
10    entity::PerEntity,
11    Block, BlockTarget, Func, Memory, MemoryArg, Module, Operator, Type,
12};
13use waffle_ast::bulk_memory_lowering::LowerBulkMemory;
14pub struct Robake {
15    pub rom: PerEntity<Memory, BTreeMap<u64, u8>>,
16    pub rf: Option<Func>,
17}
18impl Robake {
19    pub fn from_vals(m: &Module, i: PerEntity<Memory, BTreeSet<u64>>) -> Self {
20        let mut x: PerEntity<Memory, BTreeMap<u64, u8>> = PerEntity::default();
21        for (m, d) in m.memories.entries() {
22            x[m] = i[m]
23                .iter()
24                .cloned()
25                .filter_map(|a| {
26                    Some((
27                        a,
28                        d.segments
29                            .iter()
30                            .find_map(|s| {
31                                let a = a.checked_sub(s.offset as u64)?;
32                                s.data.get(a as usize).cloned()
33                            })
34                            .unwrap_or(0),
35                    ))
36                })
37                .collect();
38        }
39        return Self { rom: x, rf: None };
40    }
41    pub fn from_map(m: &Module, x: impl Iterator<Item = u8>, mem: Memory) -> Self {
42        let mut i: PerEntity<Memory, BTreeSet<u64>> = PerEntity::default();
43        i[mem].extend(
44            x.batching(|mut i| i.next_array())
45                .map(|a| u64::from_le_bytes(a))
46                .tuples()
47                .flat_map(|(a, b)| a..b),
48        );
49        return Self::from_vals(m, i);
50    }
51    pub fn bake(&mut self, m: &mut Module) -> anyhow::Result<()> {
52        obf_mod(m, &mut LowerBulkMemory {})?;
53        // obf_mod(
54        //     m,
55        //     &mut waffle_ast::bulk_memory_lowering::Reload {
56        //         wrapped: DontObf {},
57        //     },
58        // )?;
59        return obf_mod(m, self);
60    }
61    fn redo(
62        &mut self,
63        o: waffle::Operator,
64        f: &mut waffle::FunctionBody,
65        b: waffle::Block,
66        args: &[waffle::Value],
67        types: &[waffle::Type],
68        module: &mut waffle::Module,
69        memory: MemoryArg,
70        mut unpack: impl FnMut(&mut (dyn Iterator<Item = Option<u8>> + '_)) -> Option<Option<u64>>,
71    ) -> anyhow::Result<(waffle::Value, waffle::Block)> {
72        let n = f.add_block();
73        let fb = f.add_block();
74        {
75            let (a, fb) = DontObf {}.obf(o, f, fb, args, types, module)?;
76            f.set_terminator(
77                fb,
78                waffle::Terminator::Br {
79                    target: BlockTarget {
80                        block: n,
81                        args: vec![a],
82                    },
83                },
84            );
85        };
86        let bs: OnceMap<u64, Box<Block>, RandomState> = Default::default();
87        // let mut mr = self.rom[memory.memory]
88        //     .iter()
89        //     .filter_map(|(a, b)| Some((a.checked_sub(memory.offset)?, *b)))
90        //     .collect::<BTreeMap<_, _>>();
91        let base = self.rom[memory.memory].keys().cloned().min().unwrap_or(0);
92        let cases = (base..(match self.rom[memory.memory].keys().cloned().max() {
93            None => 0,
94            Some(a) => a + 1,
95        }))
96            .filter_map(|a| a.checked_sub(memory.offset))
97            .map(|a| self.rom[memory.memory].get(&(a + memory.offset)).cloned())
98            .collect_vec()
99            .into_iter()
100            .batching(move |a| unpack(a))
101            .map(|a| match a {
102                None => BlockTarget {
103                    block: fb,
104                    args: vec![],
105                },
106                Some(a) => BlockTarget {
107                    block: *bs.insert(a as u64, |_| {
108                        // let i = a as usize;
109                        let b = f.add_block();
110                        let i = f.add_op(b, Operator::I64Const { value: a }, &[], &[Type::I64]);
111                        let i = if types[0] == Type::I32 {
112                            f.add_op(b, Operator::I32WrapI64, &[i], &[Type::I32])
113                        } else {
114                            i
115                        };
116                        let i = match self.rf {
117                            None => i,
118                            Some(r) => f.add_op(
119                                b,
120                                Operator::Call { function_index: r },
121                                &[i],
122                                &[Type::I32],
123                            ),
124                        };
125                        f.set_terminator(
126                            b,
127                            waffle::Terminator::Br {
128                                target: BlockTarget {
129                                    block: n,
130                                    args: vec![i],
131                                },
132                            },
133                        );
134                        Box::new(b)
135                    }),
136                    args: vec![],
137                },
138            })
139            .collect::<Vec<_>>();
140        let mut r = args[0];
141        let k = f.add_op(
142            b,
143            if module.memories[memory.memory].memory64 {
144                Operator::I64Const { value: base }
145            } else {
146                Operator::I32Const {
147                    value: (base & 0xffffffff) as u32,
148                }
149            },
150            &[],
151            if module.memories[memory.memory].memory64 {
152                &[Type::I64]
153            } else {
154                &[Type::I32]
155            },
156        );
157        r = f.add_op(
158            b,
159            if module.memories[memory.memory].memory64 {
160                Operator::I64Sub
161            } else {
162                Operator::I32Sub
163            },
164            &[r, k],
165            if module.memories[memory.memory].memory64 {
166                &[Type::I64]
167            } else {
168                &[Type::I32]
169            },
170        );
171        f.set_terminator(
172            b,
173            waffle::Terminator::Select {
174                value: r,
175                targets: cases,
176                default: BlockTarget {
177                    block: fb,
178                    args: vec![],
179                },
180            },
181        );
182        return Ok((f.add_blockparam(n, types[0]), n));
183    }
184}
185macro_rules! ttr {
186    ($a:expr) => {
187        match $a {
188            Some(Some(b)) => b,
189            Some(None) => return Some(None),
190            None => return None,
191        }
192    };
193}
194impl Obfuscate for Robake {
195    fn obf(
196        &mut self,
197        o: waffle::Operator,
198        f: &mut waffle::FunctionBody,
199        b: waffle::Block,
200        args: &[waffle::Value],
201        types: &[waffle::Type],
202        module: &mut waffle::Module,
203    ) -> anyhow::Result<(waffle::Value, waffle::Block)> {
204        if let Operator::I32Load8U { memory } = o.clone() {
205            return self.redo(o, f, b, args, types, module, memory, |a| {
206                a.next().map(|a| a.map(|b| b as u64))
207            });
208        }
209        if let Operator::I32Load8S { memory } = o.clone() {
210            return self.redo(o, f, b, args, types, module, memory, |a| {
211                a.next().map(|a| a.map(|b| b as i8 as i64 as u64))
212            });
213        }
214        if let Operator::I32Load16U { memory } = o.clone() {
215            return self.redo(o, f, b, args, types, module, memory, |a| {
216                Some(Some(
217                    u16::from_le_bytes([ttr!(a.next()), ttr!(a.next())]) as u64
218                ))
219            });
220        }
221        if let Operator::I32Load16S { memory } = o.clone() {
222            return self.redo(o, f, b, args, types, module, memory, |a| {
223                Some(Some(
224                    u16::from_le_bytes([ttr!(a.next()), ttr!(a.next())]) as i16 as i64 as u64,
225                ))
226            });
227        }
228        if let Operator::I64Load8U { memory } = o.clone() {
229            return self.redo(o, f, b, args, types, module, memory, |a| {
230                a.next().map(|a| a.map(|b| b as u64))
231            });
232        }
233        if let Operator::I64Load8S { memory } = o.clone() {
234            return self.redo(o, f, b, args, types, module, memory, |a| {
235                a.next().map(|a| a.map(|b| b as i8 as i64 as u64))
236            });
237        }
238        if let Operator::I64Load16U { memory } = o.clone() {
239            return self.redo(o, f, b, args, types, module, memory, |a| {
240                Some(Some(
241                    u16::from_le_bytes([ttr!(a.next()), ttr!(a.next())]) as u64
242                ))
243            });
244        }
245        if let Operator::I64Load16S { memory } = o.clone() {
246            return self.redo(o, f, b, args, types, module, memory, |a| {
247                Some(Some(
248                    u16::from_le_bytes([ttr!(a.next()), ttr!(a.next())]) as i16 as i64 as u64,
249                ))
250            });
251        }
252        if let Operator::I64Load32U { memory } = o.clone() {
253            return self.redo(o, f, b, args, types, module, memory, |a| {
254                Some(Some(u32::from_le_bytes([
255                    ttr!(a.next()),
256                    ttr!(a.next()),
257                    ttr!(a.next()),
258                    ttr!(a.next()),
259                ]) as u64))
260            });
261        }
262        if let Operator::I64Load32S { memory } = o.clone() {
263            return self.redo(o, f, b, args, types, module, memory, |a| {
264                Some(Some(u32::from_le_bytes([
265                    ttr!(a.next()),
266                    ttr!(a.next()),
267                    ttr!(a.next()),
268                    ttr!(a.next()),
269                ]) as i32 as i64 as u64))
270            });
271        }
272        if let Operator::I32Load { memory } = o.clone() {
273            return self.redo(o, f, b, args, types, module, memory, |a| {
274                Some(Some(u32::from_le_bytes([
275                    ttr!(a.next()),
276                    ttr!(a.next()),
277                    ttr!(a.next()),
278                    ttr!(a.next()),
279                ]) as u64))
280            });
281        }
282        if let Operator::I64Load { memory } = o.clone() {
283            return self.redo(o, f, b, args, types, module, memory, |a| {
284                Some(Some(u64::from_le_bytes([
285                    ttr!(a.next()),
286                    ttr!(a.next()),
287                    ttr!(a.next()),
288                    ttr!(a.next()),
289                    ttr!(a.next()),
290                    ttr!(a.next()),
291                    ttr!(a.next()),
292                    ttr!(a.next()),
293                ])))
294            });
295        }
296        return DontObf {}.obf(o, f, b, args, types, module);
297    }
298}
portal_pc_waffle_robake/lib.rs

portal_pc_waffle_robake/
lib.rs
