From 99999346ed725fc78f1f3d485beaae8e224d5a23 Mon Sep 17 00:00:00 2001
From: Elnath <elnathbeta@gmail.com>
Date: Sat, 17 May 2025 18:09:57 +0200
Subject: [PATCH] Stepi now returning running state for handling program exit

---
 Cargo.lock          |  7 +++++++
 Cargo.toml          |  1 +
 src/debug_target.rs | 20 +++++++++++++-------
 src/main.rs         |  9 ++++++++-
 4 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1456507..f451893 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -37,6 +37,7 @@ name = "bdb"
 version = "0.0.0"
 dependencies = [
  "color-eyre",
+ "either",
  "nix",
  "thiserror",
 ]
@@ -95,6 +96,12 @@ dependencies = [
  "tracing-error",
 ]
 
+[[package]]
+name = "either"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
+
 [[package]]
 name = "eyre"
 version = "0.6.12"
diff --git a/Cargo.toml b/Cargo.toml
index 796a5f9..07d60a5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,6 +8,7 @@ publish = false
 nix = { version = "0.29.0", features = ["process", "ptrace"] }
 thiserror = "2.0.12"
 color-eyre = "0.6.3"
+either = "1.15.0"
 
 [build-dependencies]
 color-eyre = "0.6.3"
\ No newline at end of file
diff --git a/src/debug_target.rs b/src/debug_target.rs
index 4e7685a..724081d 100644
--- a/src/debug_target.rs
+++ b/src/debug_target.rs
@@ -1,5 +1,5 @@
+use either::{Either, Left, Right};
 use nix::libc::user_regs_struct;
-use nix::sys::signal::Signal;
 use nix::sys::wait::{waitid, Id, WaitPidFlag, WaitStatus};
 use nix::unistd::Pid;
 
@@ -40,6 +40,7 @@ impl DebugState for Stopped {
 	fn pid(&self) -> Pid { self.pid }
 }
 
+#[allow(dead_code)]
 impl DebugTarget<Stopped> {
 	pub fn new(pid: Pid) -> Result<Self, WaitError> {
 		waitid(Id::Pid(pid), WaitPidFlag::WSTOPPED).map_err(WaitError)?;
@@ -51,12 +52,9 @@ impl DebugTarget<Stopped> {
 		Ok(DebugTarget { state: Running { pid: self.state.pid } })
 	}
 
-	pub fn stepi(self) -> Result<Self, DebugError> {
+	pub fn stepi(self) -> Result<DebugTarget<Running>, PTraceError> {
 		nix::sys::ptrace::step(self.state.pid, None).map_err(PTraceError)?;
-		match waitid(Id::Pid(self.state.pid), WaitPidFlag::WSTOPPED).map_err(WaitError)? {
-			WaitStatus::PtraceEvent(_pid, Signal::SIGTRAP, _c_event) => Ok(self),
-			status => Err(DebugError::UnexpectedWaitStatus(status)),
-		}
+		Ok(DebugTarget { state: Running { pid: self.state.pid } })
 	}
 
 	pub fn get_registers(&self) -> Result<user_regs_struct, PTraceError> {
@@ -72,8 +70,16 @@ impl DebugState for Running {
 	fn pid(&self) -> Pid { self.pid }
 }
 
-
+#[allow(dead_code)]
 impl DebugTarget<Running> {
+	pub fn wait_for_something(self) -> Result<Either<DebugTarget<Stopped>, i32>, DebugError> {
+		match waitid(Id::Pid(self.state.pid), WaitPidFlag::WSTOPPED | WaitPidFlag::WEXITED).map_err(WaitError)? {
+			WaitStatus::Exited(_pid, exit_code) => Ok(Right(exit_code)),
+			WaitStatus::PtraceEvent(_pid, _signal, _c_event) => Ok(Left(DebugTarget { state: Stopped { pid: self.state.pid } })),
+			status => Err(DebugError::UnexpectedWaitStatus(status)),
+		}
+	}
+
 	pub fn wait_for_exit(self) -> Result<i32, DebugError> {
 		match waitid(Id::Pid(self.state.pid), WaitPidFlag::WEXITED).map_err(WaitError)? {
 			WaitStatus::Exited(_pid, exit_code) => Ok(exit_code),
diff --git a/src/main.rs b/src/main.rs
index 9fec6ea..e21a701 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,7 @@ mod debug_target;
 
 use crate::debug_target::{DebugTarget, Stopped};
 use color_eyre::eyre::eyre;
+use either::Either;
 use nix::libc::user_regs_struct;
 use nix::sys::ptrace::*;
 use nix::sys::signal::Signal::*;
@@ -17,7 +18,13 @@ fn single_step_all(mut target: DebugTarget<Stopped>) -> color_eyre::Result<()> {
 		instruction_number += 1;
 		let regs = target.get_registers()?;
 		println!("🔎 [==> {}] rip= {:#016x}, rax = 0x{rax:x} ({rax})", instruction_number, regs.rip, rax = regs.rax);
-		target = target.stepi()?;
+		match target.stepi()?.wait_for_something()? {
+			Either::Left(t) => target = t,
+			Either::Right(exit_code) => {
+				println!("👋 Child exited with code {exit_code}");
+				return Ok(());
+			}
+		}
 	}
 }