Protect patching with an open handle

src/codeCache.cpp

@@ -29,12 +29,14 @@ 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);
     _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;
@@ -311,3 +313,16 @@ size_t CodeCache::usedMemory() {
     }
     return bytes;
 }
+
+bool CodeCache::isValidHandle(void* handle) const {
+#ifdef __linux__
+    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 _image_base == (const char*)map->l_addr;
+    }
+    return false;
+#else
+    return handle != NULL;
+#endif
+}

src/codeCache.h

@@ -7,6 +7,11 @@
 #define _CODECACHE_H
 
 #include <jvmti.h>
+#include <dlfcn.h>
+
+#ifdef __linux__
+#include <link.h>
+#endif
 
 
 #define NO_MIN_ADDRESS  ((const void*)-1)
@@ -107,6 +112,8 @@ class CodeCache {
     const void* _min_address;
     const void* _max_address;
     const char* _text_base;
+    // TODO: Cleanup usages of lib.image_base
+    const char* _image_base;
 
     unsigned int _plt_offset;
     unsigned int _plt_size;
@@ -131,7 +138,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 +155,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;
     }
@@ -201,6 +213,8 @@ class CodeCache {
     void setDwarfTable(FrameDesc* table, int length);
     FrameDesc* findFrameDesc(const void* pc);
 
+    bool isValidHandle(void* handle) const;
+
     size_t usedMemory();
 };
 

src/hooks.cpp

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

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>
 
@@ -153,6 +154,10 @@ void MallocTracer::patchLibraries() {
 
     while (_patched_libs < native_lib_count) {
         CodeCache* cc = (*native_libs)[_patched_libs++];
+        void* handle;
+        if (!Symbols::isSafeToPatch(cc, &handle)) {
+            continue;
+        }
 
         cc->patchImport(im_malloc, (void*)malloc_hook);
         cc->patchImport(im_realloc, (void*)realloc_hook);
@@ -169,6 +174,10 @@ void MallocTracer::patchLibraries() {
             cc->patchImport(im_calloc, (void*)calloc_hook);
             cc->patchImport(im_posix_memalign, (void*)posix_memalign_hook);
         }
+
+        if (handle != NULL) {
+            dlclose(handle);
+        }
     }
 }
 

src/symbols.h

@@ -8,13 +8,24 @@
 
 #include "codeCache.h"
 #include "mutex.h"
+#include <dlfcn.h>
 
 
 class Symbols {
   private:
     static Mutex _parse_lock;
     static bool _have_kernel_symbols;
     static bool _libs_limit_reported;
+    static const void* _main_phdr;
+    static const char* _ld_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;
+    }
 
   public:
     static void parseKernelSymbols(CodeCache* cc);
@@ -23,6 +34,32 @@ class Symbols {
     static bool haveKernelSymbols() {
         return _have_kernel_symbols;
     }
+
+    static bool isSafeToPatch(CodeCache* cc, void **handle_ptr) {
+      *handle_ptr = NULL;
+
+      if (isMainExecutable(cc->imageBase(), cc->maxAddress()) || isLoader(cc->imageBase())) {
+        return true;
+      }
+
+      // Protect library from unloading while parsing in-memory ELF program headers.
+      // Also, dlopen() ensures the library is fully loaded.
+      *handle_ptr = dlopen(cc->name(), RTLD_LAZY | RTLD_NOLOAD);
+      if (cc->isValidHandle(*handle_ptr)) {
+        // Up to the user to dlclose the handle
+        return true;
+      }
+
+      if (*handle_ptr != NULL) {
+        dlclose(*handle_ptr);
+        *handle_ptr = NULL;
+      }
+      return false;
+    }
+
+    Symbols();
 };
 
+static Symbols symbols;
+
 #endif // _SYMBOLS_H

src/symbols_linux.cpp

@@ -13,7 +13,6 @@
 #include <sys/types.h>
 #include <sys/stat.h>
 #include <sys/mman.h>
-#include <dlfcn.h>
 #include <elf.h>
 #include <errno.h>
 #include <unistd.h>
@@ -64,6 +63,18 @@ static void applyPatch(CodeCache* cc) {}
 
 #endif
 
+Symbols::Symbols() {
+    _ld_base = (const char*)getauxval(AT_BASE);
+    if (!_ld_base) {
+        Log::warn("Cannot determine base address of the loader");
+    }
+
+    _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);
+}
 
 class SymbolDesc {
   private:
@@ -653,6 +664,8 @@ void ElfParser::addRelocationSymbols(ElfSection* reltab, const char* plt) {
 Mutex Symbols::_parse_lock;
 bool Symbols::_have_kernel_symbols = false;
 bool Symbols::_libs_limit_reported = false;
+const void* Symbols::_main_phdr = NULL;
+const char* Symbols::_ld_base = NULL;
 static std::unordered_set<u64> _parsed_inodes;
 
 void Symbols::parseKernelSymbols(CodeCache* cc) {
@@ -776,23 +789,12 @@ 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);
+        CodeCache* cc = new CodeCache(lib.file, array->count(), false, lib.map_start, lib.map_end, lib.image_base);
 
         // Strip " (deleted)" suffix so that removed library can be reopened
         size_t len = strlen(lib.file);
@@ -812,18 +814,9 @@ 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);
-
+            void* handle;
             // 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) {
+            if (Symbols::isSafeToPatch(cc, &handle)) {
                 ElfParser::parseProgramHeaders(cc, lib.image_base, lib.map_end, OS::isMusl());
             }
 

src/symbols_macos.cpp

@@ -6,14 +6,17 @@
 #ifdef __APPLE__
 
 #include <unordered_set>
-#include <dlfcn.h>
 #include <string.h>
 #include <mach-o/dyld.h>
 #include <mach-o/loader.h>
 #include <mach-o/nlist.h>
 #include "symbols.h"
 #include "log.h"
 
+Symbols::Symbols() {
+    _ld_base = NULL;
+    _main_phdr = NULL;
+}
 
 class MachOParser {
   private:
@@ -127,6 +130,8 @@ class MachOParser {
 Mutex Symbols::_parse_lock;
 bool Symbols::_have_kernel_symbols = false;
 bool Symbols::_libs_limit_reported = false;
+const void* Symbols::_main_phdr = NULL;
+const char* Symbols::_ld_base = NULL;
 static std::unordered_set<const void*> _parsed_libraries;
 
 void Symbols::parseKernelSymbols(CodeCache* cc) {