deploy-ng: allow deploy-stub to set MAIN_DIR via blobinfo

direct/src/showutil/FreezeTool.py

@@ -1710,7 +1710,7 @@ class Freezer:
 
         # Determine the format of the header and module list entries depending
         # on the platform.
-        num_pointers = 10
+        num_pointers = 11
         stub_data = bytearray(stub_file.read())
         bitnesses = self._get_executable_bitnesses(stub_data)
 
@@ -1804,6 +1804,7 @@ class Freezer:
                 field_offsets.get('prc_encryption_key', 0),
                 field_offsets.get('prc_executable_patterns', 0),
                 field_offsets.get('prc_executable_args_envvar', 0),
+                field_offsets.get('main_dir', 0),
                 0)
 
             # Now, find the location of the 'blobinfo' symbol in the binary,

dtool/src/interrogatedb/py_panda.cxx

@@ -695,43 +695,45 @@ PyObject *Dtool_PyModuleInitHelper(LibraryDef *defs[], const char *modulename) {
         << "Python " << version << "\n";
     }
 
-    // Grab the __main__ module.
-    PyObject *main_module = PyImport_ImportModule("__main__");
-    if (main_module == NULL) {
-      interrogatedb_cat.warning() << "Unable to import __main__\n";
-    }
+    if (!ExecutionEnvironment::has_environment_variable("MAIN_DIR")) {
+      // Grab the __main__ module.
+      PyObject *main_module = PyImport_ImportModule("__main__");
+      if (main_module == NULL) {
+        interrogatedb_cat.warning() << "Unable to import __main__\n";
+      }
 
-    // Extract the __file__ attribute, if present.
-    Filename main_dir;
-    PyObject *file_attr = PyObject_GetAttrString(main_module, "__file__");
-    if (file_attr == NULL) {
-      // Must be running in the interactive interpreter.  Use the CWD.
-      main_dir = ExecutionEnvironment::get_cwd();
-    } else {
+      // Extract the __file__ attribute, if present.
+      Filename main_dir;
+      PyObject *file_attr = PyObject_GetAttrString(main_module, "__file__");
+      if (file_attr == NULL) {
+        // Must be running in the interactive interpreter.  Use the CWD.
+        main_dir = ExecutionEnvironment::get_cwd();
+      } else {
 #if PY_MAJOR_VERSION >= 3
-      Py_ssize_t length;
-      wchar_t *buffer = PyUnicode_AsWideCharString(file_attr, &length);
-      if (buffer != NULL) {
-        main_dir = Filename::from_os_specific_w(std::wstring(buffer, length));
-        main_dir.make_absolute();
-        main_dir = main_dir.get_dirname();
-        PyMem_Free(buffer);
-      }
+        Py_ssize_t length;
+        wchar_t *buffer = PyUnicode_AsWideCharString(file_attr, &length);
+        if (buffer != NULL) {
+          main_dir = Filename::from_os_specific_w(std::wstring(buffer, length));
+          main_dir.make_absolute();
+          main_dir = main_dir.get_dirname();
+          PyMem_Free(buffer);
+        }
 #else
-      char *buffer;
-      Py_ssize_t length;
-      if (PyString_AsStringAndSize(file_attr, &buffer, &length) != -1) {
-        main_dir = Filename::from_os_specific(std::string(buffer, length));
-        main_dir.make_absolute();
-        main_dir = main_dir.get_dirname();
-      }
+        char *buffer;
+        Py_ssize_t length;
+        if (PyString_AsStringAndSize(file_attr, &buffer, &length) != -1) {
+          main_dir = Filename::from_os_specific(std::string(buffer, length));
+          main_dir.make_absolute();
+          main_dir = main_dir.get_dirname();
+        }
 #endif
-      else {
-        interrogatedb_cat.warning() << "Invalid string for __main__.__file__\n";
+        else {
+          interrogatedb_cat.warning() << "Invalid string for __main__.__file__\n";
+        }
       }
+      ExecutionEnvironment::shadow_environment_variable("MAIN_DIR", main_dir.to_os_specific());
+      PyErr_Clear();
     }
-    ExecutionEnvironment::shadow_environment_variable("MAIN_DIR", main_dir.to_os_specific());
-    PyErr_Clear();
     initialized_main_dir = true;
   }
 

dtool/src/prc/configPageManager.cxx

@@ -114,6 +114,7 @@ reload_implicit_pages() {
     const char *prc_encryption_key;
     const char *prc_executable_patterns;
     const char *prc_executable_args_envvar;
+    const char *main_dir;
   };
 #ifdef _MSC_VER
   const BlobInfo *blobinfo = (const BlobInfo *)GetProcAddress(GetModuleHandle(NULL), "blobinfo");
@@ -128,6 +129,10 @@ reload_implicit_pages() {
     blobinfo = nullptr;
   }
 
+  if (blobinfo != nullptr && blobinfo->num_pointers >= 11 && blobinfo->main_dir != nullptr) {
+    ExecutionEnvironment::shadow_environment_variable("MAIN_DIR", blobinfo->main_dir);
+  }
+
   // PRC_PATTERNS lists one or more filename templates separated by spaces.
   // Pull them out and store them in _prc_patterns.
   _prc_patterns.clear();

pandatool/src/deploy-stub/deploy-stub.c

@@ -2,30 +2,35 @@
 
 #include "Python.h"
 #ifdef _WIN32
-#include "malloc.h"
+#  include "malloc.h"
 #else
-#include <sys/mman.h>
+#  include <sys/mman.h>
 #endif
 
 #ifdef __FreeBSD__
-#include <sys/sysctl.h>
+#  include <sys/sysctl.h>
 #endif
 
 #ifdef __APPLE__
-#include <mach-o/dyld.h>
+#  include <mach-o/dyld.h>
+#  include <libgen.h>
 #endif
 
 #include <stdio.h>
 #include <stdint.h>
 
 #if PY_MAJOR_VERSION >= 3
-#include <locale.h>
+#  include <locale.h>
 
-#if PY_MINOR_VERSION < 5
-#define Py_DecodeLocale _Py_char2wchar
-#endif
+#  if PY_MINOR_VERSION < 5
+#    define Py_DecodeLocale _Py_char2wchar
+#  endif
 #endif
 
+/* Leave room for future expansion.  We only read pointer 0, but there are
+   other pointers that are being read by configPageManager.cxx. */
+#define MAX_NUM_POINTERS 24
+
 /* Define an exposed symbol where we store the offset to the module data. */
 #ifdef _MSC_VER
 __declspec(dllexport)
@@ -40,18 +45,15 @@ volatile struct {
   uint16_t codepage;
   uint16_t flags;
   uint64_t reserved;
-
-  // Leave room for future expansion.  We only read pointer 0, but there are
-  // other pointers that are being read by configPageManager.cxx.
-  void *pointers[24];
+  void *pointers[MAX_NUM_POINTERS];
 
   // The reason we initialize it to -1 is because otherwise, smart linkers may
   // end up putting it in the .bss section for zero-initialized data.
 } blobinfo = {(uint64_t)-1};
 
 #ifdef MS_WINDOWS
-#define WIN32_LEAN_AND_MEAN
-#include <windows.h>
+#  define WIN32_LEAN_AND_MEAN
+#  include <windows.h>
 
 extern void PyWinFreeze_ExeInit(void);
 extern void PyWinFreeze_ExeTerm(void);
@@ -61,7 +63,7 @@ static struct _inittab extensions[] = {
 };
 
 #if PY_MAJOR_VERSION >= 3
-#define WIN_UNICODE
+#  define WIN_UNICODE
 #endif
 #endif
 
@@ -91,6 +93,22 @@ static int supports_code_page(UINT cp) {
 }
 #endif
 
+/**
+ * Sets the main_dir field of the blobinfo structure, but only if it wasn't
+ * already set.
+ */
+static void set_main_dir(char *main_dir) {
+  if (blobinfo.num_pointers >= 10) {
+    if (blobinfo.num_pointers == 10) {
+      ++blobinfo.num_pointers;
+      blobinfo.pointers[10] = NULL;
+    }
+    if (blobinfo.pointers[10] == NULL) {
+      blobinfo.pointers[10] = main_dir;
+    }
+  }
+}
+
 /* Main program */
 
 #ifdef WIN_UNICODE
@@ -356,7 +374,11 @@ int main(int argc, char *argv[]) {
     // Offset the pointers in the header using the base mmap address.
     if (blobinfo.version > 0 && blobinfo.num_pointers > 0) {
       uint32_t i;
+      assert(blobinfo.num_pointers <= MAX_NUM_POINTERS);
       for (i = 0; i < blobinfo.num_pointers; ++i) {
+        // Only offset if the pointer is non-NULL.  Except for the first
+        // pointer, which may never be NULL and usually (but not always)
+        // points to the beginning of the blob.
         if (i == 0 || blobinfo.pointers[i] != 0) {
           blobinfo.pointers[i] = (void *)((uintptr_t)blobinfo.pointers[i] + (uintptr_t)blob);
         }