Rust proc self memory search iterator

memsearch.rs

use std::{fs::{File}, io::{BufReader, self, Read, Seek, ErrorKind, BufRead}, vec};

struct MemSearch<R>
where
	R: Read
{
	inner: R,
	matches: Vec<Vec<u8>>,
	addr: usize,
	end_addr: Option<usize>,
	max_m: usize
}

impl<R> MemSearch<R>
where
	R: Read
{
	fn new(mut buffer: R, begin_addr: usize, end_addr: Option<usize>, matches: Vec<Vec<u8>>) -> Self {
		let max_m = matches.iter()
			.map(|x| x.iter())
			.map(Iterator::count)
			.max().unwrap_or(0);
		Self { 
			inner: buffer,
			matches,
			addr: begin_addr,
			end_addr,
			max_m,
		}
	}
}

impl<R> Iterator for MemSearch<R>
where
	R: Read
{
	type Item = io::Result<Vec<(usize, usize, Vec<u8>)>>;
	
	fn next(&mut self) -> Option<io::Result<Vec<(usize, usize, Vec<u8>)>>> {
		let mut buf: [u8; 0x1000] = [0; 0x1000];
		
		// we might want to move this to `MemSearch`.
		let mut start = 0;
		let mut end = buf.len();
		// let mut first_read = true;
		loop {
			// check if we already over the end_addr
			if let Some(end_addr) = self.end_addr {
				if self.addr > end_addr || self.addr+end > end_addr {
					return None;
				}
			}

			return match self.inner.read(&mut buf[start..end]) {
				Ok(size) => {
					// let xchk = |x| (self.off > x, self.off < x);// self.off > x && self.off + size < x;
					// unsafe { println!(">> readed {size} @ soft off = {:#x} ||| a={:?} b={:?} c={:?}", self.off, xchk(a_ptr), xchk(b_ptr), xchk(c_ptr)) };
					
					let found_safe = buf.len()-self.max_m;
					let mut save_last_bytes = true;
					let found: Vec<_> = self.matches
						.iter()
						.map(|m| {
							let m_slice = m.as_slice();
							// println!(">>> est wsz={}", buf.windows(m.len()).len());
							let found: Vec<_> = buf
								.windows(m.len())
								.enumerate()
								.filter(|(ipos, wnd)| *wnd == m_slice)
								.map(|(lpos, wnd)| {
									let pos = self.addr+lpos;	
									if lpos > found_safe {
										// we found something, no need to save the last bytes
										save_last_bytes = false;
									}
									(pos, lpos, m.clone())
								})
								.collect();
							found
						})
						.filter(|found| !found.is_empty())
						.flatten()
						.collect();

					start = if save_last_bytes {
						buf[..].copy_within(found_safe.., 0);
						self.max_m
					} else { 0 };

					if let Some(end_addr) = self.end_addr {
						end = if self.addr+end > end_addr { self.addr - end_addr } else { end };
					}
					self.addr += size;

					Some(Ok(found))
				},
				Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
				// Err(ref e) if e.kind() == ErrorKind::Uncategorized => None,
				Err(e) => None,
			}
		}

	}
}

const PROT_READ: u64 = 1 << 0;
const PROT_WRITE: u64 = 1 << 1;
const PROT_EXEC: u64 = 1 << 2;
const PROT_PRIVATE: u64 = 1 << 5;
const PROT_SHARED: u64 = 1 << 6;

#[derive(Debug)]
struct MapRegion {
	name: String,
	area: (usize, usize),
	prot: u64,
}

impl MapRegion {
	fn get_prot_str(&self) -> String {
		let mut b = [b'-'; 4];
		for ts in vec![
			(PROT_READ,    0, b'r'),
			(PROT_WRITE,   1, b'w'),
			(PROT_EXEC,    2, b'x'),
			(PROT_SHARED,  3, b's'),
			(PROT_PRIVATE, 3, b'p'),
		] {
			if self.prot & ts.0 != 0 {
				b[ts.1] = ts.2;
			}
		}
		unsafe { std::str::from_utf8_unchecked(&b[..]).to_string() }
	}
}

fn get_maps() -> io::Result<Vec<MapRegion>> {
	let f = File::open("/proc/self/maps")?;
	let mut regions = vec![];
	
	for line in BufReader::new(f).lines() {
		let line = line?;
		// println!("{}", line);
		let region: Vec<&str> = line.split_ascii_whitespace()
			.filter(|x| !x.is_empty())
			.collect();
		
		let range = match region.get(0) {
			None => continue,
			Some(range_str) => {
				let range_str = range_str.split_once("-").unwrap();
				(
					usize::from_str_radix(range_str.0, 16).unwrap(),
					usize::from_str_radix(range_str.1, 16).unwrap(),
				)
			},
		};

		let prot = match region.get(1) {
			None => continue,
			Some(prot_str) => {
				let prot_b = prot_str.as_bytes();
				let mut prot = 0;
				for tc in vec![
					(0, b'r', PROT_READ), 
					(1, b'w', PROT_WRITE),
					(2, b'x', PROT_EXEC),
					(3, b'p', PROT_PRIVATE),
					(3, b's', PROT_SHARED),
				] {
					if prot_b[tc.0] == tc.1 {
						prot |= tc.2;
					}
				};
				prot
			},
		};

		let mname = match region.last() {
			None => continue,
			Some(mname) => mname.to_string(),
		};

		// println!(">> {} {} {}", region[1], prot, prot & PROT_READ == 0);
		regions.push(MapRegion {
			name: mname,
			area: range,
			prot,
		});
		// println!("line = {:?} {:?}", range, prot);
	}

	Ok(regions)
}

static mut a_ptr: usize = 0;
static mut b_ptr: usize = 0;
static mut c_ptr: usize = 0;
// static mut g_begin: usize = 0;
// static mut g_end: usize = 0;
// fn chk_btw(x: usize) -> bool { unsafe { x > g_begin && x < g_end } }

fn main() {
	static a: &str = "Nyan cat uwu doing nyan-nyan thing foo bar baz there";
	static b: &str = "aaaaaa Nyan uwu cat doing nyan-nyan thing Foo bar baz";
	static c: &str = "bbbbbbbbbbbbbFoo bar Nyan uwu cat doing nyan-nyan thing baz";

	unsafe {
		a_ptr = a.as_ptr() as usize;
		b_ptr = b.as_ptr() as usize;
		c_ptr = c.as_ptr() as usize;
		println!(">> static ptr: {:#x}, {:#x}, {:#x}", a_ptr, b_ptr, c_ptr);
	}

	let mut f = File::open("/proc/self/mem").unwrap();
	let matches = vec![
		b"Nyan".to_vec(),
		b"uwu".to_vec(),
	];
	for (ir, region) in get_maps().unwrap()
		.into_iter()
		.filter(|reg| reg.prot & (PROT_READ|PROT_PRIVATE) != 0) // filter READ & PRIVATE only
		.into_iter().enumerate()
	{
		// if ir > 2 { break; } // debugging

		let begin = region.area.0; // region begin addr
		let end = region.area.1;   // region end addr

		// unsafe {
		// 	g_begin = begin;
		// 	g_end = end;
		// 	println!("> chk reg=({:#x},{:#x}) a={} b={} c={}", 
		// 		region.area.0, region.area.1,
		// 		chk_btw(a_ptr),
		// 		chk_btw(b_ptr),
		// 		chk_btw(c_ptr),
		// 	)
		// };

		f.seek(io::SeekFrom::Start(begin as u64)).unwrap();
		for matched in MemSearch::new(&f, begin, Some(end), matches.clone()) {
			match matched {
				Ok(results)  => {
					if results.is_empty() { continue; }
					
					println!("* area=({:#x},{:#x}) prot={} name={:?}", region.area.0, region.area.1, region.get_prot_str(), region.name);
					for (i, result) in results.into_iter().enumerate() {
						let dump = result.0 as *const [u8; 0x8];

						println!("\t#{}: found {:?} addr={:#x} pos={:?}\n\t\tdmp={:?}\n\t\tstr={:?}", 
							i,
							result.2, result.0, result.1,
							unsafe {*dump},
							
							// we can get a panic due to string character encoding so if you got one
							// comment the `from_utf8_unchecked` below.
							// "<enable the comment>",
							unsafe { std::str::from_utf8_unchecked(&(*dump)[..]) },
						);
					}
				}
				Err(err) => {
					println!("Ouchy, error! region={:?} err={:?}", region, err);
					break;
				}
			}
		}
		// break;
	}
}