Override dlsym instead of dlopen to correctly honour RPATH/RUNPATHS

We need to override dlopen/dlclose to account for new shared libraries
being loaded in the process memory space. This is needed so we can
correctly track allocations in those libraries by overriding their PLT
entries and also so we can properly map the addresses of the symbols in
those libraries when we resolve later native traces. Unfortunately, we
can't just override dlopen directly because of the following edge case:
when a shared library dlopen's another by name (e.g.
dlopen("libfoo.so")), the dlopen call will honor the RPATH/RUNPATH of
the calling library if it's set. Some libraries set an RPATH/RUNPATH
based on $ORIGIN (the path of the calling library) to load dependencies
from a relative directory based on the location of the calling library.
This means that if we override dlopen, we'll end up loading the library
from the wrong path or more likely, not loading it at all because the
dynamic loader will think the memray extenion it's the calling library
and the RPATH of the real calling library will not be honoured.

To work around this, we override dlsym instead and override the symbols
in the loaded libraries only the first time we have seen a handle passed
to dlsym. This works because for a symbol from a given dlopen-ed library
to appear in a call stack, *something* from that library has to be
dlsym-ed first. The only exception to this are static initializers, but
we cannot track those anyway by overriding dlopen as they run within the
dlopen call itself.

Signed-off-by: Pablo Galindo <[email protected]>

news/525.bugfix.rst

@@ -0,0 +1 @@
+Fix a bug that was causing ``dlopen`` to not load shared libraries that have an RPATH/RUNPATH set.

src/memray/_memray/hooks.cpp

@@ -1,5 +1,7 @@
 #include <cassert>
 #include <cstdio>
+#include <shared_mutex>
+#include <unordered_set>
 
 #include "hooks.h"
 #include "tracking_api.h"
@@ -304,19 +306,76 @@ posix_memalign(void** memptr, size_t alignment, size_t size) noexcept
     return ret;
 }
 
+// We need to override dlopen/dlclose to account for new shared libraries being
+// loaded in the process memory space. This is needed so we can correctly track
+// allocations in those libraries by overriding their PLT entries and also so we
+// can properly map the addresses of the symbols in those libraries when we
+// resolve later native traces. Unfortunately, we can't just override dlopen
+// directly because of the following edge case: when a shared library dlopen's
+// another by name (e.g. dlopen("libfoo.so")), the dlopen call will honor the
+// RPATH/RUNPATH of the calling library if it's set. Some libraries set an
+// RPATH/RUNPATH based on $ORIGIN (the path of the calling library) to load
+// dependencies from a relative directory based on the location of the calling
+// library. This means that if we override dlopen, we'll end up loading the
+// library from the wrong path or more likely, not loading it at all because the
+// dynamic loader will think the memray extenion it's the calling library and
+// the RPATH of the real calling library will not be honoured.
+//
+// To work around this, we override dlsym instead and override the symbols in
+// the loaded libraries only the first time we have seen a handle passed to
+// dlsym. This works because for a symbol from a given dlopen-ed library to
+// appear in a call stack, *something* from that library has to be dlsym-ed
+// first. The only exception to this are static initializers, but we cannot
+// track those anyway by overriding dlopen as they run within the dlopen call
+// itself.
+
+class DlsymCache
+{
+  public:
+    void addEntry(const void* handle)
+    {
+        std::unique_lock lock(mutex_);
+        d_handles.insert(handle);
+    }
+
+    bool lookup(const void* handle) const
+    {
+        std::shared_lock lock(mutex_);
+        auto it = d_handles.find(handle);
+        if (it != d_handles.end()) {
+            return true;
+        }
+        return false;
+    }
+
+    void deleteEntry(const void* handle)
+    {
+        std::unique_lock lock(mutex_);
+        d_handles.erase(handle);
+    }
+
+  private:
+    mutable std::shared_mutex mutex_;
+    std::set<const void*> d_handles;
+};
+
+static DlsymCache dlsym_cache;
+
 void*
-dlopen(const char* filename, int flag) noexcept
+dlsym(void* handle, const char* symbol) noexcept
 {
-    assert(hooks::dlopen);
+    assert(hooks::dlsym);
     void* ret;
     {
         tracking_api::RecursionGuard guard;
-        ret = hooks::dlopen(filename, flag);
+        ret = hooks::dlsym(handle, symbol);
     }
-    if (ret) {
+    if (ret && !dlsym_cache.lookup(handle)) {
+        dlsym_cache.addEntry(handle);
         tracking_api::Tracker::invalidate_module_cache();
-        if (filename
-            && (nullptr != strstr(filename, "/_greenlet.") || nullptr != strstr(filename, "/greenlet.")))
+        if (symbol
+            && (nullptr != strstr(symbol, "PyInit_greenlet")
+                || nullptr != strstr(symbol, "PyInit__greenlet")))
         {
             tracking_api::Tracker::beginTrackingGreenlets();
         }
@@ -334,6 +393,7 @@ dlclose(void* handle) noexcept
         tracking_api::RecursionGuard guard;
         ret = hooks::dlclose(handle);
     }
+    dlsym_cache.deleteEntry(handle);
     tracking_api::NativeTrace::flushCache();
     if (!ret) tracking_api::Tracker::invalidate_module_cache();
     return ret;

src/memray/_memray/hooks.h

@@ -43,7 +43,7 @@
     FOR_EACH_HOOKED_FUNCTION(aligned_alloc)                                                             \
     FOR_EACH_HOOKED_FUNCTION(mmap)                                                                      \
     FOR_EACH_HOOKED_FUNCTION(munmap)                                                                    \
-    FOR_EACH_HOOKED_FUNCTION(dlopen)                                                                    \
+    FOR_EACH_HOOKED_FUNCTION(dlsym)                                                                     \
     FOR_EACH_HOOKED_FUNCTION(dlclose)                                                                   \
     FOR_EACH_HOOKED_FUNCTION(PyGILState_Ensure)                                                         \
     MEMRAY_PLATFORM_HOOKED_FUNCTIONS
@@ -175,7 +175,7 @@ void*
 pvalloc(size_t size) noexcept;
 
 void*
-dlopen(const char* filename, int flag) noexcept;
+dlsym(void* handle, const char* symbol) noexcept;
 
 int
 dlclose(void* handle) noexcept;

src/memray/commands/_attach.lldb

@@ -15,7 +15,7 @@ breakpoint set -b malloc -b calloc -b realloc -b free -b PyMem_Malloc -b PyMem_C
 # Set commands to execute when breakpoint is reached
 breakpoint command add -e true
 breakpoint disable
-expr auto $dlsym = (void* (*)(void*, const char*))&dlsym
+expr auto $dlsym = (void* (*)(void*, const char*))dlsym.d_original
 expr auto $dlopen = $dlsym($rtld_default, "dlopen")
 expr auto $dlerror = $dlsym($rtld_default, "dlerror")
 expr auto $dll = ((void*(*)(const char*, int))$dlopen)($libpath, $rtld_now)

tests/integration/rpath_extension/ext.c

@@ -0,0 +1,46 @@
+#include <Python.h>
+#include <dlfcn.h>
+
+
+static PyObject *hello_world(PyObject *self, PyObject *args) {
+    // Load the shared library
+    void *lib_handle = dlopen("sharedlib.so", RTLD_LAZY);
+
+    if (!lib_handle) {
+        PyErr_SetString(PyExc_RuntimeError, dlerror());
+        return NULL;
+    }
+
+    // Get the function pointer
+    void (*my_shared_function)() = dlsym(lib_handle, "my_shared_function");
+    if (!my_shared_function) {
+        PyErr_SetString(PyExc_RuntimeError, dlerror());
+        dlclose(lib_handle);
+        return NULL;
+    }
+
+    // Call the function
+    my_shared_function();
+
+    // Close the shared library
+    dlclose(lib_handle);
+
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef methods[] = {
+    {"hello_world", hello_world, METH_NOARGS, "Print Hello, World!"},
+    {NULL, NULL, 0, NULL}
+};
+
+static struct PyModuleDef module = {
+    PyModuleDef_HEAD_INIT,
+    "ext",
+    NULL,
+    -1,
+    methods
+};
+
+PyMODINIT_FUNC PyInit_ext(void) {
+    return PyModule_Create(&module);
+}

tests/integration/rpath_extension/setup.py

@@ -0,0 +1,32 @@
+import subprocess
+
+from setuptools import Extension
+from setuptools import setup
+from setuptools.command.build_ext import build_ext
+
+# Compile the shared library before building the extension
+subprocess.run(
+    ["gcc", "-shared", "-o", "sharedlibs/sharedlib.so", "sharedlibs/sharedlib.c"]
+)
+
+
+class CustomBuildExt(build_ext):
+    def run(self):
+        # Run the original build_ext command
+        build_ext.run(self)
+
+        # Add code to compile the shared library
+
+
+setup(
+    name="ext",
+    version="1.0",
+    ext_modules=[
+        Extension(
+            "ext",
+            sources=["ext.c"],
+            extra_link_args=["-Wl,-rpath,$ORIGIN/sharedlibs"],
+        )
+    ],
+    cmdclass={"build_ext": CustomBuildExt},
+)

tests/integration/rpath_extension/sharedlibs/sharedlib.c

@@ -0,0 +1,5 @@
+#include <stdio.h>
+
+void my_shared_function() {
+    printf("This is a function from your_shared_lib!\n");
+}

tests/integration/test_extensions.py

@@ -14,6 +14,7 @@
 HERE = Path(__file__).parent
 TEST_MULTITHREADED_EXTENSION = HERE / "multithreaded_extension"
 TEST_MISBEHAVING_EXTENSION = HERE / "misbehaving_extension"
+TEST_RPATH_EXTENSION = HERE / "rpath_extension"
 
 
 @pytest.mark.valgrind
@@ -107,7 +108,7 @@ def allocating_function():  # pragma: no cover
     func, filename, line = bottom_frame
     assert func == "allocating_function"
     assert filename.endswith(__file__)
-    assert line == 82
+    assert line == 83
 
     frees = [
         event
@@ -170,7 +171,7 @@ def foo2():
     func, filename, line = bottom_frame
     assert func == "test_extension_that_uses_pygilstate_ensure"
     assert filename.endswith(__file__)
-    assert line == 153
+    assert line == 154
 
     # We should have 2 frames here: this function calling `allocator.valloc`,
     # and `allocator.valloc` calling the C `valloc`.
@@ -185,7 +186,7 @@ def foo2():
     func, filename, line = caller
     assert func == "test_extension_that_uses_pygilstate_ensure"
     assert filename.endswith(__file__)
-    assert line == 154
+    assert line == 155
 
     frees = [
         event
@@ -241,7 +242,7 @@ def allocating_function():
     func, filename, line = bottom_frame
     assert func == "test_native_dlopen"
     assert filename.endswith(__file__)
-    assert line == 225
+    assert line == 226
 
     frees = [
         event
@@ -354,3 +355,34 @@ def test_hard_exit(tmpdir, py_finalize):
 
     # THEN
     # No assertions, just check that the program exits without error
+
+
+@pytest.mark.skipif(
+    sys.platform == "darwin", reason="Test requires a linker that supports -rpath"
+)
+def test_dlopen_with_rpath(tmpdir, monkeypatch):
+    # GIVEN
+    output = Path(tmpdir) / "test.bin"
+    extension_name = "sharedlibs"
+    extension_path = tmpdir / extension_name
+    shutil.copytree(TEST_RPATH_EXTENSION, extension_path)
+    subprocess.run(
+        [sys.executable, str(extension_path / "setup.py"), "build_ext", "--inplace"],
+        check=True,
+        cwd=extension_path,
+        capture_output=True,
+    )
+
+    # WHEN
+    with monkeypatch.context() as ctx:
+        ctx.setattr(sys, "path", [*sys.path, str(extension_path)])
+        from ext import hello_world  # type: ignore
+
+        try:
+            hello_world()
+        except RuntimeError:
+            pytest.skip("Test requires a linker that supports -rpath")
+
+        # THEN
+        with Tracker(output):
+            hello_world()