fix a bug (#93) in PtraceDetach when threads are enabled

src/ptrace.cc

@@ -237,8 +237,6 @@ long PtraceCallFunction(pid_t pid, unsigned long addr) {
       return -1;
     }
 
-    // std::cerr << "probe point is at " << reinterpret_cast<void *>(probe_)
-    //           << "\n";
     long code = 0;
     uint8_t *new_code_bytes = (uint8_t *)&code;
     new_code_bytes[0] = 0xff;  // CALL
@@ -260,34 +258,53 @@ long PtraceCallFunction(pid_t pid, unsigned long addr) {
   return newregs.rax;
 };
 
-void PtraceCleanup(pid_t pid) {
-  if (probe_ == 0) {
-    return;
-  }
-
-  user_regs_struct oldregs = PtraceGetRegs(pid);
-  long orig_code = PtracePeek(pid, oldregs.rip);
-
-  user_regs_struct newregs = oldregs;
-  newregs.rax = SYS_munmap;
-  newregs.rdi = probe_;         // addr
-  newregs.rsi = getpagesize();  // len
-
-  PauseChildThreads(pid);
-
-  PtracePoke(pid, oldregs.rip, syscall_x86);
-  PtraceSetRegs(pid, newregs);
-  PtraceSingleStep(pid);
-  PtracePoke(pid, oldregs.rip, orig_code);
-  PtraceSetRegs(pid, oldregs);
+// Like PtraceDetach(), but ignore errors.
+static inline void SafeDetach(pid_t pid) noexcept {
+  ptrace(PTRACE_DETACH, pid, 0, 0);
+}
 
-  ResumeChildThreads(pid);
-  PtraceDetach(pid);
+void PtraceCleanup(pid_t pid) noexcept {
+  // Clean up the memory area allocated by AllocPage().
+  if (probe_ != 0) {
+    try {
+      const user_regs_struct oldregs = PtraceGetRegs(pid);
+      const long orig_code = PtracePeek(pid, oldregs.rip);
+
+      user_regs_struct newregs = oldregs;
+      newregs.rax = SYS_munmap;
+      newregs.rdi = probe_;         // addr
+      newregs.rsi = getpagesize();  // len
+
+      // Prepare to run munmap(2) syscall.
+      PauseChildThreads(pid);
+      PtracePoke(pid, oldregs.rip, syscall_x86);
+      PtraceSetRegs(pid, newregs);
+
+      // Call the syscall and check the return value.
+      PtraceSingleStep(pid);
+      if (PtraceGetRegs(pid).rax == 0) {
+        probe_ = 0;
+      } else {
+        std::cerr << "Warning: failed to munmap(2) trampoline page! Please "
+                     "report this on GitHub.\n";
+      }
+
+      // Clean up and resume the child threads.
+      PtracePoke(pid, oldregs.rip, orig_code);
+      PtraceSetRegs(pid, oldregs);
+      ResumeChildThreads(pid);
+
+    } catch (...) {
+      // If the process has already exited, then we'll get a ptrace error, which
+      // can be safely ignored. This *should* happen at the initial
+      // PtraceGetRegs() call, but we wrap the entire block to be safe.
+    }
+  }
 
-  probe_ = 0;
+  SafeDetach(pid);
 }
 #else
-void PtraceCleanup(pid_t pid) { PtraceDetach(pid); }
+void PtraceCleanup(pid_t pid) noexcept { SafeDetach(pid); }
 #endif
 
 }  // namespace pyflame

src/ptrace.h

@@ -57,10 +57,10 @@ void PtraceCont(pid_t pid);
 void PtraceSingleStep(pid_t pid);
 
 #if ENABLE_THREADS
-// call a function pointer
+// Call a function pointer.
 long PtraceCallFunction(pid_t pid, unsigned long addr);
 #endif
 
-// maybe dealloc the page allocated in PtraceCallFunction();
-void PtraceCleanup(pid_t pid);
+// Detach, and maybe dealloc the page allocated in PtraceCallFunction();
+void PtraceCleanup(pid_t pid) noexcept;
 }  // namespace pyflame

tests/test_end_to_end.py

@@ -271,12 +271,16 @@ def test_sample_not_python(not_python):
     assert proc.returncode == 1
 
 
-@pytest.mark.parametrize('force_abi', [(False, ), (True, )])
-def test_trace(force_abi):
+@pytest.mark.parametrize('force_abi,trace_threads',
+                         [(False, False), (False, True), (True, False),
+                          (True, True)])
+def test_trace(force_abi, trace_threads):
     args = ['./src/pyflame']
     if force_abi:
         abi_string = '%d%d' % sys.version_info[:2]
         args.extend(['--abi', abi_string])
+    if trace_threads:
+        args.append('--threads')
     args.extend(['-t', python_command(), 'tests/exit_early.py', '-s'])
 
     proc = subprocess.Popen(