Add loadNativeObj and unloadNativeObj

(This change is originally written by niteria)

This adds two functions:
* `loadNativeObj`
* `unloadNativeObj`
and implements them for Linux.

They are useful if you want to load a shared object with Haskell code
using the system linker and have GHC call dlclose() after the
code is no longer referenced from the heap.

Using the system linker allows you to load the shared object
above outside the low-mem region. It also loads the DWARF sections
in a way that `perf` understands.

`dl_iterate_phdr` is what makes this implementation Linux specific.

5 files changed, 274 insertions, 15 deletions

diff --git a/includes/rts/Linker.h b/includes/rts/Linker.h
index 06c9402b4f..1f3719c0c7 100644
--- a/includes/rts/Linker.h
+++ b/includes/rts/Linker.h
@@ -76,6 +76,19 @@ HsInt loadArchive( pathchar *path );
 /* resolve all the currently unlinked objects in memory */
 HsInt resolveObjs( void );
 
+/* Load an .so using the system linker.
+   Returns a handle that can be passed to dlsym() or NULL on error.
+
+   In the case of error, stores the error message in errmsg. The caller
+   is responsible for freeing it. */
+void *loadNativeObj( pathchar *path, char **errmsg );
+
+/* Mark the .so loaded with the system linker for unloading.
+   The RTS will unload it when all the references to the .so disappear from
+   the heap.
+   Takes the handle returned from loadNativeObj() as an argument. */
+HsInt unloadNativeObj( void *handle );
+
 /* load a dynamic library */
 const char *addDLL( pathchar* dll_name );
 
diff --git a/rts/CheckUnload.c b/rts/CheckUnload.c
index fcbe0f6156..8f834d13fa 100644
--- a/rts/CheckUnload.c
+++ b/rts/CheckUnload.c
@@ -238,21 +238,45 @@ static void reserveOCSectionIndices(OCSectionIndices *s_indices, int len)
 // state.
 void insertOCSectionIndices(ObjectCode *oc)
 {
-    reserveOCSectionIndices(global_s_indices, oc->n_sections);
+    // after we finish the section table will no longer be sorted.
     global_s_indices->sorted = false;
 
-    int s_i = global_s_indices->n_sections;
-    for (int i = 0; i < oc->n_sections; i++) {
-        if (oc->sections[i].kind != SECTIONKIND_OTHER) {
-            global_s_indices->indices[s_i].start = (W_)oc->sections[i].start;
-            global_s_indices->indices[s_i].end = (W_)oc->sections[i].start
-                + oc->sections[i].size;
-            global_s_indices->indices[s_i].oc = oc;
+    if (oc->type == DYNAMIC_OBJECT) {
+        // First count the ranges
+        int n_ranges = 0;
+        for (NativeCodeRange *ncr = oc->nc_ranges; ncr != NULL; ncr = ncr->next) {
+            n_ranges++;
+        }
+
+        // Next reserve the appropriate number of table entries...
+        reserveOCSectionIndices(global_s_indices, n_ranges);
+
+        // Now insert the new ranges...
+        int s_i = global_s_indices->n_sections;
+        for (NativeCodeRange *ncr = oc->nc_ranges; ncr != NULL; ncr = ncr->next) {
+            OCSectionIndex *ent = &global_s_indices->indices[s_i];
+            ent->start = (W_)ncr->start;
+            ent->end = (W_)ncr->end;
+            ent->oc = oc;
             s_i++;
         }
-    }
 
-    global_s_indices->n_sections = s_i;
+        global_s_indices->n_sections = s_i;
+    } else {
+        reserveOCSectionIndices(global_s_indices, oc->n_sections);
+        int s_i = global_s_indices->n_sections;
+        for (int i = 0; i < oc->n_sections; i++) {
+            if (oc->sections[i].kind != SECTIONKIND_OTHER) {
+                OCSectionIndex *ent = &global_s_indices->indices[s_i];
+                ent->start = (W_)oc->sections[i].start;
+                ent->end = (W_)oc->sections[i].start + oc->sections[i].size;
+                ent->oc = oc;
+                s_i++;
+            }
+        }
+
+        global_s_indices->n_sections = s_i;
+    }
 
     // Add object to 'objects' list
     if (objects != NULL) {
@@ -446,6 +470,7 @@ void checkUnload()
     ObjectCode *next = NULL;
     for (ObjectCode *oc = old_objects; oc != NULL; oc = next) {
         next = oc->next;
+        ASSERT(oc->status == OBJECT_UNLOADED);
 
         removeOCSectionIndices(s_indices, oc);
 
diff --git a/rts/Linker.c b/rts/Linker.c
index f6a38e08dd..96d25fb741 100644
--- a/rts/Linker.c
+++ b/rts/Linker.c
@@ -64,6 +64,7 @@
 #  include "linker/Elf.h"
 #  include <regex.h>    // regex is already used by dlopen() so this is OK
                         // to use here without requiring an additional lib
+#  include <link.h>
 #elif defined(OBJFORMAT_PEi386)
 #  include "linker/PEi386.h"
 #  include <windows.h>
@@ -170,6 +171,8 @@ Mutex linker_mutex;
 /* Generic wrapper function to try and Resolve and RunInit oc files */
 int ocTryLoad( ObjectCode* oc );
 
+static void freeNativeCode_ELF (ObjectCode *nc);
+
 /* Link objects into the lower 2Gb on x86_64 and AArch64.  GHC assumes the
  * small memory model on this architecture (see gcc docs,
  * -mcmodel=small).
@@ -1246,6 +1249,16 @@ freePreloadObjectFile (ObjectCode *oc)
  */
 void freeObjectCode (ObjectCode *oc)
 {
+    if (oc->type == DYNAMIC_OBJECT) {
+#if defined(OBJFORMAT_ELF)
+        ACQUIRE_LOCK(&dl_mutex);
+        freeNativeCode_ELF(oc);
+        RELEASE_LOCK(&dl_mutex);
+#else
+        barf("freeObjectCode: This shouldn't happen");
+#endif
+    }
+
     freePreloadObjectFile(oc);
 
     if (oc->symbols != NULL) {
@@ -1328,7 +1341,7 @@ void freeObjectCode (ObjectCode *oc)
 }
 
 ObjectCode*
-mkOc( pathchar *path, char *image, int imageSize,
+mkOc( ObjectType type, pathchar *path, char *image, int imageSize,
       bool mapped, pathchar *archiveMemberName, int misalignment ) {
    ObjectCode* oc;
 
@@ -1336,6 +1349,7 @@ mkOc( pathchar *path, char *image, int imageSize,
    oc = stgMallocBytes(sizeof(ObjectCode), "mkOc(oc)");
 
    oc->info = NULL;
+   oc->type = type;
 
 #  if defined(OBJFORMAT_ELF)
    oc->formatName = "ELF";
@@ -1396,6 +1410,10 @@ mkOc( pathchar *path, char *image, int imageSize,
    oc->rx_m32 = m32_allocator_new(true);
 #endif
 
+   oc->l_addr = NULL;
+   oc->nc_ranges = NULL;
+   oc->dlopen_handle = NULL;
+
    IF_DEBUG(linker, debugBelch("mkOc: done\n"));
    return oc;
 }
@@ -1524,7 +1542,7 @@ preloadObjectFile (pathchar *path)
    IF_DEBUG(linker, debugBelch("loadObj: preloaded image at %p\n", (void *) image));
 
    /* FIXME (AP): =mapped= parameter unconditionally set to true */
-   oc = mkOc(path, image, fileSize, true, NULL, misalignment);
+   oc = mkOc(STATIC_OBJECT, path, image, fileSize, true, NULL, misalignment);
 
 #if defined(OBJFORMAT_MACHO)
    if (ocVerifyImage_MachO( oc ))
@@ -1943,6 +1961,180 @@ addSection (Section *s, SectionKind kind, SectionAlloc alloc,
                        size, kind ));
 }
 
+
+#  if defined(OBJFORMAT_ELF)
+static int loadNativeObjCb_(struct dl_phdr_info *info,
+    size_t _size GNUC3_ATTRIBUTE(__unused__), void *data) {
+  ObjectCode* nc = (ObjectCode*) data;
+
+  // This logic mimicks _dl_addr_inside_object from glibc
+  // For reference:
+  // int
+  // internal_function
+  // _dl_addr_inside_object (struct link_map *l, const ElfW(Addr) addr)
+  // {
+  //   int n = l->l_phnum;
+  //   const ElfW(Addr) reladdr = addr - l->l_addr;
+  //
+  //   while (--n >= 0)
+  //     if (l->l_phdr[n].p_type == PT_LOAD
+  //         && reladdr - l->l_phdr[n].p_vaddr >= 0
+  //         && reladdr - l->l_phdr[n].p_vaddr < l->l_phdr[n].p_memsz)
+  //       return 1;
+  //   return 0;
+  // }
+
+  if ((void*) info->dlpi_addr == nc->l_addr) {
+    int n = info->dlpi_phnum;
+    while (--n >= 0) {
+      if (info->dlpi_phdr[n].p_type == PT_LOAD) {
+        NativeCodeRange* ncr =
+          stgMallocBytes(sizeof(NativeCodeRange), "loadNativeObjCb_");
+        ncr->start = (void*) ((char*) nc->l_addr + info->dlpi_phdr[n].p_vaddr);
+        ncr->end = (void*) ((char*) ncr->start + info->dlpi_phdr[n].p_memsz);
+
+        ncr->next = nc->nc_ranges;
+        nc->nc_ranges = ncr;
+      }
+    }
+  }
+  return 0;
+}
+
+static void copyErrmsg(char** errmsg_dest, char* errmsg) {
+  if (errmsg == NULL) errmsg = "loadNativeObj_ELF: unknown error";
+  *errmsg_dest = stgMallocBytes(strlen(errmsg)+1, "loadNativeObj_ELF");
+  strcpy(*errmsg_dest, errmsg);
+}
+
+// need dl_mutex
+static void freeNativeCode_ELF (ObjectCode *nc) {
+  dlclose(nc->dlopen_handle);
+
+  NativeCodeRange *ncr = nc->nc_ranges;
+  while (ncr) {
+    NativeCodeRange* last_ncr = ncr;
+    ncr = ncr->next;
+    stgFree(last_ncr);
+  }
+}
+
+static void * loadNativeObj_ELF (pathchar *path, char **errmsg)
+{
+   ObjectCode* nc;
+   void *hdl, *retval;
+
+   IF_DEBUG(linker, debugBelch("loadNativeObj_ELF %" PATH_FMT "\n", path));
+
+   retval = NULL;
+   ACQUIRE_LOCK(&dl_mutex);
+
+   nc = mkOc(DYNAMIC_OBJECT, path, NULL, 0, true, NULL, 0);
+
+   foreignExportsLoadingObject(nc);
+   hdl = dlopen(path, RTLD_NOW|RTLD_LOCAL);
+   foreignExportsFinishedLoadingObject();
+   if (hdl == NULL) {
+     /* dlopen failed; save the message in errmsg */
+     copyErrmsg(errmsg, dlerror());
+     goto dlopen_fail;
+   }
+
+   struct link_map *map;
+   if (dlinfo(hdl, RTLD_DI_LINKMAP, &map) == -1) {
+     /* dlinfo failed; save the message in errmsg */
+     copyErrmsg(errmsg, dlerror());
+     goto dlinfo_fail;
+   }
+
+   nc->l_addr = (void*) map->l_addr;
+   nc->dlopen_handle = hdl;
+   hdl = NULL; // pass handle ownership to nc
+
+   dl_iterate_phdr(loadNativeObjCb_, nc);
+   if (!nc->nc_ranges) {
+     copyErrmsg(errmsg, "dl_iterate_phdr failed to find obj");
+     goto dl_iterate_phdr_fail;
+   }
+
+   insertOCSectionIndices(nc);
+
+   nc->next_loaded_object = loaded_objects;
+   loaded_objects = nc;
+
+   retval = nc->dlopen_handle;
+   goto success;
+
+dl_iterate_phdr_fail:
+   // already have dl_mutex
+   freeNativeCode_ELF(nc);
+dlinfo_fail:
+   if (hdl) dlclose(hdl);
+dlopen_fail:
+success:
+
+   RELEASE_LOCK(&dl_mutex);
+   IF_DEBUG(linker, debugBelch("loadNativeObj_ELF result=%p\n", retval));
+
+   return retval;
+}
+
+#  endif
+
+#define UNUSED(x) (void)(x)
+
+void * loadNativeObj (pathchar *path, char **errmsg)
+{
+#if defined(OBJFORMAT_ELF)
+   ACQUIRE_LOCK(&linker_mutex);
+   void *r = loadNativeObj_ELF(path, errmsg);
+   RELEASE_LOCK(&linker_mutex);
+   return r;
+#else
+   UNUSED(path);
+   UNUSED(errmsg);
+   barf("loadNativeObj: not implemented on this platform");
+#endif
+}
+
+HsInt unloadNativeObj (void *handle)
+{
+    bool unloadedAnyObj = false;
+
+    IF_DEBUG(linker, debugBelch("unloadNativeObj: %p\n", handle));
+
+    ObjectCode *prev = NULL, *next;
+    for (ObjectCode *nc = loaded_objects; nc; nc = next) {
+        next = nc->next_loaded_object; // we might move nc
+
+        if (nc->type == DYNAMIC_OBJECT && nc->dlopen_handle == handle) {
+            nc->status = OBJECT_UNLOADED;
+            n_unloaded_objects += 1;
+
+            // dynamic objects have no symbols
+            ASSERT(nc->symbols == NULL);
+            freeOcStablePtrs(nc);
+
+            // Remove object code from root set
+            if (prev == NULL) {
+              loaded_objects = nc->next_loaded_object;
+            } else {
+              prev->next_loaded_object = nc->next_loaded_object;
+            }
+            unloadedAnyObj = true;
+        } else {
+            prev = nc;
+        }
+    }
+
+    if (unloadedAnyObj) {
+        return 1;
+    } else {
+        errorBelch("unloadObjNativeObj_ELF: can't find `%p' to unload", handle);
+        return 0;
+    }
+}
+
 /* -----------------------------------------------------------------------------
  * Segment management
  */
diff --git a/rts/LinkerInternals.h b/rts/LinkerInternals.h
index 6b726f7a27..93e949f2f2 100644
--- a/rts/LinkerInternals.h
+++ b/rts/LinkerInternals.h
@@ -31,6 +31,13 @@ typedef struct _Symbol
     SymbolAddr *addr;
 } Symbol_t;
 
+typedef struct NativeCodeRange_ {
+  void *start, *end;
+
+  /* Allow a chain of these things */
+  struct NativeCodeRange_ *next;
+} NativeCodeRange;
+
 /* Indication of section kinds for loaded objects.  Needed by
    the GC for deciding whether or not a pointer on the stack
    is a code pointer.
@@ -157,6 +164,13 @@ typedef struct {
 #endif
 } SymbolExtra;
 
+typedef enum {
+    /* Objects that were loaded by this linker */
+    STATIC_OBJECT,
+
+    /* Objects that were loaded by dlopen */
+    DYNAMIC_OBJECT,
+} ObjectType;
 
 /* Top-level structure for an object module.  One of these is allocated
  * for each object file in use.
@@ -165,7 +179,8 @@ typedef struct _ObjectCode {
     OStatus    status;
     pathchar  *fileName;
     int        fileSize;     /* also mapped image size when using mmap() */
-    char*      formatName;            /* eg "ELF32", "DLL", "COFF", etc. */
+    char*      formatName;   /* e.g. "ELF32", "DLL", "COFF", etc. */
+    ObjectType type;         /* who loaded this object? */
 
     /* If this object is a member of an archive, archiveMemberName is
      * like "libarchive.a(object.o)". Otherwise it's NULL.
@@ -267,6 +282,19 @@ typedef struct _ObjectCode {
      * (read-only/executable) code. */
     m32_allocator *rw_m32, *rx_m32;
 #endif
+
+    /*
+     * The following are only valid if .type == DYNAMIC_OBJECT
+     */
+
+    /* handle returned from dlopen */
+    void *dlopen_handle;
+
+    /* base virtual address of the loaded code */
+    void *l_addr;
+
+    /* virtual memory ranges of loaded code */
+    NativeCodeRange *nc_ranges;
 } ObjectCode;
 
 #define OC_INFORMATIVE_FILENAME(OC)             \
@@ -275,6 +303,7 @@ typedef struct _ObjectCode {
       (OC)->fileName                            \
     )
 
+
 #if defined(THREADED_RTS)
 extern Mutex linker_mutex;
 #endif
@@ -360,7 +389,7 @@ resolveSymbolAddr (pathchar* buffer, int size,
 
 HsInt isAlreadyLoaded( pathchar *path );
 HsInt loadOc( ObjectCode* oc );
-ObjectCode* mkOc( pathchar *path, char *image, int imageSize,
+ObjectCode* mkOc( ObjectType type, pathchar *path, char *image, int imageSize,
                   bool mapped, pathchar *archiveMemberName,
                   int misalignment
                   );
diff --git a/rts/linker/LoadArchive.c b/rts/linker/LoadArchive.c
index 0ad3d94725..55081489f5 100644
--- a/rts/linker/LoadArchive.c
+++ b/rts/linker/LoadArchive.c
@@ -521,7 +521,7 @@ static HsInt loadArchive_ (pathchar *path)
             pathprintf(archiveMemberName, size, WSTR("%" PATH_FMT "(%.*s)"),
                        path, (int)thisFileNameSize, fileName);
 
-            ObjectCode *oc = mkOc(path, image, memberSize, false, archiveMemberName,
+            ObjectCode *oc = mkOc(STATIC_OBJECT, path, image, memberSize, false, archiveMemberName,
                                   misalignment);
 #if defined(OBJFORMAT_MACHO)
             ocInit_MachO( oc );