Reset NetHack by overriding rw segments of dynamic library in memory.

Based on investigation in
  https://gist.github.com/heiner/bc78064fec32174e1a216dbd5fbc6503

Fixes #254.

CMakeLists.txt

@@ -119,17 +119,11 @@ target_link_directories(nethack PUBLIC /usr/local/lib)
 
 target_link_libraries(nethack PUBLIC m fcontext bz2)
 
-# dlopen wrapper library
-add_library(nethackdl STATIC "sys/unix/nledl.c")
-target_include_directories(nethackdl PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
-target_link_libraries(nethackdl PUBLIC dl)
-
-# rlmain C++ (test) binary
-add_executable(rlmain "sys/unix/rlmain.cc")
-set_target_properties(rlmain PROPERTIES CXX_STANDARD 11)
-target_link_libraries(rlmain PUBLIC nethackdl)
-target_include_directories(rlmain PUBLIC ${NLE_INC_GEN})
-add_dependencies(rlmain util) # For pm.h.
+# rlmain C++ (test) binary add_executable(rlmain "sys/unix/rlmain.cc")
+# set_target_properties(rlmain PROPERTIES CXX_STANDARD 11)
+# target_link_libraries(rlmain PUBLIC nethackdl)
+# target_include_directories(rlmain PUBLIC ${NLE_INC_GEN})
+# add_dependencies(rlmain util) # For pm.h.
 
 # pybind11 python library.
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/pybind11)
@@ -141,7 +135,8 @@ pybind11_add_module(
   src/drawing.c
   src/objects.c
   $<TARGET_OBJECTS:tile>)
-target_link_libraries(_pynethack PUBLIC nethackdl)
+target_link_libraries(_pynethack PUBLIC dl)
 set_target_properties(_pynethack PROPERTIES CXX_STANDARD 14)
-target_include_directories(_pynethack PUBLIC ${NLE_INC_GEN})
+target_include_directories(
+  _pynethack PUBLIC ${NLE_INC_GEN} ${CMAKE_CURRENT_SOURCE_DIR}/include)
 add_dependencies(_pynethack util) # For pm.h.

include/dloverride.h

@@ -0,0 +1,234 @@
+/* Copyright (c) Facebook, Inc. and its affiliates. */
+
+/* Mechanism to reset a loaded dynamic library. */
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <assert.h>
+#include <stdio.h>
+
+#include <dlfcn.h>
+#include <errno.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+
+#ifdef __linux__
+#include <elf.h>
+
+#define PAGE_SIZE 4096
+#define PAGE_MASK (PAGE_SIZE - 1)
+#define PAGE_START(x) ((x) & ~PAGE_MASK)
+#define PAGE_OFFSET(x) ((x) &PAGE_MASK)
+#define PAGE_END(x) PAGE_START((x) + (PAGE_SIZE - 1))
+
+#if __LP64__
+#define Elf_Ehdr Elf64_Ehdr
+#define Elf_Phdr Elf64_Phdr
+
+#else /* __LP64__ */
+#define Elf_Ehdr Elf32_Ehdr
+#define Elf_Phdr Elf32_Phdr
+#endif /* __LP64__ */
+
+#elif __APPLE__
+#include <mach-o/dyld.h>
+
+#if __LP64__
+#define LC_SEGMENT_COMMAND LC_SEGMENT_64
+#define MH_MAGIC_NUMBER MH_MAGIC_64
+
+struct macho_header : public mach_header_64 {
+};
+struct macho_segment_command : public segment_command_64 {
+};
+struct macho_section : public section_64 {
+};
+#else /* __LP64__ */
+#define LC_SEGMENT_COMMAND LC_SEGMENT
+#define MH_MAGIC_NUMBER MH_MAGIC
+
+struct macho_header : public mach_header {
+};
+struct macho_segment_command : public segment_command {
+};
+struct macho_section : public section {
+};
+#endif /* __LP64__ */
+
+#endif /* __linux__, __APPLE__ */
+
+class DL
+{
+  public:
+    DL(const char *filename, const char *symbol)
+        : handle_(dlopen(filename, RTLD_LAZY))
+    {
+        if (!handle_) {
+            throw std::runtime_error(std::string("dlopen failed on ")
+                                     + filename + ": " + dlerror());
+        }
+
+        void *ptr = dlsym(handle_, symbol);
+        if (!ptr) {
+            throw std::runtime_error(dlerror());
+        }
+        Dl_info dlinfo;
+        if (dladdr(ptr, &dlinfo) == 0) {
+            throw std::runtime_error("dladdr failed");
+        }
+        if (!dlinfo.dli_sname) {
+            throw std::runtime_error("No matching addr found.");
+        }
+#ifdef __linux__
+        hdr_ = (Elf_Ehdr *) dlinfo.dli_fbase;
+        if (memcmp(hdr_->e_ident, ELFMAG, SELFMAG) != 0) {
+            throw std::runtime_error("Illegal elf header");
+        }
+
+        size_t offset = ~0;
+
+        size_t phoff = hdr_->e_phoff;
+        for (size_t i = 0; i != hdr_->e_phnum; ++i) {
+            Elf_Phdr *ph = (Elf_Phdr *) ((uint8_t *) hdr_ + phoff);
+
+            if (ph->p_type == PT_LOAD) {
+                offset = std::min(offset, (size_t) ph->p_vaddr);
+            }
+            segs_.push_back(ph);
+
+            phoff += hdr_->e_phentsize;
+        }
+
+        baseaddr_ = (uint8_t *) hdr_ - PAGE_START(offset);
+
+#elif __APPLE__
+        hdr_ = (macho_header *) dlinfo.dli_fbase;
+        if (hdr_->magic != MH_MAGIC_NUMBER) {
+            throw std::runtime_error(
+                "Illegal magic integer " + std::to_string(hdr_->magic)
+                + ", expected " + std::to_string(MH_MAGIC_NUMBER));
+        }
+        if (hdr_->filetype != MH_DYLIB) {
+            throw std::runtime_error(
+                std::string("Expected MH_DYLIB file type but got "
+                            + std::to_string(hdr_->filetype)));
+        }
+
+        const load_command *cmds = (load_command *) (hdr_ + 1);
+        const load_command *cmd = cmds;
+
+        for (uint32_t i = 0; i < hdr_->ncmds; ++i) {
+            if (cmd->cmd != LC_SEGMENT_COMMAND)
+                continue;
+
+            const auto *seg = (macho_segment_command *) cmd;
+            if (seg->nsects)
+                segs_.push_back(seg);
+            cmd = (const load_command *) (((uint8_t *) cmd) + cmd->cmdsize);
+        }
+
+        baseaddr_ = (uint8_t *) hdr_ - segs_[0]->vmaddr;
+#endif /* __linux__, __APPLE__ */
+    }
+
+    ~DL()
+    {
+        if (handle_)
+            dlclose(handle_);
+    }
+
+    DL(DL &&dl) noexcept
+    {
+        *this = std::move(dl);
+    }
+    DL(const DL &) = delete;
+    DL &operator=(const DL &) = delete;
+    DL &
+    operator=(DL &&dl) noexcept
+    {
+        if (this == &dl)
+            return *this;
+        if (handle_)
+            dlclose(handle_);
+        handle_ = std::exchange(dl.handle_, nullptr);
+        segs_ = std::move(dl.segs_);
+        hdr_ = dl.hdr_;
+        baseaddr_ = dl.baseaddr_;
+        return *this;
+    }
+
+#ifdef __linux__
+    bool
+    is_overridable(const Elf_Phdr *ph) const
+    {
+        return ph->p_type == PT_LOAD && ph->p_flags & PF_R
+               && ph->p_flags & PF_W;
+    }
+#elif __APPLE__
+    bool
+    is_overridable(const macho_segment_command *seg) const
+    {
+        return strcmp(seg->segname, SEG_DATA) == 0;
+    }
+#endif
+
+    template <typename F>
+    void
+    for_changing_sections(F &&f)
+    {
+        for (const auto *seg : segs_) {
+            if (is_overridable(seg))
+                f(seg);
+        }
+    }
+
+    template <typename T, typename... Ts>
+    auto
+    func(const char *symbol) -> decltype((T(*)(Ts...)) nullptr)
+    {
+        void *ptr = dlsym(handle_, symbol);
+        if (!ptr) {
+            throw std::runtime_error(dlerror());
+        }
+        return (T(*)(Ts...)) ptr;
+    }
+
+#ifdef __linux__
+    uint8_t *
+    mem_addr(const Elf_Phdr *ph) const
+    {
+        size_t start = PAGE_END((size_t) (baseaddr_ + ph->p_vaddr));
+        return (uint8_t *) start;
+    }
+    size_t
+    mem_size(const Elf_Phdr *ph) const
+    {
+        return ph->p_memsz;
+    }
+#elif __APPLE__
+    uint8_t *
+    mem_addr(const macho_segment_command *seg) const
+    {
+        return baseaddr_ + seg->vmaddr;
+    }
+    size_t
+    mem_size(const macho_segment_command *seg) const
+    {
+        return seg->vmsize;
+    }
+#endif
+
+  private:
+    void *handle_{ nullptr };
+#ifdef __linux__
+    std::vector<const Elf_Phdr *> segs_;
+    const Elf_Ehdr *hdr_;
+#elif __APPLE__
+    std::vector<const macho_segment_command *> segs_;
+    const macho_header *hdr_;
+#endif
+    uint8_t *baseaddr_;
+};

setup.py

@@ -42,7 +42,7 @@ def build_extension(self, ext):
         hackdir_path = os.getenv("HACKDIR", output_path.joinpath("nethackdir"))
 
         os.makedirs(self.build_temp, exist_ok=True)
-        build_type = "Debug" if self.debug else "Release"
+        build_type = "Debug"  # if self.debug else "Release"
 
         generator = "Ninja" if spawn.find_executable("ninja") else "Unix Makefiles"
 

src/nle.c

@@ -419,12 +419,12 @@ nle_start(nle_obs *obs, FILE *ttyrec, nle_seeds_init_t *seed_init,
 
     nle_ctx_t *nle = init_nle(ttyrec, obs);
     nle_seeds_init = seed_init;
+    current_nle_ctx = nle;
 
     nle->stack = create_fcontext_stack(STACK_SIZE);
     nle->generatorcontext =
         make_fcontext(nle->stack.sptr, nle->stack.ssize, mainloop);
 
-    current_nle_ctx = nle;
     fcontext_transfer_t t = jump_fcontext(nle->generatorcontext, NULL);
     nle->generatorcontext = t.ctx;
     nle->done = (t.data == NULL);