From 727217ad422e5a757a77ec14ac93fe71e9208d74 Mon Sep 17 00:00:00 2001
From: jD91mZM2 <me@krake.one>
Date: Mon, 15 Jun 2020 17:14:52 +0200
Subject: [PATCH] Fix acid test-bench issues

---
 src/arch/x86_64/interrupt/exception.rs |  3 +--
 src/ptrace.rs                          | 17 +++++++-----
 src/scheme/proc.rs                     | 36 ++++++++++++++------------
 3 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/src/arch/x86_64/interrupt/exception.rs b/src/arch/x86_64/interrupt/exception.rs
index e132739..bd3f4a4 100644
--- a/src/arch/x86_64/interrupt/exception.rs
+++ b/src/arch/x86_64/interrupt/exception.rs
@@ -48,13 +48,12 @@ interrupt_stack!(non_maskable, stack, {
 });
 
 interrupt_stack!(breakpoint, stack, {
-    println!("Breakpoint trap");
-
     let guard = ptrace::set_process_regs(stack);
 
     if ptrace::breakpoint_callback(PTRACE_STOP_BREAKPOINT, None).is_none() {
         drop(guard);
 
+        println!("Breakpoint trap");
         stack.dump();
         ksignal(SIGTRAP);
     }
diff --git a/src/ptrace.rs b/src/ptrace.rs
index 521f154..51f4e78 100644
--- a/src/ptrace.rs
+++ b/src/ptrace.rs
@@ -191,8 +191,8 @@ struct Breakpoint {
     flags: PtraceFlags
 }
 
-/// Continue the process with the specified ID
-pub fn cont(pid: ContextId) {
+/// Clear any breakpoints for the process with the specified ID
+pub fn clear_breakpoint(pid: ContextId) {
     let sessions = sessions();
     let session = match sessions.get(&pid) {
         Some(session) => session,
@@ -203,6 +203,15 @@ pub fn cont(pid: ContextId) {
     // Remove the breakpoint to make sure any yet unreached but
     // obsolete breakpoints don't stop the program.
     data.breakpoint = None;
+}
+
+// TODO: All these small functions should be moved to be on the session instance
+pub fn notify(pid: ContextId) {
+    let sessions = sessions();
+    let session = match sessions.get(&pid) {
+        Some(session) => session,
+        None => return
+    };
 
     session.tracee.notify();
 }
@@ -218,10 +227,6 @@ pub fn set_breakpoint(pid: ContextId, flags: PtraceFlags, should_continue: bool)
         reached: false,
         flags
     });
-
-    if should_continue {
-        session.tracee.notify();
-    }
 }
 
 /// Wait for the tracee to stop. If an event occurs, it returns a copy
diff --git a/src/scheme/proc.rs b/src/scheme/proc.rs
index c99e525..b198c69 100644
--- a/src/scheme/proc.rs
+++ b/src/scheme/proc.rs
@@ -47,8 +47,9 @@ fn with_context_mut<F, T>(pid: ContextId, callback: F) -> Result<T>
     }
     callback(&mut context)
 }
-fn try_stop_context<F, T>(pid: ContextId, restart_after: bool, mut callback: F) -> Result<T>
-    where F: FnMut(&mut Context) -> Result<T>
+fn try_stop_context<F, T>(pid: ContextId, mut callback: F) -> Result<T>
+where
+    F: FnMut(&mut Context) -> Result<T>,
 {
     // Stop process
     let (was_stopped, mut running) = with_context_mut(pid, |context| {
@@ -72,9 +73,7 @@ fn try_stop_context<F, T>(pid: ContextId, restart_after: bool, mut callback: F)
 
         let ret = callback(context);
 
-        if !was_stopped || restart_after {
-            context.ptrace_stop = false;
-        }
+        context.ptrace_stop = was_stopped;
 
         ret
     })
@@ -342,7 +341,7 @@ impl Scheme for ProcScheme {
                         let fx = context.arch.get_fx_regs().unwrap_or_default();
                         Ok((Output { float: fx }, mem::size_of::<FloatRegisters>()))
                     })?,
-                    RegsKind::Int => try_stop_context(info.pid, false, |context| match unsafe { ptrace::regs_for(&context) } {
+                    RegsKind::Int => try_stop_context(info.pid, |context| match unsafe { ptrace::regs_for(&context) } {
                         None => {
                             println!("{}:{}: Couldn't read registers from stopped process", file!(), line!());
                             Err(Error::new(ENOTRECOVERABLE))
@@ -443,7 +442,7 @@ impl Scheme for ProcScheme {
                         *(buf as *const _ as *const IntRegisters)
                     };
 
-                    try_stop_context(info.pid, false, |context| match unsafe { ptrace::regs_for_mut(context) } {
+                    try_stop_context(info.pid, |context| match unsafe { ptrace::regs_for_mut(context) } {
                         None => {
                             println!("{}:{}: Couldn't read registers from stopped process", file!(), line!());
                             Err(Error::new(ENOTRECOVERABLE))
@@ -470,9 +469,7 @@ impl Scheme for ProcScheme {
                 let should_continue = !op.contains(PTRACE_FLAG_WAIT) || op.intersects(PTRACE_STOP_MASK);
 
                 if op.contains(PTRACE_STOP_SINGLESTEP) {
-                    // `true` to `try_stop_context` means we restart after, no
-                    // matter if it was or wasn't stopped before
-                    try_stop_context(info.pid, true, |context| {
+                    try_stop_context(info.pid, |context| {
                         match unsafe { ptrace::regs_for_mut(context) } {
                             None => {
                                 println!("{}:{}: Couldn't read registers from stopped process", file!(), line!());
@@ -484,17 +481,24 @@ impl Scheme for ProcScheme {
                             }
                         }
                     })?;
-                } else {
-                    // disable ptrace stop
-                    with_context_mut(info.pid, |context| {
-                        context.ptrace_stop = false;
-                        Ok(())
-                    })?;
                 }
 
                 // Set next breakpoint, and potentially restart tracee
                 if op.intersects(PTRACE_STOP_MASK) {
                     ptrace::set_breakpoint(info.pid, op, should_continue);
+                } else if should_continue {
+                    ptrace::clear_breakpoint(info.pid);
+                }
+
+                if should_continue {
+                    // disable the ptrace_stop flag, which is used in some cases
+                    with_context_mut(info.pid, |context| {
+                        context.ptrace_stop = false;
+                        Ok(())
+                    })?;
+
+                    // and notify the tracee's WaitCondition, which is used in other cases
+                    ptrace::notify(info.pid);
                 }
 
                 // And await the tracee, if requested to