Added one-time breakpoints and continue from them

src/debug_target/mod.rs

@@ -23,6 +23,12 @@ pub enum DebugError {
 
 	#[error("Child stopped with status {0:?}, but was not expecting to catch this one")]
 	UnexpectedWaitStatus(nix::sys::wait::WaitStatus),
+
+	#[error("Tried to insert a breakpoint at {address:#x} but one was already set for the same address")]
+	DuplicateBreakpoint { address: u64 },
+
+	#[error("Reference a breakpoint at {address:#x}, but not breakpoint exists there")]
+	NonExistingBreakpoint { address: u64 },
 }
 
 

src/debug_target/running_target.rs

@@ -2,15 +2,17 @@ use crate::debug_target::{DebugError, ExitedTarget, StoppedTarget, WaitError};
 use either::{Either, Left, Right};
 use nix::sys::wait::{waitid, Id, WaitPidFlag, WaitStatus};
 use nix::unistd::Pid;
+use std::collections::HashMap;
 
 pub struct RunningTarget {
 	pub pid: Pid,
+	pub breakpoints: HashMap<u64, u8>,
 }
 
 #[allow(dead_code)]
 impl RunningTarget {
 	fn to_stopped(self) -> StoppedTarget {
-		StoppedTarget { pid: self.pid }
+		StoppedTarget { pid: self.pid, breakpoints: self.breakpoints }
 	}
 
 	pub fn wait_for_something(self) -> Result<Either<StoppedTarget, ExitedTarget>, DebugError> {

src/debug_target/stopped_target.rs

@@ -1,33 +1,49 @@
 use crate::debug_target::{DebugError, PTraceError, RunningTarget, WaitError};
 use crate::syscall_info::{syscall_info, SyscallInfo, SyscallInfoError};
-use libc::user_regs_struct;
+use libc::{c_long, user_regs_struct};
 use nix::sys::wait::{waitid, Id, WaitPidFlag};
 use nix::unistd::Pid;
+use std::collections::HashMap;
+use std::ffi::c_void;
 
 pub struct StoppedTarget {
 	pub pid: Pid,
+	pub breakpoints: HashMap<u64, u8>,
 }
 
 
 #[allow(dead_code)]
 impl StoppedTarget {
 	fn to_running(self) -> RunningTarget {
-		RunningTarget { pid: self.pid }
+		RunningTarget { pid: self.pid, breakpoints: self.breakpoints }
 	}
 
 	pub fn new(pid: Pid) -> Result<Self, DebugError> {
 		waitid(Id::Pid(pid), WaitPidFlag::WSTOPPED).map_err(WaitError)?;
 		// Needed for waiting on syscalls and apparently also for getting syscall info (can not get it to work without this)
 		nix::sys::ptrace::setoptions(pid, nix::sys::ptrace::Options::PTRACE_O_TRACESYSGOOD).map_err(PTraceError)?;
-		Ok(Self { pid })
+		Ok(Self { pid, breakpoints: HashMap::new() })
 	}
 
-	pub fn cont(self) -> Result<RunningTarget, PTraceError> {
+	pub fn on_breakpoint(&self) -> Result<bool, PTraceError> {
+		let rip = self.get_registers()?.rip;
+		Ok(self.breakpoints.contains_key(&(rip - 1)))
+	}
+
+	pub fn cont(mut self) -> Result<RunningTarget, DebugError> {
+		let mut registers = self.get_registers()?;
+		if self.breakpoints.contains_key(&(registers.rip - 1)) { // We are on a breakpoint, we remove it
+			self = self.remove_breakpoint(registers.rip - 1)?;
+			registers.rip -= 1;
+			nix::sys::ptrace::setregs(self.pid, registers).map_err(PTraceError)?;
+		}
+
 		nix::sys::ptrace::cont(self.pid, None).map_err(PTraceError)?;
 		Ok(self.to_running())
 	}
 
 	pub fn stepi(self) -> Result<RunningTarget, PTraceError> {
+		todo!("Take into account breakpoints also for stepi");
 		nix::sys::ptrace::step(self.pid, None).map_err(PTraceError)?;
 		Ok(self.to_running())
 	}
@@ -44,4 +60,31 @@ impl StoppedTarget {
 	pub fn get_syscall_info(&self) -> Result<SyscallInfo, SyscallInfoError> {
 		Ok(syscall_info(self.pid.as_raw())?)
 	}
+
+	#[cfg(target_endian = "little")] // With a bit more work it could be implemented on both architectures, but I only have little-endian computers 🤷
+	pub fn add_breakpoint(mut self, address: u64) -> Result<Self, DebugError> {
+		if self.breakpoints.contains_key(&address) {
+			return Err(DebugError::DuplicateBreakpoint { address });
+		}
+		let orig_bytes = nix::sys::ptrace::read(self.pid, address as *mut c_void).map_err(PTraceError)?;
+		let target_byte: u8 = (orig_bytes & 0xFF as c_long) as u8;
+		let new_content = (orig_bytes & (!0xff as c_long)) | (0xCC as c_long);
+		nix::sys::ptrace::write(self.pid, address as *mut c_void, new_content).map_err(PTraceError)?;
+
+		self.breakpoints.insert(address, target_byte);
+		Ok(self)
+	}
+
+	#[cfg(target_endian = "little")] // Same as add_breakpoint
+	pub fn remove_breakpoint(mut self, address: u64) -> Result<Self, DebugError> {
+		match self.breakpoints.remove(&address) {
+			None => Err(DebugError::NonExistingBreakpoint { address }),
+			Some(original_byte) => {
+				let content = nix::sys::ptrace::read(self.pid, address as *mut c_void).map_err(PTraceError)?;
+				let new_content = (content & (!0xff as c_long)) | (original_byte as c_long);
+				nix::sys::ptrace::write(self.pid, address as *mut c_void, new_content).map_err(PTraceError)?;
+				Ok(self)
+			}
+		}
+	}
 }

src/main.rs

@@ -47,9 +47,6 @@ fn strace(mut target: StoppedTarget) -> color_eyre::Result<()> {
 
 #[allow(dead_code)]
 fn breakpoint_fun(child_pid: Pid) -> color_eyre::Result<()> {
-	println!("⏳️ Waiting for child to be ready");
-	waitid(Id::Pid(child_pid), WaitPidFlag::WSTOPPED)?;
-
 	let address: u64 = 0x0000000000401019;
 	println!("🚧 Setting breakpoint at location 0x{address:x}");
 	let orig_bytes: [u8; 8] = read(child_pid, address as *mut c_void).expect("Breakpoint memory read").to_le_bytes();
@@ -109,26 +106,25 @@ fn main() -> color_eyre::Result<()> {
 			let target = StoppedTarget::new(child_pid)?;
 			println!("✔️ Child ready!");
 
-
+			let target = target.add_breakpoint(0x401019)?;
-			// println!("🔎 rip: {:#x}", target.get_registers()?.rip);
+			let target = target.cont()?.wait_for_something()?;
-			//
+			match target {
-			//
+				Either::Left(t) => {
-			// println!("⚙️ Executing until next syscall");
+					println!("🔎 rip: {:#x}", t.get_registers()?.rip);
-			// let target = target.cont_syscall()?.wait_for_syscall()?;
+					if t.on_breakpoint()? {
-			// println!("{:?}", target.get_syscall_info()?);
+						println!("🚧 We are on a breakpoint!")
-			// let target = target.cont_syscall()?.wait_for_syscall()?;
+					}
-			// println!("{:?}", target.get_syscall_info()?);
+					t.cont()?.wait_for_exit()?;
-			// println!("🔎 rip: {}", target.get_registers()?.rip);
+				}
-			//
+				Either::Right(ExitedTarget { exit_code, .. }) => {
-			// println!("⚙️ Continuing execution");
+					println!("👋 Child exited with code {exit_code}");
-			// let exit_code = target.cont()?.wait_for_exit()?;
+				}
-			// println!("👋 Child exited with code {exit_code}");
+			}
-			//
+			Ok(())
-			// Ok(())
 
 
 			// single_step_all(target)
-			strace(target)
+			// strace(target)
 			// breakpoint_fun(child_pid)
 		}
 		Err(e) => {