1use serde::de::DeserializeOwned;
4use serde::{Deserialize, Serialize};
5use std::io::{self, Read, Write};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
8pub struct Pid(pub u32);
9
10impl std::fmt::Display for Pid {
11 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12 write!(f, "{}", self.0)
13 }
14}
15
16impl From<u32> for Pid {
17 fn from(v: u32) -> Self {
18 Pid(v)
19 }
20}
21
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23pub struct ProcessInfo {
24 pub pid: Pid,
25 pub name: String,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct ModuleInfo {
30 pub name: String,
31 pub base: u64,
32 pub size: u64,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36pub struct MemRegion {
37 pub base: u64,
38 pub size: u64,
39 pub readable: bool,
40 pub writable: bool,
41 pub executable: bool,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
45pub struct Diagnostics {
46 pub connector: String,
47 pub reads: u64,
48 pub writes: u64,
49 pub unsupported_ops: u64,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
53pub struct GuestInjectInfo {
54 pub method: String,
55 pub pid: Pid,
56 pub remote_base: Option<u64>,
57 pub notes: Vec<String>,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub enum ProtoError {
62 NoSuchProcess {
63 pid: Option<u32>,
64 name: Option<String>,
65 },
66 NoSuchModule {
67 pid: u32,
68 module: String,
69 },
70 ReadFailed {
71 addr: u64,
72 len: u64,
73 reason: String,
74 },
75 WriteFailed {
76 addr: u64,
77 reason: String,
78 },
79 Unsupported {
80 op: String,
81 },
82 Backend {
83 message: String,
84 },
85}
86
87impl std::fmt::Display for ProtoError {
88 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89 match self {
90 ProtoError::NoSuchProcess { pid, name } => {
91 write!(f, "no such process (pid={pid:?}, name={name:?})")
92 }
93 ProtoError::NoSuchModule { pid, module } => {
94 write!(f, "no such module {module:?} in pid {pid}")
95 }
96 ProtoError::ReadFailed { addr, len, reason } => {
97 write!(f, "read of {len} bytes at {addr:#x} failed: {reason}")
98 }
99 ProtoError::WriteFailed { addr, reason } => {
100 write!(f, "write at {addr:#x} failed: {reason}")
101 }
102 ProtoError::Unsupported { op } => {
103 write!(
104 f,
105 "unsupported operation: {op} requires guest execution, which memflow cannot perform"
106 )
107 }
108 ProtoError::Backend { message } => write!(f, "backend error: {message}"),
109 }
110 }
111}
112
113impl std::error::Error for ProtoError {}
114
115#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
116pub enum Request {
117 Ping,
118 ListProcesses,
119 ProcessByPid(Pid),
120 ProcessByName(String),
121 ModuleList(Pid),
122 ModuleByName(Pid, String),
123 ModuleExports(Pid, String),
124 Read {
125 pid: Pid,
126 addr: u64,
127 len: u64,
128 },
129 Write {
130 pid: Pid,
131 addr: u64,
132 data: Vec<u8>,
133 },
134 MemoryMap(Pid),
135 Diagnostics,
136 Scan {
137 pid: Pid,
138 pattern: String,
139 },
140 Resolve {
141 pid: Pid,
142 base: u64,
143 offsets: Vec<u64>,
144 },
145 GuestInject {
146 config_toml: String,
147 payload_image: Vec<u8>,
148 },
149 ReportUnsupported {
150 op: String,
151 },
152}
153
154impl Request {
155 pub fn retry_safe(&self) -> bool {
156 matches!(
157 self,
158 Request::Ping
159 | Request::ListProcesses
160 | Request::ProcessByPid(_)
161 | Request::ProcessByName(_)
162 | Request::ModuleList(_)
163 | Request::ModuleByName(_, _)
164 | Request::ModuleExports(_, _)
165 | Request::Read { .. }
166 | Request::MemoryMap(_)
167 | Request::Diagnostics
168 | Request::Scan { .. }
169 | Request::Resolve { .. }
170 )
171 }
172}
173
174#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
175pub enum Response {
176 Pong,
177 Processes(Vec<ProcessInfo>),
178 Process(ProcessInfo),
179 Modules(Vec<ModuleInfo>),
180 Module(ModuleInfo),
181 Exports(Vec<(String, u64)>),
182 Data(Vec<u8>),
183 Written(u64),
184 MemoryMap(Vec<MemRegion>),
185 Diagnostics(Diagnostics),
186 Err(ProtoError),
187 ScanHits(Vec<u64>),
188 Resolved { address: u64, value: Vec<u8> },
189 GuestInjected(GuestInjectInfo),
190}
191
192pub const MAX_MSG_LEN: u32 = 64 * 1024 * 1024;
193
194fn enc_err(e: bincode::Error) -> io::Error {
195 io::Error::new(io::ErrorKind::InvalidData, e)
196}
197
198pub fn write_msg<W: Write, T: Serialize>(w: &mut W, msg: &T) -> io::Result<()> {
199 let bytes = bincode::serialize(msg).map_err(enc_err)?;
200 if bytes.len() as u64 > MAX_MSG_LEN as u64 {
201 return Err(io::Error::new(
202 io::ErrorKind::InvalidData,
203 format!(
204 "message too large: {} bytes (max {MAX_MSG_LEN})",
205 bytes.len()
206 ),
207 ));
208 }
209 w.write_all(&(bytes.len() as u32).to_le_bytes())?;
210 w.write_all(&bytes)?;
211 w.flush()
212}
213
214pub fn read_msg<R: Read, T: DeserializeOwned>(r: &mut R) -> io::Result<T> {
215 let mut len_buf = [0u8; 4];
216 r.read_exact(&mut len_buf)?;
217 let len = u32::from_le_bytes(len_buf);
218 if len > MAX_MSG_LEN {
219 return Err(io::Error::new(
220 io::ErrorKind::InvalidData,
221 format!("incoming message length {len} exceeds max {MAX_MSG_LEN}"),
222 ));
223 }
224 let mut buf = vec![0u8; len as usize];
225 r.read_exact(&mut buf)?;
226 bincode::deserialize(&buf).map_err(enc_err)
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232 use std::io::Cursor;
233
234 fn roundtrip_req(req: Request) {
235 let mut buf = Vec::new();
236 write_msg(&mut buf, &req).unwrap();
237 let mut cur = Cursor::new(buf);
238 let got: Request = read_msg(&mut cur).unwrap();
239 assert_eq!(req, got);
240 }
241
242 fn roundtrip_resp(resp: Response) {
243 let mut buf = Vec::new();
244 write_msg(&mut buf, &resp).unwrap();
245 let mut cur = Cursor::new(buf);
246 let got: Response = read_msg(&mut cur).unwrap();
247 assert_eq!(resp, got);
248 }
249
250 #[test]
251 fn requests_roundtrip() {
252 roundtrip_req(Request::Ping);
253 roundtrip_req(Request::ListProcesses);
254 roundtrip_req(Request::ProcessByPid(Pid(1234)));
255 roundtrip_req(Request::ProcessByName("target.exe".into()));
256 roundtrip_req(Request::ModuleList(Pid(1)));
257 roundtrip_req(Request::ModuleByName(Pid(1), "ntdll.dll".into()));
258 roundtrip_req(Request::ModuleExports(Pid(1), "kernel32.dll".into()));
259 roundtrip_req(Request::Read {
260 pid: Pid(7),
261 addr: 0x1400010000,
262 len: 64,
263 });
264 roundtrip_req(Request::Write {
265 pid: Pid(7),
266 addr: 0xdead,
267 data: vec![1, 2, 3, 4],
268 });
269 roundtrip_req(Request::MemoryMap(Pid(7)));
270 roundtrip_req(Request::Diagnostics);
271 roundtrip_req(Request::Scan {
272 pid: Pid(7),
273 pattern: "DE CA ?? 00".into(),
274 });
275 roundtrip_req(Request::Resolve {
276 pid: Pid(7),
277 base: 0x1000,
278 offsets: vec![0x10, 0x18],
279 });
280 roundtrip_req(Request::GuestInject {
281 config_toml: "[injection]\ndomain = \"guest\"\nmethod = \"manual-map\"\n".into(),
282 payload_image: vec![0x4d, 0x5a],
283 });
284 roundtrip_req(Request::ReportUnsupported {
285 op: "VirtualAllocEx".into(),
286 });
287 }
288
289 #[test]
290 fn responses_roundtrip() {
291 roundtrip_resp(Response::Pong);
292 roundtrip_resp(Response::Processes(vec![ProcessInfo {
293 pid: Pid(1234),
294 name: "target.exe".into(),
295 }]));
296 roundtrip_resp(Response::Process(ProcessInfo {
297 pid: Pid(1),
298 name: "a".into(),
299 }));
300 roundtrip_resp(Response::Modules(vec![ModuleInfo {
301 name: "m.dll".into(),
302 base: 0x1400000000,
303 size: 0x80000,
304 }]));
305 roundtrip_resp(Response::Module(ModuleInfo {
306 name: "m".into(),
307 base: 0,
308 size: 1,
309 }));
310 roundtrip_resp(Response::Exports(vec![("add".into(), 0x1000)]));
311 roundtrip_resp(Response::Data(vec![0xde, 0xad, 0xbe, 0xef]));
312 roundtrip_resp(Response::Written(64));
313 roundtrip_resp(Response::MemoryMap(vec![MemRegion {
314 base: 0x1000,
315 size: 0x2000,
316 readable: true,
317 writable: true,
318 executable: false,
319 }]));
320 roundtrip_resp(Response::Diagnostics(Diagnostics::default()));
321 roundtrip_resp(Response::Err(ProtoError::Unsupported {
322 op: "VirtualAllocEx".into(),
323 }));
324 roundtrip_resp(Response::ScanHits(vec![0x1400010100, 0x1400010200]));
325 roundtrip_resp(Response::Resolved {
326 address: 0x1400010290,
327 value: vec![0x39, 5, 0, 0],
328 });
329 roundtrip_resp(Response::GuestInjected(GuestInjectInfo {
330 method: "manual-map".into(),
331 pid: Pid(7),
332 remote_base: Some(0x1000),
333 notes: vec!["ok".into()],
334 }));
335 }
336
337 #[test]
338 fn oversized_length_prefix_is_rejected() {
339 let mut bytes = (u32::MAX).to_le_bytes().to_vec();
340 bytes.extend_from_slice(&[0u8; 8]);
341 let mut cur = Cursor::new(bytes);
342 let got: io::Result<Request> = read_msg(&mut cur);
343 assert!(got.is_err());
344 }
345
346 #[test]
347 fn truncated_stream_reports_eof() {
348 let buf: Vec<u8> = vec![10, 0, 0, 0, 1, 2];
349 let mut cur = Cursor::new(buf);
350 let got: io::Result<Request> = read_msg(&mut cur);
351 assert_eq!(got.unwrap_err().kind(), io::ErrorKind::UnexpectedEof);
352 }
353
354 #[test]
355 fn two_messages_back_to_back() {
356 let mut buf = Vec::new();
357 write_msg(&mut buf, &Request::Ping).unwrap();
358 write_msg(&mut buf, &Request::ProcessByPid(Pid(9))).unwrap();
359 let mut cur = Cursor::new(buf);
360 assert_eq!(read_msg::<_, Request>(&mut cur).unwrap(), Request::Ping);
361 assert_eq!(
362 read_msg::<_, Request>(&mut cur).unwrap(),
363 Request::ProcessByPid(Pid(9))
364 );
365 }
366
367 #[test]
368 fn mutating_requests_are_not_retry_safe() {
369 assert!(
370 !Request::Write {
371 pid: Pid(7),
372 addr: 0x1000,
373 data: vec![1, 2, 3],
374 }
375 .retry_safe()
376 );
377 assert!(
378 !Request::GuestInject {
379 config_toml: String::new(),
380 payload_image: Vec::new(),
381 }
382 .retry_safe()
383 );
384 assert!(!Request::ReportUnsupported { op: "x".into() }.retry_safe());
385 assert!(Request::ListProcesses.retry_safe());
386 assert!(
387 Request::Read {
388 pid: Pid(7),
389 addr: 0x1000,
390 len: 8,
391 }
392 .retry_safe()
393 );
394 }
395}