From 3c7ad0eab2a3c0942c353531d10bb693f8560169 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C5=8Dshin?= <git@wholezero.org>
Date: Sun, 17 Dec 2023 13:38:24 +0000
Subject: [PATCH] TryPath, verify __program_executable_name

---
 libc/calls/getprogramexecutablename.greg.c | 78 ++++++++++++----------
 1 file changed, 41 insertions(+), 37 deletions(-)

diff --git a/libc/calls/getprogramexecutablename.greg.c b/libc/calls/getprogramexecutablename.greg.c
index 05c690b98..4816bf69c 100644
--- a/libc/calls/getprogramexecutablename.greg.c
+++ b/libc/calls/getprogramexecutablename.greg.c
@@ -50,15 +50,46 @@ static inline int IsAlpha(int c) {
   return ('A' <= c && c <= 'Z') || ('a' <= c && c <= 'z');
 }
 
+// if q exists then turn it into an absolute path. we also try adding
+// a .com suffix since the ape auto-appends it when resolving
+static int TryPath(const char *q, int com) {
+  if (!q) return 0;
+  char *p = g_prog.u.buf;
+  char *e = p + sizeof(g_prog.u.buf);
+  int c, f_ok;
+  if ((f_ok = !sys_faccessat(AT_FDCWD, q, F_OK, 0))) {
+    com = 0;
+  }
+  if (*q != '/') {
+    if (q[0] == '.' && q[1] == '/') {
+      q += 2;
+    }
+    int got = __getcwd(p, e - p - 1 - com * 4);   // for / and .com
+    if (got != -1) {
+      p += got - 1;
+      *p++ = '/';
+    }
+  }
+  while ((c = *q++)) {
+    if (p + 1 + com * 4 < e) {  // for nul and .com
+      *p++ = c;
+    }
+  }
+  if (f_ok) {
+    *p = 0;
+    return 1;
+  }
+  if (!com) return 0;
+  p = WRITE32LE(p, READ32LE(".com"));
+  *p = 0;
+  return !sys_faccessat(AT_FDCWD, g_prog.u.buf, F_OK, 0);
+}
+
 static inline void InitProgramExecutableNameImpl(void) {
   size_t n;
   ssize_t got;
   char c, *q, *b;
 
-  if (__program_executable_name) {
-    /* already set by the loader */
-    return;
-  }
   if (IsWindows()) {
     int n = GetModuleFileName(0, g_prog.u.buf16, ARRAYLEN(g_prog.u.buf16));
     for (int i = 0; i < n; ++i) {
@@ -116,43 +147,16 @@ static inline void InitProgramExecutableNameImpl(void) {
     goto UseEmpty;
   }
 
-  // if argv[0] exists then turn it into an absolute path. we also try
-  // adding a .com suffix since the ape auto-appends it when resolving
-  if ((q = __argv[0])) {
-    char *p = g_prog.u.buf;
-    char *e = p + sizeof(g_prog.u.buf);
-    if (*q != '/') {
-      if (q[0] == '.' && q[1] == '/') {
-        q += 2;
-      }
-      int got = __getcwd(p, e - p - 1 - 4);  // for / and .com
-      if (got != -1) {
-        p += got - 1;
-        *p++ = '/';
-      }
-    }
-    while ((c = *q++)) {
-      if (p + 1 + 4 < e) {  // for nul and .com
-        *p++ = c;
-      }
-    }
-    *p = 0;
-    if (!sys_faccessat(AT_FDCWD, g_prog.u.buf, F_OK, 0)) goto UseBuf;
-    p = WRITE32LE(p, READ32LE(".com"));
-    *p = 0;
-    if (!sys_faccessat(AT_FDCWD, g_prog.u.buf, F_OK, 0)) goto UseBuf;
-  }
-
-  /* the previous loader supplied the full program path as the first
-     environment variable. we also try "_". */
-  if ((q = __getenv(__envp, "COSMOPOLITAN_PROGRAM_EXECUTABLE").s) ||
-      (q = __getenv(__envp, "_").s)) {
-    goto CopyString;
+  // Try what the loader supplied first. Fall back to argv[0],
+  // then argv[0].com, then $_, then $_.com.
+  if (TryPath(__program_executable_name, 0) ||
+      TryPath(__argv[0], 0) || TryPath(__argv[0], 1) ||
+      TryPath((q = __getenv(__envp, "_").s), 0) || TryPath(q, 1)) {
+    goto UseBuf;
   }
 
   // give up and just copy argv[0] into it
   if ((q = __argv[0])) {
-  CopyString:
     char *p = g_prog.u.buf;
     char *e = p + sizeof(g_prog.u.buf);
     while ((c = *q++)) {