Fix snapshot dump crash on stale jmethodID (PROF-14003)

ddprof-lib/src/main/cpp/findLibraryImpl.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef _FINDLIBRARYIMPL_H
+#define _FINDLIBRARYIMPL_H
+
+// Signal-handler-safe last-hit-index cache for findLibraryByAddress.
+//
+// Templated on CacheArray and CacheEntry so that the exact same algorithm
+// can be exercised from benchmarks (using lightweight fake types) without
+// pulling in JVM headers, while production uses CodeCacheArray/CodeCache.
+//
+// Requirements on CacheArray:
+//   int          count()     const  — number of live entries
+//   CacheEntry*  operator[](int i) const  — entry at index i (may be nullptr)
+//
+// Requirements on CacheEntry:
+//   bool  contains(const void* address) const
+//
+// Signal-safety: the last-hit index is a plain static volatile int.
+// DTLS initialisation for shared libraries calls calloc internally; if a
+// profiler signal fires on a thread whose TLS block has not been set up yet
+// while that thread is inside malloc, any thread_local access deadlocks on
+// the allocator lock.  A plain static volatile int avoids TLS entirely.
+// A benign race on the index is acceptable — the worst case is a cache miss
+// that falls through to the O(n) linear scan.
+
+template<typename CacheEntry, typename CacheArray>
+static inline CacheEntry* findLibraryByAddressImpl(const CacheArray& libs, const void* address) {
+    static volatile int last_hit = 0;
+    const int count = libs.count();
+    int hint = last_hit;
+    if (hint < count) {
+        CacheEntry* lib = libs[hint];
+        if (lib != nullptr && lib->contains(address)) {
+            return lib;
+        }
+    }
+    for (int i = 0; i < count; i++) {
+        CacheEntry* lib = libs[i];
+        if (lib != nullptr && lib->contains(address)) {
+            last_hit = i;
+            return lib;
+        }
+    }
+    return nullptr;
+}
+
+#endif // _FINDLIBRARYIMPL_H

ddprof-lib/src/main/cpp/flightRecorder.cpp

@@ -176,14 +176,18 @@ void Lookup::fillJavaMethodInfo(MethodInfo *mi, jmethodID method,
     bool entry = false;
     if (VMMethod::check_jmethodID(method) &&
         jvmti->GetMethodDeclaringClass(method, &method_class) == 0 &&
+        // GetMethodDeclaringClass may return a jclass wrapping a stale/garbage oop when the class was
+        // unloaded between sample capture and dump (TOCTOU race with class unloading). Guard against
+        // null handles before calling GetClassSignature.
+        method_class != NULL &&
         // On some older versions of J9, the JVMTI call to GetMethodDeclaringClass will return OK = 0, but when a
         // classloader is unloaded they free all JNIIDs. This means that anyone holding on to a jmethodID is
         // pointing to corrupt data and the behaviour is undefined.
         // The behaviour is adjusted so that when asgct() is used or if `-XX:+KeepJNIIDs` is specified,
         // when a classloader is unloaded, the jmethodIDs are not freed, but instead marked as -1.
-        // The nested check below is to mitigate these crashes.
-        // In more recent versions, the condition above will short-circuit safely.
-        ((!VM::isOpenJ9() || method_class != reinterpret_cast<jclass>(-1)) && jvmti->GetClassSignature(method_class, &class_name, NULL) == 0) &&
+        // The check below mitigates these crashes on J9.
+        (!VM::isOpenJ9() || method_class != reinterpret_cast<jclass>(-1)) &&
+        jvmti->GetClassSignature(method_class, &class_name, NULL) == 0 &&
         jvmti->GetMethodName(method, &method_name, &method_sig, NULL) == 0) {
 
       if (first_time) {

ddprof-lib/src/main/cpp/hotspot/vmStructs.cpp

@@ -952,6 +952,12 @@ bool VMMethod::check_jmethodID_hotspot(jmethodID id) {
         }
     }
     if (VMStructs::class_loader_data_offset() >= 0) {
+        // Verify the Klass at cpool_holder is readable over the full range we're about to access,
+        // catching freed/reclaimed Klass memory between check_jmethodID and the JVMTI calls (PROF-14003).
+        if (!SafeAccess::isReadableRange(cpool_holder,
+                                         (size_t)VMStructs::class_loader_data_offset() + sizeof(void*))) {
+            return false;
+        }
         cl_data = (const char *)SafeAccess::load(
             (void **)(cpool_holder + VMStructs::class_loader_data_offset()));
         if (cl_data == NULL) {

ddprof-lib/src/main/cpp/libraries.cpp

@@ -1,5 +1,6 @@
 #include "codeCache.h"
 #include "common.h"
+#include "findLibraryImpl.h"
 #include "hotspot/vmStructs.h"
 #include "libraries.h"
 #include "libraryPatcher.h"
@@ -100,12 +101,5 @@ CodeCache *Libraries::findLibraryByName(const char *lib_name) {
 }
 
 CodeCache *Libraries::findLibraryByAddress(const void *address) const {
-  const int native_lib_count = _native_libs.count();
-  for (int i = 0; i < native_lib_count; i++) {
-    CodeCache *lib = _native_libs[i];
-    if (lib != NULL && lib->contains(address)) {
-      return lib;
-    }
-  }
-  return NULL;
+  return findLibraryByAddressImpl<CodeCache>(_native_libs, address);
 }

ddprof-lib/src/test/cpp/libraries_ut.cpp

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <gtest/gtest.h>
+#include "../../main/cpp/codeCache.h"
+#include "../../main/cpp/findLibraryImpl.h"
+#include "../../main/cpp/gtest_crash_handler.h"
+
+static constexpr char LIBRARIES_TEST_NAME[] = "LibrariesTest";
+
+// Three static buffers to serve as fake "text segments".
+// Using static storage ensures the addresses are stable and non-overlapping.
+static char g_lib0_text[1024];
+static char g_lib1_text[1024];
+static char g_lib2_text[1024];
+
+class LibrariesTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        installGtestCrashHandler<LIBRARIES_TEST_NAME>();
+
+        _lib0 = new CodeCache("libfake0.so", 0, g_lib0_text, g_lib0_text + sizeof(g_lib0_text));
+        _lib1 = new CodeCache("libfake1.so", 1, g_lib1_text, g_lib1_text + sizeof(g_lib1_text));
+        _lib2 = new CodeCache("libfake2.so", 2, g_lib2_text, g_lib2_text + sizeof(g_lib2_text));
+
+        _libs.add(_lib0);
+        _libs.add(_lib1);
+        _libs.add(_lib2);
+    }
+
+    void TearDown() override {
+        restoreDefaultSignalHandlers();
+        delete _lib0;
+        delete _lib1;
+        delete _lib2;
+    }
+
+    CodeCacheArray _libs;
+    CodeCache* _lib0 = nullptr;
+    CodeCache* _lib1 = nullptr;
+    CodeCache* _lib2 = nullptr;
+};
+
+TEST_F(LibrariesTest, KnownAddress) {
+    // Address inside lib1's range should return lib1.
+    const void* addr = g_lib1_text + 512;
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(result, _lib1);
+}
+
+TEST_F(LibrariesTest, NullAddress) {
+    // NULL should return nullptr without crashing.
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, nullptr);
+    EXPECT_EQ(result, nullptr);
+}
+
+TEST_F(LibrariesTest, InvalidAddress) {
+    // An address outside all known ranges should return nullptr.
+    const void* addr = reinterpret_cast<const void*>(0x1);
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(result, nullptr);
+}
+
+TEST_F(LibrariesTest, FirstLibrary) {
+    const void* addr = g_lib0_text + 100;
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(result, _lib0);
+}
+
+TEST_F(LibrariesTest, LastLibrary) {
+    const void* addr = g_lib2_text + 900;
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(result, _lib2);
+}
+
+TEST_F(LibrariesTest, CacheHitOnRepeatedLookup) {
+    // First lookup primes the cache for lib1.
+    const void* addr = g_lib1_text + 256;
+    CodeCache* first = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(first, _lib1);
+
+    // Second lookup with the same range should still return lib1 (cache hit path).
+    CodeCache* second = findLibraryByAddressImpl<CodeCache>(_libs, g_lib1_text + 700);
+    EXPECT_EQ(second, _lib1);
+}
+
+TEST_F(LibrariesTest, CacheMissFallsBackToLinearScan) {
+    // Prime cache for lib0.
+    findLibraryByAddressImpl<CodeCache>(_libs, g_lib0_text + 10);
+
+    // Now look up an address in lib2 — cache miss, linear scan must find it.
+    const void* addr = g_lib2_text + 500;
+    CodeCache* result = findLibraryByAddressImpl<CodeCache>(_libs, addr);
+    EXPECT_EQ(result, _lib2);
+}