From 8d51a1e00b831370304f7d9b863db8d279b0849d Mon Sep 17 00:00:00 2001
From: Robert O'Callahan <robert@ocallahan.org>
Date: Tue, 5 Sep 2023 14:06:37 +1200
Subject: [PATCH] Fix the unexpected-exit path in
 `Task::exit_syscall_and_prepare_restart`

---
 src/Monkeypatcher.cc |  4 +++-
 src/Task.cc          | 11 +++++++++--
 src/Task.h           |  7 +++++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/Monkeypatcher.cc b/src/Monkeypatcher.cc
index c34293b1327..b12cd2eefaa 100644
--- a/src/Monkeypatcher.cc
+++ b/src/Monkeypatcher.cc
@@ -1218,7 +1218,9 @@ bool Monkeypatcher::try_patch_syscall_aarch64(RecordTask* t, bool entering_sysca
   auto success = patch_syscall_with_hook(*this, t, syscall_hooks[0], ip, 4, 0);
   if (!success && entering_syscall) {
     // Need to reenter the syscall to undo exit_syscall_and_prepare_restart
-    t->enter_syscall();
+    if (!t->enter_syscall()) {
+      return false;
+    }
   }
 
   if (!success) {
diff --git a/src/Task.cc b/src/Task.cc
index bb962e84098..b33f74e75b1 100644
--- a/src/Task.cc
+++ b/src/Task.cc
@@ -972,12 +972,12 @@ bool Task::exit_syscall_and_prepare_restart() {
   // This exits the hijacked SYS_gettid.  Now the tracee is
   // ready to do our bidding.
   if (!exit_syscall()) {
-    // The tracee suddenly exited. To get this to replay correctly, we need to
+    // The tracee unexpectedly exited. To get this to replay correctly, we need to
     // make it look like we really entered the syscall. Then
     // handle_ptrace_exit_event will record something appropriate.
     r.set_original_syscallno(syscallno);
     r.set_syscall_result(-ENOSYS);
-    set_regs(r);
+    override_regs_during_exit(r);
     return false;
   }
   LOG(debug) << "exit_syscall_and_prepare_restart done";
@@ -1578,6 +1578,13 @@ void Task::set_regs(const Registers& regs) {
   }
 }
 
+void Task::override_regs_during_exit(const Registers& regs) {
+  ASSERT(this, !is_stopped_);
+  orig_syscallno_dirty = true;
+  registers_dirty = true;
+  registers = regs;
+}
+
 void Task::flush_regs() {
   if (registers_dirty) {
     LOG(debug) << "Flushing registers for tid " << tid << " " << registers;
diff --git a/src/Task.h b/src/Task.h
index 62f3a53758d..c3cdbf1a1bd 100644
--- a/src/Task.h
+++ b/src/Task.h
@@ -598,6 +598,13 @@ class Task {
   /** Set the tracee's registers to |regs|. Lazy. */
   void set_regs(const Registers& regs);
 
+  /** Set the tracee's registers to |regs|. The task
+   * is known to be exiting (running towards, or actually in, a
+   * not-yet-reported PTRACE_EVENT_EXIT or reap) so its real registers
+   * won't change underneath us, but we need to override those
+   * registers with our values for recording purposes. */
+  void override_regs_during_exit(const Registers& regs);
+
   /** Ensure registers are flushed back to the underlying task. */
   void flush_regs();