Protect patching with an open handle

.github/workflows/test-and-publish-nightly.yml

@@ -99,7 +99,7 @@ jobs:
             name: ubuntu-latest
         java-version: [11, 17, 21, 24]
         java-distribution: [corretto]
-        image: [public.ecr.aws/async-profiler/asprof-builder-x86:latest]
+        image: ["public.ecr.aws/async-profiler/asprof-builder-x86:latest"]
         include:
           - runson:
               display: macos-arm64

Makefile

@@ -233,6 +233,7 @@ build-test-bins:
 	$(CC) -o $(TEST_BIN_DIR)/native_api -Isrc test/test/c/native_api.c -ldl
 	$(CC) -o $(TEST_BIN_DIR)/profile_with_dlopen -Isrc test/test/nativemem/profile_with_dlopen.c -ldl
 	$(CC) -o $(TEST_BIN_DIR)/preload_malloc -Isrc test/test/nativemem/preload_malloc.c -ldl
+	$(CC) -o $(TEST_BIN_DIR)/nativemem_known_lib_crash -Isrc test/test/nativemem/nativemem_known_lib_crash.c -ldl
 	$(CXX) -o $(TEST_BIN_DIR)/non_java_app $(INCLUDES) $(CPP_TEST_INCLUDES) test/test/nonjava/non_java_app.cpp $(LIBS)
 
 test-cpp: build-test-cpp

src/codeCache.cpp

@@ -29,12 +29,18 @@ size_t NativeFunc::usedMemory(const char* name) {
 
 
 CodeCache::CodeCache(const char* name, short lib_index, bool imports_patchable,
-                     const void* min_address, const void* max_address) {
+                     const void* min_address, const void* max_address,
+                     const char* image_base) {
     _name = NativeFunc::create(name, -1);
+
+    // Strip " (deleted)" suffix so that removed library can be reopened
+    size_t len = strlen(name);
+
     _lib_index = lib_index;
     _min_address = min_address;
     _max_address = max_address;
     _text_base = NULL;
+    _image_base = image_base;
 
     _plt_offset = 0;
     _plt_size = 0;
@@ -234,7 +240,6 @@ void CodeCache::addImport(void** entry, const char* name) {
 void** CodeCache::findImport(ImportId id) {
     if (!_imports_patchable) {
         makeImportsPatchable();
-        _imports_patchable = true;
     }
     return _imports[id][PRIMARY];
 }
@@ -270,6 +275,8 @@ void CodeCache::makeImportsPatchable() {
         uintptr_t patch_end = (uintptr_t)max_import & ~OS::page_mask;
         mprotect((void*)patch_start, patch_end - patch_start + OS::page_size, PROT_READ | PROT_WRITE);
     }
+
+    _imports_patchable = true;
 }
 
 void CodeCache::setDwarfTable(FrameDesc* table, int length) {

src/codeCache.h

@@ -100,13 +100,18 @@ class CodeBlob {
 
 class FrameDesc;
 
+class UnloadProtection;
+
 class CodeCache {
+  friend UnloadProtection;
+
   private:
     char* _name;
     short _lib_index;
     const void* _min_address;
     const void* _max_address;
     const char* _text_base;
+    const char* _image_base;
 
     unsigned int _plt_offset;
     unsigned int _plt_size;
@@ -131,7 +136,8 @@ class CodeCache {
               short lib_index = -1,
               bool imports_patchable = false,
               const void* min_address = NO_MIN_ADDRESS,
-              const void* max_address = NO_MAX_ADDRESS);
+              const void* max_address = NO_MAX_ADDRESS,
+              const char* image_base = NULL);
 
     ~CodeCache();
 
@@ -147,6 +153,10 @@ class CodeCache {
         return _max_address;
     }
 
+    const char* imageBase() const {
+        return _image_base;
+    }
+
     bool contains(const void* address) const {
         return address >= _min_address && address < _max_address;
     }

src/hooks.cpp

@@ -13,6 +13,7 @@
 #include "cpuEngine.h"
 #include "mallocTracer.h"
 #include "profiler.h"
+#include "symbols.h"
 
 
 #define ADDRESS_OF(sym) ({ \
@@ -182,6 +183,11 @@ void Hooks::patchLibraries() {
 
     while (_patched_libs < native_lib_count) {
         CodeCache* cc = (*native_libs)[_patched_libs++];
+        UnloadProtection handle(cc);
+        if (!handle.isValid()) {
+            continue;
+        }
+
         if (!cc->contains((const void*)Hooks::init)) {
             // Let libasyncProfiler always use original dlopen
             cc->patchImport(im_dlopen, (void*)dlopen_hook);

src/mallocTracer.cpp

@@ -10,6 +10,7 @@
 #include "os.h"
 #include "profiler.h"
 #include "tsc.h"
+#include "symbols.h"
 #include <dlfcn.h>
 #include <string.h>
 
@@ -120,12 +121,16 @@ void MallocTracer::patchLibraries() {
 
     while (_patched_libs < native_lib_count) {
         CodeCache* cc = (*native_libs)[_patched_libs++];
-
         if (cc->contains((const void*)MallocTracer::initialize)) {
             // Let libasyncProfiler always use original allocation methods
             continue;
         }
 
+        UnloadProtection handle(cc);
+        if (!handle.isValid()) {
+            continue;
+        }
+
         cc->patchImport(im_malloc, (void*)malloc_hook);
         cc->patchImport(im_realloc, (void*)realloc_hook);
         cc->patchImport(im_free, (void*)free_hook);

src/symbols.h

@@ -25,4 +25,18 @@ class Symbols {
     }
 };
 
+class UnloadProtection {
+  private:
+    void* _lib_handle;
+    bool _valid;
+
+  public:
+    UnloadProtection(const CodeCache *cc);
+    ~UnloadProtection();
+
+    UnloadProtection& operator=(const UnloadProtection& other) = delete;
+
+    bool isValid() const { return _valid; }
+};
+
 #endif // _SYMBOLS_H

src/symbols_linux.cpp

@@ -13,14 +13,14 @@
 #include <sys/types.h>
 #include <sys/stat.h>
 #include <sys/mman.h>
-#include <dlfcn.h>
 #include <elf.h>
 #include <errno.h>
 #include <unistd.h>
 #include <fcntl.h>
 #include <link.h>
 #include <linux/limits.h>
 #include <sys/auxv.h>
+#include <link.h>
 #include "symbols.h"
 #include "dwarf.h"
 #include "fdtransferClient.h"
@@ -52,8 +52,11 @@ static void applyPatch(CodeCache* cc) {
     if (patch_libnet) {
         size_t len = strlen(cc->name());
         if (len >= 10 && strcmp(cc->name() + len - 10, "/libnet.so") == 0) {
-            cc->patchImport(im_poll, (void*)poll_hook);
-            patch_libnet = false;
+            UnloadProtection handle(cc);
+            if (handle.isValid()) {
+                cc->patchImport(im_poll, (void*)poll_hook);
+                patch_libnet = false;
+            }
         }
     }
 }
@@ -64,6 +67,25 @@ static void applyPatch(CodeCache* cc) {}
 
 #endif
 
+static const void* getMainPhdr() {
+    void* main_phdr = NULL;
+    dl_iterate_phdr([](struct dl_phdr_info* info, size_t size, void* data) {
+        *(const void**)data = info->dlpi_phdr;
+        return 1;
+    }, &main_phdr);
+    return main_phdr;
+}
+
+static const void* _main_phdr = getMainPhdr();
+static const char* _ld_base = (const char*)getauxval(AT_BASE);
+
+static bool isMainExecutable(const char* image_base, const void* map_end) {
+    return _main_phdr != NULL && _main_phdr >= image_base && _main_phdr < map_end;
+}
+
+static bool isLoader(const char* image_base) {
+    return _ld_base != NULL && _ld_base == image_base;
+}
 
 class SymbolDesc {
   private:
@@ -754,6 +776,14 @@ static void collectSharedLibraries(std::unordered_map<u64, SharedLibrary>& libs,
     fclose(f);
 }
 
+// Strip " (deleted)" suffix so that removed library can be reopened
+void stripDeletedSuffix(char* name) {
+    size_t len = strlen(name);
+    if (len > 10 && strcmp(name + len - 10, " (deleted)") == 0) {
+        name[len - 10] = 0;
+    }
+}
+
 void Symbols::parseLibraries(CodeCacheArray* array, bool kernel_symbols) {
     MutexLocker ml(_parse_lock);
 
@@ -776,29 +806,13 @@ void Symbols::parseLibraries(CodeCacheArray* array, bool kernel_symbols) {
     std::unordered_map<u64, SharedLibrary> libs;
     collectSharedLibraries(libs, MAX_NATIVE_LIBS - array->count());
 
-    const char* ld_base = (const char*)getauxval(AT_BASE);
-    if (!ld_base) {
-        Log::warn("Cannot determine base address of the loader");
-    }
-
-    const void* main_phdr = NULL;
-    dl_iterate_phdr([](struct dl_phdr_info* info, size_t size, void* data) {
-        *(const void**)data = info->dlpi_phdr;
-        return 1;
-    }, &main_phdr);
-
     for (auto& it : libs) {
         u64 inode = it.first;
         _parsed_inodes.insert(inode);
 
         SharedLibrary& lib = it.second;
-        CodeCache* cc = new CodeCache(lib.file, array->count(), false, lib.map_start, lib.map_end);
-
-        // Strip " (deleted)" suffix so that removed library can be reopened
-        size_t len = strlen(lib.file);
-        if (len > 10 && strcmp(lib.file + len - 10, " (deleted)") == 0) {
-            lib.file[len - 10] = 0;
-        }
+        CodeCache* cc = new CodeCache(lib.file, array->count(), false, lib.map_start, lib.map_end, lib.image_base);
+        stripDeletedSuffix(lib.file);
 
         if (strchr(lib.file, ':') != NULL) {
             // Do not try to parse pseudofiles like anon_inode:name, /memfd:name
@@ -812,24 +826,10 @@ void Symbols::parseLibraries(CodeCacheArray* array, bool kernel_symbols) {
             // Parse debug symbols first
             ElfParser::parseFile(cc, lib.image_base, lib.file, true);
 
-            dlerror();  // reset any error from previous dl function calls
-
-            // Protect library from unloading while parsing in-memory ELF program headers.
-            // Also, dlopen() ensures the library is fully loaded.
-            // Main executable and ld-linux interpreter cannot be dlopen'ed, but dlerror() returns NULL for them on some systems.
-            void* handle = dlopen(lib.file, RTLD_LAZY | RTLD_NOLOAD);
-
-            // Parse main executable and the loader (ld.so) regardless of dlopen result, since they cannot be unloaded.
-            bool is_main_exe = main_phdr >= lib.image_base && main_phdr < lib.map_end;
-            bool is_loader = ld_base == lib.image_base;
-
-            if (handle != NULL || is_main_exe || is_loader) {
+            UnloadProtection handle(cc);
+            if (handle.isValid()) {
                 ElfParser::parseProgramHeaders(cc, lib.image_base, lib.map_end, OS::isMusl());
             }
-
-            if (handle != NULL) {
-                dlclose(handle);
-            }
         }
 
         free(lib.file);
@@ -845,4 +845,39 @@ void Symbols::parseLibraries(CodeCacheArray* array, bool kernel_symbols) {
     }
 }
 
+static bool isValidHandle(const CodeCache* cc, void* handle) {
+    struct link_map* map;
+    // validate that the current loaded library is the same library that was observed during the /proc/self/maps processing
+    if (handle != NULL && dlinfo(handle, RTLD_DI_LINKMAP, &map) == 0) {
+        return cc->imageBase() == (const char*)map->l_addr;
+    }
+    return false;
+}
+
+UnloadProtection::UnloadProtection(const CodeCache *cc) {
+    _lib_handle = NULL;
+    _valid = false;
+
+    if (OS::isMusl() || isMainExecutable(cc->imageBase(), cc->maxAddress()) || isLoader(cc->imageBase())) {
+        _valid = true;
+        return;
+    }
+
+    char* nameCopy = strdup(cc->name());
+    stripDeletedSuffix(nameCopy);
+
+    // Protect library from unloading while parsing in-memory ELF program headers.
+    // Also, dlopen() ensures the library is fully loaded.
+    _lib_handle = dlopen(nameCopy, RTLD_LAZY | RTLD_NOLOAD);
+    _valid = isValidHandle(cc, _lib_handle);
+
+    free(nameCopy);
+}
+
+UnloadProtection::~UnloadProtection() {
+    if (_lib_handle != NULL) {
+        dlclose(_lib_handle);
+    }
+}
+
 #endif // __linux__

src/symbols_macos.cpp

@@ -14,6 +14,18 @@
 #include "symbols.h"
 #include "log.h"
 
+UnloadProtection::UnloadProtection(const CodeCache *cc) {
+    // Protect library from unloading while parsing in-memory ELF program headers.
+    // Also, dlopen() ensures the library is fully loaded.
+    _lib_handle = dlopen(cc->name(), RTLD_LAZY | RTLD_NOLOAD);
+    _valid = _lib_handle != NULL;
+}
+
+UnloadProtection::~UnloadProtection() {
+    if (_lib_handle != NULL) {
+        dlclose(_lib_handle);
+    }
+}
 
 class MachOParser {
   private:
@@ -152,22 +164,16 @@ void Symbols::parseLibraries(CodeCacheArray* array, bool kernel_symbols) {
         }
 
         const char* path = _dyld_get_image_name(i);
-
-        // Protect the library from unloading while parsing symbols
-        void* handle = dlopen(path, RTLD_LAZY | RTLD_NOLOAD);
-        if (handle == NULL) {
-            continue;
-        }
-
         CodeCache* cc = new CodeCache(path, count, true);
-        MachOParser parser(cc, image_base);
-        if (!parser.parse()) {
-            Log::warn("Could not parse symbols from %s", path);
+        UnloadProtection handle(cc);
+        if (handle.isValid()) {
+            MachOParser parser(cc, image_base);
+            if (!parser.parse()) {
+                Log::warn("Could not parse symbols from %s", path);
+            }
+            cc->sort();
+            array->add(cc);
         }
-        dlclose(handle);
-
-        cc->sort();
-        array->add(cc);
     }
 }
 

test/test/nativemem/NativememTests.java

@@ -199,4 +199,10 @@ public void preloadMalloc(TestProcess p) throws Exception {
 
         Assert.isEqual(sizeCounts.getOrDefault((long) MALLOC_SIZE, 0L), 1);
     }
+
+    @Test(os = Os.LINUX, sh = "%testbin/nativemem_known_lib_crash %f.jfr", output = true, env = {"LD_LIBRARY_PATH=build/test/lib:build/lib"})
+    public void nativememKnownLibCrash(TestProcess p) throws Exception {
+        p.waitForExit();
+        Assert.isEqual(p.exitCode(), 0);
+    }
 }