Properly handle launching script executable (not DLL)

Author: zooba

src/pymanager/_launch.cpp

@@ -34,49 +34,48 @@ dup_handle(HANDLE input, HANDLE *output)
 
 
 int
-launch(const wchar_t *executable, const wchar_t *insert_args, int skip_argc, DWORD *exitCode)
-{
+launch(
+    const wchar_t *executable,
+    const wchar_t *origCmdLine,
+    const wchar_t *insert_args,
+    int skip_argc,
+    DWORD *exitCode
+) {
     HANDLE job;
     JOBOBJECT_EXTENDED_LIMIT_INFORMATION info;
     DWORD info_len;
     STARTUPINFOW si;
     PROCESS_INFORMATION pi;
     int lastError = 0;
-    const wchar_t *arg_space = L" ";
-    LPCWSTR origCmdLine = GetCommandLineW();
     const wchar_t *cmdLine = NULL;
 
-    if (insert_args == NULL) {
-        insert_args = L"";
+    if (origCmdLine[0] == L'"') {
+        cmdLine = wcschr(origCmdLine + 1, L'"');
+    } else {
+        cmdLine = wcschr(origCmdLine, L' ');
     }
 
-    size_t n = wcslen(executable) + wcslen(origCmdLine) + wcslen(insert_args) + 5;
+    size_t n = wcslen(executable) + wcslen(origCmdLine) + wcslen(insert_args) + 6;
     wchar_t *newCmdLine = (wchar_t *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, n * sizeof(wchar_t));
     if (!newCmdLine) {
         lastError = GetLastError();
         goto exit;
     }
 
-    if (origCmdLine[0] == L'"') {
-        cmdLine = wcschr(origCmdLine + 1, L'"');
-    } else {
-        cmdLine = wcschr(origCmdLine, L' ');
-    }
-
+    // Skip any requested args, deliberately leaving any trailing spaces
+    // (we'll skip one later one and add our own space, and preserve multiple)
     while (skip_argc-- > 0) {
         wchar_t c;
         while (*++cmdLine && *cmdLine == L' ') { }
         while (*++cmdLine && *cmdLine != L' ') { }
     }
 
-    if (!insert_args || !*insert_args) {
-        arg_space = L"";
-    }
-    if (cmdLine && *cmdLine) {
-        swprintf_s(newCmdLine, n + 1, L"\"%s\"%s%s %s", executable, arg_space, insert_args, cmdLine + 1);
-    } else {
-        swprintf_s(newCmdLine, n + 1, L"\"%s\"%s%s", executable, arg_space, insert_args);
-    }
+    swprintf_s(newCmdLine, n, L"\"%s\"%s%s%s%s",
+               executable,
+               (insert_args && *insert_args) ? L" ": L"",
+               (insert_args && *insert_args) ? insert_args : L"",
+               (cmdLine && *cmdLine) ? L" " : L"",
+               (cmdLine && *cmdLine) ? cmdLine + 1 : L"");
 
 #if defined(_WINDOWS)
     /*

src/pymanager/_launch.h

@@ -1 +1,7 @@
-int launch(const wchar_t *executable, const wchar_t *insert_args, int skip_argc, DWORD *exitCode);
+int launch(
+    const wchar_t *executable,
+    const wchar_t *origCmdLine,
+    const wchar_t *insert_args,
+    int skip_argc,
+    DWORD *exitCode
+);

src/pymanager/launcher.cpp

@@ -143,7 +143,7 @@ get_executable(wchar_t *executable, unsigned int bufferSize)
 
 
 int
-insert_script(int *argc, wchar_t ***argv)
+get_script(wchar_t **result_path)
 {
     HANDLE ph = GetProcessHeap();
     wchar_t *path = NULL;
@@ -194,22 +194,9 @@ insert_script(int *argc, wchar_t ***argv)
     }
     FindClose(fh);
 
-    // Create a new argv that will be used to launch the script.
-    wchar_t **argv2 = (wchar_t **)HeapAlloc(ph, HEAP_ZERO_MEMORY, sizeof(wchar_t *) * (*argc + 1));
-    if (!argv2) {
-        HeapFree(ph, 0, path);
-        return HRESULT_FROM_WIN32(GetLastError());
-    }
-
     // Deliberately letting our memory leak - it'll be cleaned up when the
     // process ends, and this is not a loop.
-    argv2[0] = (*argv)[0];
-    argv2[1] = path;
-    for (int i = 1; i < (*argc); ++i) {
-        argv2[i + 1] = (*argv)[i];
-    }
-    *argv = argv2;
-    *argc += 1;
+    *result_path = path;
     return 0;
 }
 
@@ -280,21 +267,36 @@ wmain(int argc, wchar_t **argv)
 {
     int exit_code;
     wchar_t executable[MAXLEN];
+    wchar_t *script;
 
     int err = get_executable(executable, MAXLEN);
     if (err) {
         return print_error(err, L"Failed to get target path");
     }
 
-    err = insert_script(&argc, &argv);
+    err = get_script(&script);
     if (err) {
-        return print_error(err, L"Failed to insert script path");
+        return print_error(err, L"Failed to get script path");
     }
 
     void *main_func = NULL;
     err = try_load_python3_dll(executable, MAXLEN, (void **)&main_func);
     switch (err) {
     case 0:
+        if (script) {
+            wchar_t **argv2 = (wchar_t **)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY,
+                (argc + 1) * sizeof(wchar_t *));
+            if (!argv2) {
+                return HRESULT_FROM_WIN32(GetLastError());
+            }
+            argv2[0] = argv[0];
+            argv2[1] = script;
+            for (int i = 1; i < argc; ++i) {
+                argv2[i + 1] = argv[i];
+            }
+            argv = argv2;
+            argc += 1;
+        }
         err = launch_by_dll(main_func, executable, argc, argv, &exit_code);
         if (!err) {
             return exit_code;
@@ -324,7 +326,8 @@ wmain(int argc, wchar_t **argv)
         break;
     }
 
-    err = launch(executable, NULL, 0, (DWORD *)&exit_code);
+    err = launch(executable, GetCommandLineW(), script, 0, (DWORD *)&exit_code);
+
     if (!err) {
         return exit_code;
     }

src/pymanager/main.cpp

@@ -630,7 +630,7 @@ wmain(int argc, wchar_t **argv)
     }
     #endif
 
-    err = launch(executable.c_str(), args.c_str(), skip_argc, &exitCode);
+    err = launch(executable.c_str(), GetCommandLineW(), args.c_str(), skip_argc, &exitCode);
 
     // TODO: Consider sharing print_error() with launcher.cpp
     // This will ensure error messages are aligned whether we're launching