FIx 1

Author: pablogsal

src/memray/_memray/ghost_stack/src/ghost_stack.cpp

@@ -133,18 +133,41 @@ class GhostStackImpl {
     /**
      * Reset the shadow stack, restoring all original return addresses.
      *
-     * This is the normal reset path - it restores the original return addresses
-     * to the stack before clearing the shadow stack entries.
+     * On ARM64, stale trampolines may still fire after reset() because the LR
+     * register may have already been loaded with the trampoline address before
+     * we restored the stack location. We keep entries_ around to handle these
+     * stale trampolines gracefully.
+     *
+     * We restore ALL entries (not just 0 to tail-1) but only if the location
+     * still contains the trampoline address. This handles the case where a
+     * location was reused by a new frame after its original trampoline fired.
      */
     void reset() {
         if (trampolines_installed_) {
-            size_t tail = tail_.load(std::memory_order_acquire);
-            // With reversed order, iterate from 0 to tail (all entries below tail)
-            for (size_t i = 0; i < tail; ++i) {
-                *entries_[i].location = entries_[i].return_address;
+            uintptr_t tramp_addr = reinterpret_cast<uintptr_t>(ghost_ret_trampoline);
+
+            // Restore ALL entries whose locations still contain the trampoline.
+            // This handles both pending entries AND already-fired entries whose
+            // locations haven't been reused by new frames.
+            for (size_t i = 0; i < entries_.size(); ++i) {
+                uintptr_t current_value = *entries_[i].location;
+                // Strip PAC bits before comparison - on ARM64 with PAC enabled,
+                // the value read from stack may be PAC-signed while tramp_addr is not
+                uintptr_t stripped_value = ptrauth_strip(current_value);
+                if (stripped_value == tramp_addr) {
+                    *entries_[i].location = entries_[i].return_address;
+                }
             }
+
+            // Mark trampolines as not installed, but DON'T clear entries_!
+            // On ARM64, stale trampolines may still fire because LR was loaded
+            // before we restored the stack. Keep entries_ so we can still
+            // return the correct address.
+            trampolines_installed_ = false;
+
+            // Increment epoch to signal state change
+            epoch_.fetch_add(1, std::memory_order_release);
         }
-        clear_entries();
     }
 
 public: