Make Simulate app scan for plugin libraries on launch.

PiperOrigin-RevId: 478502151
Change-Id: I9057e74af4d207cfc551e701134820100bdeee78

simulate/main.cc

@@ -13,19 +13,37 @@
 // limitations under the License.
 
 #include <chrono>
+#include <cstdint>
 #include <cstdio>
 #include <cstring>
 #include <iostream>
 #include <memory>
 #include <mutex>
+#include <new>
 #include <string>
 #include <thread>
+#include <type_traits>
+#include <vector>
 
 #include <mujoco/mujoco.h>
 #include "glfw_dispatch.h"
 #include "simulate.h"
 #include "array_safety.h"
 
+extern "C" {
+#if defined(_WIN32) || defined(__CYGWIN__)
+  #include <windows.h>
+#else
+  #if defined(__APPLE__)
+    #include <mach-o/dyld.h>
+  #endif
+  #include <dirent.h>
+  #include <dlfcn.h>
+  #include <sys/errno.h>
+  #include <unistd.h>
+#endif
+}
+
 namespace {
 namespace mj = ::mujoco;
 namespace mju = ::mujoco::sample_util;
@@ -46,6 +64,213 @@ mjtNum* ctrlnoise = nullptr;
 
 
 
+//---------------------------------------- plugin handling -----------------------------------------
+
+// return the path to the directory containing the current executable
+// used to determine the location of auto-loaded plugin libraries
+std::string getExecutableDir() {
+#if defined(_WIN32) || defined(__CYGWIN__)
+  constexpr char kPathSep = '\\';
+  std::string realpath = [&]() -> std::string {
+    std::unique_ptr<char[]> realpath(nullptr);
+    DWORD buf_size = 128;
+    bool success = false;
+    while (!success) {
+      realpath.reset(new(std::nothrow) char[buf_size]);
+      if (!realpath) {
+        std::cerr << "cannot allocate memory to store executable path\n";
+        return "";
+      }
+
+      DWORD written = GetModuleFileNameA(nullptr, realpath.get(), buf_size);
+      if (written < buf_size) {
+        success = true;
+      } else if (written == buf_size) {
+        // realpath is too small, grow and retry
+        buf_size *=2;
+      } else {
+        std::cerr << "failed to retrieve executable path: " << GetLastError() << "\n";
+        return "";
+      }
+    }
+    return realpath.get();
+  }();
+#else
+  constexpr char kPathSep = '/';
+#if defined(__APPLE__)
+  std::unique_ptr<char[]> buf(nullptr);
+  {
+    std::uint32_t buf_size = 0;
+    _NSGetExecutablePath(nullptr, &buf_size);
+    buf.reset(new char[buf_size]);
+    if (!buf) {
+      std::cerr << "cannot allocate memory to store executable path\n";
+      return "";
+    }
+    if (_NSGetExecutablePath(buf.get(), &buf_size)) {
+      std::cerr << "unexpected error from _NSGetExecutablePath\n";
+    }
+  }
+  const char* path = buf.get();
+#else
+  const char* path = "/proc/self/exe";
+#endif
+  std::string realpath = [&]() -> std::string {
+    std::unique_ptr<char[]> realpath(nullptr);
+    std::uint32_t buf_size = 128;
+    bool success = false;
+    while (!success) {
+      realpath.reset(new(std::nothrow) char[buf_size]);
+      if (!realpath) {
+        std::cerr << "cannot allocate memory to store executable path\n";
+        return "";
+      }
+
+      std::size_t written = readlink(path, realpath.get(), buf_size);
+      if (written < buf_size) {
+        realpath.get()[written] = '\0';
+        success = true;
+      } else if (written == -1) {
+        if (errno == EINVAL) {
+          // path is already not a symlink, just use it
+          return path;
+        }
+
+        std::cerr << "error while resolving executable path: " << strerror(errno) << '\n';
+        return "";
+      } else {
+        // realpath is too small, grow and retry
+        buf_size *= 2;
+      }
+    }
+    return realpath.get();
+  }();
+#endif
+
+  if (realpath.empty()) {
+    return "";
+  }
+
+  for (std::size_t i = realpath.size() - 1; i > 0; --i) {
+    if (realpath.c_str()[i] == kPathSep) {
+      return realpath.substr(0, i);
+    }
+  }
+
+  // don't scan through the entire file system's root
+  return "";
+}
+
+
+#if defined(_WIN32) || defined(__CYGWIN__)
+using unique_dlhandle = std::unique_ptr<std::remove_pointer_t<HMODULE>, decltype(&FreeLibrary)>;
+#else
+using unique_dlhandle = std::unique_ptr<void, decltype(&dlclose)>;
+#endif
+
+
+// scan for libraries in the plugin directory to load additional plugins
+std::vector<unique_dlhandle> scanPluginLibraries() {
+  // check and print plugins that are linked directly into the executable
+  int nplugin = mjp_pluginCount();
+  if (nplugin) {
+    std::printf("Built-in plugins:\n");
+    for (int i = 0; i < nplugin; ++i) {
+      std::printf("    %s\n", mjp_getPluginAtSlot(i)->name);
+    }
+  }
+
+  // define platform-specific strings
+#if defined(_WIN32) || defined(__CYGWIN__)
+  const std::string sep = "\\";
+  const std::string dso_suffix = ".dll";
+#else
+  const std::string sep = "/";
+  #if defined(__APPLE__)
+    const std::string dso_suffix = ".dylib";
+  #else
+    const std::string dso_suffix = ".so";
+  #endif
+#endif
+
+  // output vectors containing DSO handles
+  std::vector<unique_dlhandle> dso_handles;
+
+  // platform-independent routine for checking and printing plugins registered by a dynamic library
+  const auto check_and_print_plugins = [&](const std::string& name, unique_dlhandle&& dlhandle) {
+    if (!dlhandle) {
+      return;
+    }
+
+    const int nplugin_new = mjp_pluginCount();
+    if (nplugin_new > nplugin) {
+      dso_handles.push_back(std::move(dlhandle));
+
+      // print all newly registered plugins
+      std::printf("Plugins registered by library '%s':\n", name.c_str());
+      for (int i = nplugin; i < nplugin_new; ++i) {
+        std::printf("    %s\n", mjp_getPluginAtSlot(i)->name);
+      }
+
+      // update counter for plugins registered so far
+      nplugin = nplugin_new;
+    }
+  };
+
+  // try to open the ${EXECDIR}/plugin directory
+  // ${EXECDIR} is the directory containing the simulate binary itself
+  const std::string executable_dir = getExecutableDir();
+  if (executable_dir.empty()) {
+    return dso_handles;
+  }
+
+  const std::string plugin_dir = getExecutableDir() + sep + "plugin";
+
+#if defined(_WIN32) || defined(__CYGWIN__)
+  WIN32_FIND_DATAA find_data;
+  HANDLE hfile = FindFirstFileA((plugin_dir + sep + "*.dll").c_str(), &find_data);
+  if (!hfile) {
+    return dso_handles;
+  }
+
+  // go through each file in the directory
+  bool keep_going = true;
+  while (keep_going) {
+    const std::string name(find_data.cFileName);
+    // load the library and check for plugins
+    const std::string dso_path = plugin_dir + sep + name;
+    check_and_print_plugins(
+        name, unique_dlhandle(LoadLibraryA(dso_path.c_str()), &FreeLibrary));
+    keep_going = FindNextFileA(hfile, &find_data);
+  }
+  FindClose(hfile);
+#else
+  DIR* dirp = opendir(plugin_dir.c_str());
+  if (!dirp) {
+    return dso_handles;
+  }
+
+  // go through each entry in the directory
+  for (struct dirent* dp; (dp = readdir(dirp));) {
+    // only look at regular files (i.e. skip symlinks, pipes, directories, etc.)
+    if (dp->d_type == DT_REG) {
+      const std::string name(dp->d_name);
+      if (name.size() > dso_suffix.size() &&
+          name.substr(name.size() - dso_suffix.size()) == dso_suffix) {
+        // load the library and check for plugins
+        const std::string dso_path = plugin_dir + sep + name;
+        check_and_print_plugins(
+            name, unique_dlhandle(dlopen(dso_path.c_str(), RTLD_NOW | RTLD_LOCAL), &dlclose));
+      }
+    }
+  }
+  closedir(dirp);
+#endif
+
+  return dso_handles;
+}
+
+
 //------------------------------------------- simulation -------------------------------------------
 
 
@@ -285,6 +510,9 @@ int main(int argc, const char** argv) {
     mju_error("Headers and library have different versions");
   }
 
+  // scan for libraries in the plugin directory to load additional plugins
+  std::vector<unique_dlhandle> dso_handles = scanPluginLibraries();
+
   // simulate object encapsulates the UI
   auto sim = std::make_unique<mj::Simulate>();