gh-100247: Fix py.exe launcher not using entire shebang command for finding custom commands (GH-100944)

1 files changed, 85 insertions, 65 deletions

diff --git a/PC/launcher2.c b/PC/launcher2.c
index 9b3db04..8371c60 100644
--- a/PC/launcher2.c
+++ b/PC/launcher2.c
@@ -16,6 +16,7 @@
 #include <stdio.h>
 #include <stdbool.h>
 #include <tchar.h>
+#include <assert.h>
 
 #define MS_WINDOWS
 #include "patchlevel.h"
@@ -37,6 +38,7 @@
 #define RC_INSTALLING       111
 #define RC_NO_PYTHON_AT_ALL 112
 #define RC_NO_SHEBANG       113
+#define RC_RECURSIVE_SHEBANG 114
 
 static FILE * log_fp = NULL;
 
@@ -702,16 +704,23 @@ _decodeShebang(SearchInfo *search, const char *buffer, int bufferLength, bool on
 
 
 bool
-_shebangStartsWith(const wchar_t *buffer, int bufferLength, const wchar_t *prefix, const wchar_t **rest)
+_shebangStartsWith(const wchar_t *buffer, int bufferLength, const wchar_t *prefix, const wchar_t **rest, int *firstArgumentLength)
 {
     int prefixLength = (int)wcsnlen_s(prefix, MAXLEN);
-    if (bufferLength < prefixLength) {
+    if (bufferLength < prefixLength || !_startsWithArgument(buffer, bufferLength, prefix, prefixLength)) {
         return false;
     }
     if (rest) {
         *rest = &buffer[prefixLength];
     }
-    return _startsWithArgument(buffer, bufferLength, prefix, prefixLength);
+    if (firstArgumentLength) {
+        int i = prefixLength;
+        while (i < bufferLength && !isspace(buffer[i])) {
+            i += 1;
+        }
+        *firstArgumentLength = i - prefixLength;
+    }
+    return true;
 }
 
 
@@ -723,26 +732,27 @@ searchPath(SearchInfo *search, const wchar_t *shebang, int shebangLength)
     }
 
     wchar_t *command;
-    if (!_shebangStartsWith(shebang, shebangLength, L"/usr/bin/env ", &command)) {
+    int commandLength;
+    if (!_shebangStartsWith(shebang, shebangLength, L"/usr/bin/env ", &command, &commandLength)) {
         return RC_NO_SHEBANG;
     }
 
-    wchar_t filename[MAXLEN];
-    int lastDot = 0;
-    int commandLength = 0;
-    while (commandLength < MAXLEN && command[commandLength] && !isspace(command[commandLength])) {
-        if (command[commandLength] == L'.') {
-            lastDot = commandLength;
-        }
-        filename[commandLength] = command[commandLength];
-        commandLength += 1;
-    }
-
     if (!commandLength || commandLength == MAXLEN) {
         return RC_BAD_VIRTUAL_PATH;
     }
 
-    filename[commandLength] = L'\0';
+    int lastDot = commandLength;
+    while (lastDot > 0 && command[lastDot] != L'.') {
+        lastDot -= 1;
+    }
+    if (!lastDot) {
+        lastDot = commandLength;
+    }
+
+    wchar_t filename[MAXLEN];
+    if (wcsncpy_s(filename, MAXLEN, command, lastDot)) {
+        return RC_BAD_VIRTUAL_PATH;
+    }
 
     const wchar_t *ext = L".exe";
     // If the command already has an extension, we do not want to add it again
@@ -780,7 +790,7 @@ searchPath(SearchInfo *search, const wchar_t *shebang, int shebangLength)
     if (GetModuleFileNameW(NULL, filename, MAXLEN) &&
         0 == _comparePath(filename, -1, buffer, -1)) {
         debug(L"# ignoring recursive shebang command\n");
-        return RC_NO_SHEBANG;
+        return RC_RECURSIVE_SHEBANG;
     }
 
     wchar_t *buf = allocSearchInfoBuffer(search, n + 1);
@@ -994,73 +1004,78 @@ checkShebang(SearchInfo *search)
         return exitCode;
     }
 
-    // Handle some known, case-sensitive shebang templates
+    // Handle some known, case-sensitive shebangs
     const wchar_t *command;
     int commandLength;
+    // Each template must end with "python"
     static const wchar_t *shebangTemplates[] = {
-        L"/usr/bin/env ",
-        L"/usr/bin/",
-        L"/usr/local/bin/",
+        L"/usr/bin/env python",
+        L"/usr/bin/python",
+        L"/usr/local/bin/python",
         L"python",
         NULL
     };
 
     for (const wchar_t **tmpl = shebangTemplates; *tmpl; ++tmpl) {
-        if (_shebangStartsWith(shebang, shebangLength, *tmpl, &command)) {
-            commandLength = 0;
-            // Normally "python" is the start of the command, but we also need it
-            // as a shebang prefix for back-compat. We move the command marker back
-            // if we match on that one.
-            if (0 == wcscmp(*tmpl, L"python")) {
-                command -= 6;
-            }
-            while (command[commandLength] && !isspace(command[commandLength])) {
-                commandLength += 1;
-            }
-            if (!commandLength) {
-            } else if (_findCommand(search, command, commandLength)) {
+        // Just to make sure we don't mess this up in the future
+        assert(0 == wcscmp(L"python", (*tmpl) + wcslen(*tmpl) - 6));
+
+        if (_shebangStartsWith(shebang, shebangLength, *tmpl, &command, &commandLength)) {
+            // Search for "python{command}" overrides. All templates end with
+            // "python", so we prepend it by jumping back 6 characters
+            if (_findCommand(search, &command[-6], commandLength + 6)) {
                 search->executableArgs = &command[commandLength];
                 search->executableArgsLength = shebangLength - commandLength;
                 debug(L"# Treating shebang command '%.*s' as %s\n",
-                    commandLength, command, search->executablePath);
-            } else if (_shebangStartsWith(command, commandLength, L"python", NULL)) {
-                search->tag = &command[6];
-                search->tagLength = commandLength - 6;
-                // If we had 'python3.12.exe' then we want to strip the suffix
-                // off of the tag
-                if (search->tagLength > 4) {
-                    const wchar_t *suffix = &search->tag[search->tagLength - 4];
-                    if (0 == _comparePath(suffix, 4, L".exe", -1)) {
-                        search->tagLength -= 4;
-                    }
-                }
-                // If we had 'python3_d' then we want to strip the '_d' (any
-                // '.exe' is already gone)
-                if (search->tagLength > 2) {
-                    const wchar_t *suffix = &search->tag[search->tagLength - 2];
-                    if (0 == _comparePath(suffix, 2, L"_d", -1)) {
-                        search->tagLength -= 2;
-                    }
+                    commandLength + 6, &command[-6], search->executablePath);
+                return 0;
+            }
+
+            search->tag = command;
+            search->tagLength = commandLength;
+            // If we had 'python3.12.exe' then we want to strip the suffix
+            // off of the tag
+            if (search->tagLength > 4) {
+                const wchar_t *suffix = &search->tag[search->tagLength - 4];
+                if (0 == _comparePath(suffix, 4, L".exe", -1)) {
+                    search->tagLength -= 4;
                 }
-                search->oldStyleTag = true;
-                search->executableArgs = &command[commandLength];
-                search->executableArgsLength = shebangLength - commandLength;
-                if (search->tag && search->tagLength) {
-                    debug(L"# Treating shebang command '%.*s' as 'py -%.*s'\n",
-                        commandLength, command, search->tagLength, search->tag);
-                } else {
-                    debug(L"# Treating shebang command '%.*s' as 'py'\n",
-                        commandLength, command);
+            }
+            // If we had 'python3_d' then we want to strip the '_d' (any
+            // '.exe' is already gone)
+            if (search->tagLength > 2) {
+                const wchar_t *suffix = &search->tag[search->tagLength - 2];
+                if (0 == _comparePath(suffix, 2, L"_d", -1)) {
+                    search->tagLength -= 2;
                 }
+            }
+            search->oldStyleTag = true;
+            search->executableArgs = &command[commandLength];
+            search->executableArgsLength = shebangLength - commandLength;
+            if (search->tag && search->tagLength) {
+                debug(L"# Treating shebang command '%.*s' as 'py -%.*s'\n",
+                    commandLength, command, search->tagLength, search->tag);
             } else {
-                debug(L"# Found shebang command but could not execute it: %.*s\n",
+                debug(L"# Treating shebang command '%.*s' as 'py'\n",
                     commandLength, command);
             }
-            // search is done by this point
             return 0;
         }
     }
 
+    // Unrecognised executables are first tried as command aliases
+    commandLength = 0;
+    while (commandLength < shebangLength && !isspace(shebang[commandLength])) {
+        commandLength += 1;
+    }
+    if (_findCommand(search, shebang, commandLength)) {
+        search->executableArgs = &shebang[commandLength];
+        search->executableArgsLength = shebangLength - commandLength;
+        debug(L"# Treating shebang command '%.*s' as %s\n",
+            commandLength, shebang, search->executablePath);
+        return 0;
+    }
+
     // Unrecognised commands are joined to the script's directory and treated
     // as the executable path
     return _useShebangAsExecutable(search, shebang, shebangLength);
@@ -2407,7 +2422,12 @@ performSearch(SearchInfo *search, EnvironmentInfo **envs)
     // Check for a shebang line in our script file
     // (or return quickly if no script file was specified)
     exitCode = checkShebang(search);
-    if (exitCode) {
+    switch (exitCode) {
+    case 0:
+    case RC_NO_SHEBANG:
+    case RC_RECURSIVE_SHEBANG:
+        break;
+    default:
         return exitCode;
     }