1use crate::{BackendError, MemoryBackend, Result};
2use decant_protocol::{MemRegion, ModuleInfo, Pid, ProcessInfo};
3use std::collections::BTreeMap;
4use std::sync::{Arc, RwLock};
5
6const PAGE: u64 = 0x1000;
7
8#[derive(Debug, Clone)]
9struct Region {
10 base: u64,
11 size: u64,
12 readable: bool,
13 writable: bool,
14 executable: bool,
15}
16
17impl Region {
18 fn contains(&self, addr: u64) -> bool {
19 addr >= self.base && addr < self.base + self.size
20 }
21 fn to_mem_region(&self) -> MemRegion {
22 MemRegion {
23 base: self.base,
24 size: self.size,
25 readable: self.readable,
26 writable: self.writable,
27 executable: self.executable,
28 }
29 }
30}
31
32#[derive(Debug)]
33struct Process {
34 info: ProcessInfo,
35 modules: Vec<ModuleInfo>,
36 exports: BTreeMap<String, Vec<(String, u64)>>,
37 regions: Vec<Region>,
38 mem: BTreeMap<u64, u8>,
39}
40
41impl Process {
42 fn region_at(&self, addr: u64) -> Option<&Region> {
43 self.regions.iter().find(|r| r.contains(addr))
44 }
45}
46
47#[derive(Debug, Clone)]
48pub struct MockGuest {
49 inner: Arc<RwLock<Vec<Process>>>,
50}
51
52impl MockGuest {
53 pub fn builder() -> GuestBuilder {
54 GuestBuilder {
55 processes: Vec::new(),
56 }
57 }
58}
59
60pub struct GuestBuilder {
61 processes: Vec<Process>,
62}
63
64impl GuestBuilder {
65 pub fn process(self, name: &str, pid: Pid) -> ProcessBuilder {
66 ProcessBuilder {
67 parent: self,
68 proc: Process {
69 info: ProcessInfo {
70 pid,
71 name: name.to_string(),
72 },
73 modules: Vec::new(),
74 exports: BTreeMap::new(),
75 regions: Vec::new(),
76 mem: BTreeMap::new(),
77 },
78 cur_region: None,
79 }
80 }
81
82 pub fn build(self) -> MockGuest {
83 MockGuest {
84 inner: Arc::new(RwLock::new(self.processes)),
85 }
86 }
87}
88
89struct PendingRegion {
90 base: u64,
91 readable: bool,
92 writable: bool,
93 executable: bool,
94 max_written: Option<u64>,
95}
96
97pub struct ProcessBuilder {
98 parent: GuestBuilder,
99 proc: Process,
100 cur_region: Option<PendingRegion>,
101}
102
103impl ProcessBuilder {
104 pub fn module(mut self, name: &str, base: u64, size: u64) -> Self {
105 self.proc.modules.push(ModuleInfo {
106 name: name.to_string(),
107 base,
108 size,
109 });
110 self
111 }
112
113 pub fn export(mut self, module: &str, name: &str, addr: u64) -> Self {
114 self.proc
115 .exports
116 .entry(module.to_ascii_lowercase())
117 .or_default()
118 .push((name.to_string(), addr));
119 self
120 }
121
122 pub fn region(mut self, base: u64, perms: &str) -> Self {
123 self.finalize_region();
124 let p = perms.as_bytes();
125 self.cur_region = Some(PendingRegion {
126 base,
127 readable: p.first() == Some(&b'r'),
128 writable: p.get(1) == Some(&b'w'),
129 executable: p.get(2) == Some(&b'x'),
130 max_written: None,
131 });
132 self
133 }
134
135 pub fn bytes_at(mut self, addr: u64, bytes: &[u8]) -> Self {
136 {
137 let region = self
138 .cur_region
139 .as_mut()
140 .expect("bytes_at called before .region(); open a region first");
141 assert!(
142 addr >= region.base,
143 "bytes_at {addr:#x} is below the current region base {:#x}",
144 region.base
145 );
146 let end = addr + bytes.len().saturating_sub(1) as u64;
147 region.max_written = Some(region.max_written.map_or(end, |m| m.max(end)));
148 }
149 for (i, b) in bytes.iter().enumerate() {
150 self.proc.mem.insert(addr + i as u64, *b);
151 }
152 self
153 }
154
155 pub fn u64_at(self, addr: u64, value: u64) -> Self {
156 self.bytes_at(addr, &value.to_le_bytes())
157 }
158
159 pub fn u32_at(self, addr: u64, value: u32) -> Self {
160 self.bytes_at(addr, &value.to_le_bytes())
161 }
162
163 fn finalize_region(&mut self) {
164 if let Some(r) = self.cur_region.take() {
165 let span = match r.max_written {
166 Some(m) => m - r.base + 1,
167 None => 0,
168 };
169 let size = span.div_ceil(PAGE).max(1) * PAGE;
170 self.proc.regions.push(Region {
171 base: r.base,
172 size,
173 readable: r.readable,
174 writable: r.writable,
175 executable: r.executable,
176 });
177 }
178 }
179
180 pub fn done(mut self) -> GuestBuilder {
181 self.finalize_region();
182 for v in self.proc.exports.values_mut() {
183 v.sort();
184 }
185 let mut parent = self.parent;
186 parent.processes.push(self.proc);
187 parent
188 }
189}
190
191#[derive(Debug, Clone)]
192pub struct MockBackend {
193 guest: MockGuest,
194}
195
196impl MockBackend {
197 pub fn new(guest: MockGuest) -> Self {
198 MockBackend { guest }
199 }
200
201 pub fn guest(&self) -> MockGuest {
202 self.guest.clone()
203 }
204}
205
206fn idx_by_pid(procs: &[Process], pid: Pid) -> Result<usize> {
207 procs
208 .iter()
209 .position(|p| p.info.pid == pid)
210 .ok_or(BackendError::NoSuchProcess {
211 pid: Some(pid.0),
212 name: None,
213 })
214}
215
216impl MemoryBackend for MockBackend {
217 fn list_processes(&self) -> Result<Vec<ProcessInfo>> {
218 let g = self.guest.inner.read().unwrap();
219 Ok(g.iter().map(|p| p.info.clone()).collect())
220 }
221
222 fn process_by_pid(&self, pid: Pid) -> Result<ProcessInfo> {
223 let g = self.guest.inner.read().unwrap();
224 Ok(g[idx_by_pid(&g, pid)?].info.clone())
225 }
226
227 fn process_by_name(&self, name: &str) -> Result<ProcessInfo> {
228 let g = self.guest.inner.read().unwrap();
229 g.iter()
230 .find(|p| p.info.name.eq_ignore_ascii_case(name))
231 .map(|p| p.info.clone())
232 .ok_or_else(|| BackendError::NoSuchProcess {
233 pid: None,
234 name: Some(name.to_string()),
235 })
236 }
237
238 fn module_list(&self, pid: Pid) -> Result<Vec<ModuleInfo>> {
239 let g = self.guest.inner.read().unwrap();
240 Ok(g[idx_by_pid(&g, pid)?].modules.clone())
241 }
242
243 fn module_by_name(&self, pid: Pid, name: &str) -> Result<ModuleInfo> {
244 let g = self.guest.inner.read().unwrap();
245 let p = &g[idx_by_pid(&g, pid)?];
246 p.modules
247 .iter()
248 .find(|m| m.name.eq_ignore_ascii_case(name))
249 .cloned()
250 .ok_or_else(|| BackendError::NoSuchModule {
251 pid: pid.0,
252 module: name.to_string(),
253 })
254 }
255
256 fn module_exports(&self, pid: Pid, module: &str) -> Result<Vec<(String, u64)>> {
257 let g = self.guest.inner.read().unwrap();
258 let p = &g[idx_by_pid(&g, pid)?];
259 if !p
260 .modules
261 .iter()
262 .any(|m| m.name.eq_ignore_ascii_case(module))
263 {
264 return Err(BackendError::NoSuchModule {
265 pid: pid.0,
266 module: module.to_string(),
267 });
268 }
269 Ok(p.exports
270 .get(&module.to_ascii_lowercase())
271 .cloned()
272 .unwrap_or_default())
273 }
274
275 fn read(&self, pid: Pid, addr: u64, len: usize) -> Result<Vec<u8>> {
276 let g = self.guest.inner.read().unwrap();
277 let p = &g[idx_by_pid(&g, pid)?];
278 let mut out = Vec::with_capacity(len);
279 for i in 0..len as u64 {
280 let a = addr + i;
281 match p.region_at(a) {
282 Some(r) if r.readable => out.push(p.mem.get(&a).copied().unwrap_or(0)),
283 _ => {
284 return Err(BackendError::ReadFailed {
285 addr,
286 len: len as u64,
287 reason: format!("address {a:#x} is not in a readable region"),
288 });
289 }
290 }
291 }
292 Ok(out)
293 }
294
295 fn write(&self, pid: Pid, addr: u64, data: &[u8]) -> Result<usize> {
296 let mut g = self.guest.inner.write().unwrap();
297 let i = idx_by_pid(&g, pid)?;
298 for off in 0..data.len() as u64 {
299 let a = addr + off;
300 match g[i].region_at(a) {
301 Some(r) if r.writable => {}
302 _ => {
303 return Err(BackendError::WriteFailed {
304 addr,
305 reason: format!("address {a:#x} is not in a writable region"),
306 });
307 }
308 }
309 }
310 for (off, b) in data.iter().enumerate() {
311 g[i].mem.insert(addr + off as u64, *b);
312 }
313 Ok(data.len())
314 }
315
316 fn memory_map(&self, pid: Pid) -> Result<Vec<MemRegion>> {
317 let g = self.guest.inner.read().unwrap();
318 Ok(g[idx_by_pid(&g, pid)?]
319 .regions
320 .iter()
321 .map(Region::to_mem_region)
322 .collect())
323 }
324}
325
326#[cfg(test)]
327mod tests {
328 use super::*;
329
330 const MAGIC: &[u8] = b"\xDE\xCA\x47\x00MAGIC";
331
332 fn guest() -> MockGuest {
333 MockGuest::builder()
334 .process("target.exe", Pid(1234))
335 .module("target.exe", 0x1400000000, 0x80000)
336 .export("target.exe", "tick", 0x1400001000)
337 .region(0x1400010000, "rw-")
338 .bytes_at(0x1400010100, MAGIC)
339 .u64_at(0x1400010200, 0x1400010300)
340 .u32_at(0x1400010300 + 0x10, 1337)
341 .done()
342 .process("explorer.exe", Pid(4))
343 .done()
344 .build()
345 }
346
347 #[test]
348 fn enumerates_processes_and_modules() {
349 let b = MockBackend::new(guest());
350 let procs = b.list_processes().unwrap();
351 assert_eq!(procs.len(), 2);
352 assert_eq!(b.process_by_name("TARGET.EXE").unwrap().pid, Pid(1234));
353 assert_eq!(b.process_by_pid(Pid(4)).unwrap().name, "explorer.exe");
354 let mods = b.module_list(Pid(1234)).unwrap();
355 assert_eq!(mods[0].base, 0x1400000000);
356 }
357
358 #[test]
359 fn planted_bytes_read_back() {
360 let b = MockBackend::new(guest());
361 let got = b.read(Pid(1234), 0x1400010100, MAGIC.len()).unwrap();
362 assert_eq!(got, MAGIC);
363 let hop = b.read(Pid(1234), 0x1400010200, 8).unwrap();
364 assert_eq!(u64::from_le_bytes(hop.try_into().unwrap()), 0x1400010300);
365 let term = b.read(Pid(1234), 0x1400010310, 4).unwrap();
366 assert_eq!(u32::from_le_bytes(term.try_into().unwrap()), 1337);
367 }
368
369 #[test]
370 fn writes_round_trip() {
371 let b = MockBackend::new(guest());
372 let n = b
373 .write(Pid(1234), 0x1400010400, &[0xAA, 0xBB, 0xCC, 0xDD])
374 .unwrap();
375 assert_eq!(n, 4);
376 let got = b.read(Pid(1234), 0x1400010400, 4).unwrap();
377 assert_eq!(got, vec![0xAA, 0xBB, 0xCC, 0xDD]);
378 }
379
380 #[test]
381 fn write_to_readonly_region_fails() {
382 let g = MockGuest::builder()
383 .process("ro.exe", Pid(1))
384 .region(0x2000, "r-x")
385 .bytes_at(0x2000, &[0; 8])
386 .done()
387 .build();
388 let b = MockBackend::new(g);
389 assert!(b.write(Pid(1), 0x2000, &[1]).is_err());
390 }
391
392 #[test]
393 fn unwritten_bytes_in_region_read_zero() {
394 let b = MockBackend::new(guest());
395 let got = b.read(Pid(1234), 0x1400010000, 16).unwrap();
396 assert_eq!(got, vec![0u8; 16]);
397 }
398
399 #[test]
400 fn read_outside_region_fails() {
401 let b = MockBackend::new(guest());
402 assert!(b.read(Pid(1234), 0xdead_0000, 4).is_err());
403 }
404
405 #[test]
406 fn memory_map_covers_planted_addresses() {
407 let b = MockBackend::new(guest());
408 let map = b.memory_map(Pid(1234)).unwrap();
409 assert_eq!(map.len(), 1);
410 let r = map[0];
411 assert_eq!(r.base, 0x1400010000);
412 assert!(r.readable && r.writable && !r.executable);
413 assert!(r.base + r.size > 0x1400010314);
414 }
415
416 #[test]
417 fn unknown_pid_and_module_error() {
418 let b = MockBackend::new(guest());
419 assert!(b.process_by_pid(Pid(9999)).is_err());
420 assert!(b.module_by_name(Pid(1234), "nope.dll").is_err());
421 assert!(b.module_exports(Pid(1234), "nope.dll").is_err());
422 }
423
424 #[test]
425 fn declared_module_exports_returned_sorted() {
426 let b = MockBackend::new(guest());
427 let ex = b.module_exports(Pid(1234), "target.exe").unwrap();
428 assert_eq!(ex, vec![("tick".to_string(), 0x1400001000)]);
429 }
430}