Overhaul parent process detection

Read the entire process ancestry in a single snapshot and then return
the full image names to opam. This makes the shell analysis slightly
more efficient. The full image paths rather than just the executable
names pave the way for opam to be able to initialise PowerShell
properly.

Author: dra27

master_changes.md

@@ -405,6 +405,7 @@ users)
   * [BUG] Fix case insensitive variable handling [#5356 @dra27]
   * Use OCaml code to copy/move/remove directories instead of unix commands [#4823 @kit-ty-kate - fix #1073]
   * Update Windows-on-Windows detection for ARM [#5541 @dra27]
+  * Overhaul parent process detection [#5541 @dra27]
 
 ## Test
   * Update crowbar with compare functions [#4918 @rjbou]

src/core/opamStd.ml

@@ -1060,36 +1060,13 @@ module OpamSys = struct
     else
       fun x -> x
 
-  let windows_max_ancestor_depth = 5
-
-  (** [windows_ancestor_process_names] finds the names of the parent of the
-      current process and all of its ancestors up to [max_ancestor_depth]
-      in length.
-
-      The immediate parent of the current process will be first in the list.
-    *)
-  let windows_ancestor_process_names () =
-    let rec helper pid depth =
-      if depth > windows_max_ancestor_depth then []
-      else
-      try
-        OpamStubs.(getProcessName pid ::
-                   helper
-                     (getParentProcessID pid)
-                     (depth + 1))
-      with Failure _ -> []
-    in
-    lazy (
-      try
-        let parent = OpamStubs.getCurrentProcessID () in
-        helper (OpamStubs.getParentProcessID parent) 0
-      with Failure _ -> []
-    )
+  let windows_process_ancestry = Lazy.from_fun OpamStubs.getProcessAncestry
 
   type shell_choice = Accept of shell
 
   let windows_get_shell =
-    let categorize_process = function
+    let categorize_process (_, image) =
+      match String.lowercase_ascii (Filename.basename image) with
       | "powershell.exe" | "powershell_ise.exe" ->
         Some (Accept (SH_pwsh Powershell))
       | "pwsh.exe" -> Some (Accept (SH_pwsh Powershell_pwsh))
@@ -1101,9 +1078,8 @@ module OpamSys = struct
           (shell_of_string (Filename.chop_suffix name ".exe"))
     in
     lazy (
-      let ancestors = Lazy.force (windows_ancestor_process_names ()) in
-      match (List.map String.lowercase_ascii ancestors |>
-              OpamList.filter_map categorize_process) with
+      let lazy ancestors = windows_process_ancestry in
+      match OpamList.filter_map categorize_process ancestors with
       | [] -> None
       | Accept most_relevant_shell :: _ -> Some most_relevant_shell
     )
@@ -1337,7 +1313,7 @@ module Win32 = struct
   end
 
   let (set_parent_pid, parent_putenv) =
-    let ppid = ref (lazy (OpamStubs.(getCurrentProcessID () |> getParentProcessID))) in
+    let ppid = ref (OpamCompat.Lazy.map (function (_::(pid, _)::_) -> pid | _ -> 0l) OpamSys.windows_process_ancestry) in
     let parent_putenv = lazy (
       let {contents = lazy ppid} = ppid in
       let our_architecture = OpamStubs.getProcessArchitecture None in

src/core/opamStubs.dummy.ml

@@ -33,7 +33,6 @@ let getProcessArchitecture = that's_a_no_no
 let process_putenv _ = that's_a_no_no
 let shGetFolderPath _ = that's_a_no_no
 let sendMessageTimeout _ _ _ _ _ = that's_a_no_no
-let getParentProcessID = that's_a_no_no
-let getProcessName = that's_a_no_no
+let getProcessAncestry = that's_a_no_no
 let getConsoleAlias _ = that's_a_no_no
 let win_create_process _ _ _ _ _ = that's_a_no_no

src/core/opamStubs.mli

@@ -117,16 +117,11 @@ val sendMessageTimeout :
     return value from SendMessageTimeout, [snd] depends on both the message and
     [fst]. See https://msdn.microsoft.com/en-us/library/windows/desktop/ms644952.aspx *)
 
-val getParentProcessID : int32 -> int32
-(** Windows only. [getParentProcessID pid] returns the process ID of the parent
-    of [pid].
-
-    @raise Failure If walking the process tree fails to find the process. *)
-
-val getProcessName : int32 -> string
-(** Windows only. [getProcessName pid] returns the executable name of [pid].
-
-    @raise Failure If the process does not exist. *)
+val getProcessAncestry : unit -> (int32 * string) list
+(** Windows only. Returns the pid and full path to the image for each entry in
+    the ancestry list for this process, starting with the process itself. If an
+    image name can't be determined, then [""] is returned; on failure, returns
+    [[]]. *)
 
 val getConsoleAlias : string -> string -> string
 (** Windows only. [getConsoleAlias alias exeName] retrieves the value for a

src/stubs/win32/opamWin32Stubs.ml

@@ -32,6 +32,5 @@ external getProcessArchitecture : int32 option -> 'a = "OPAMW_GetProcessArchitec
 external process_putenv : int32 -> string -> string -> bool = "OPAMW_process_putenv"
 external shGetFolderPath : int -> 'a -> string = "OPAMW_SHGetFolderPath"
 external sendMessageTimeout : nativeint -> int -> int -> 'a -> 'b -> 'c -> int * 'd = "OPAMW_SendMessageTimeout_byte" "OPAMW_SendMessageTimeout"
-external getParentProcessID : int32 -> int32 = "OPAMW_GetParentProcessID"
-external getProcessName : int32 -> string = "OPAMW_GetProcessName"
+external getProcessAncestry : unit -> (int32 * string) list = "OPAMW_GetProcessAncestry"
 external getConsoleAlias : string -> string -> string = "OPAMW_GetConsoleAlias"

src/stubs/win32/opamWindows.c

@@ -70,45 +70,6 @@ static HKEY roots[] =
    HKEY_LOCAL_MACHINE,
    HKEY_USERS};
 
-/*
- * OPAMW_process_putenv is implemented using Process Injection.
- * Idea inspired by Bill Stewart's editvar
- *   (see http://www.westmesatech.com/editv.html)
- * Full technical details at http://www.codeproject.com/Articles/4610/Three-Ways-to-Inject-Your-Code-into-Another-Proces#section_3
- */
-
-static char* getProcessInfo(HANDLE hProcessSnapshot,
-                            DWORD processId,
-                            PROCESSENTRY32 *entry)
-{
-  entry->dwSize = sizeof(PROCESSENTRY32);
-
-  if (hProcessSnapshot == INVALID_HANDLE_VALUE)
-    return "getProcessInfo: could not create snapshot";
-
-  /*
-   * Locate our process
-   */
-  if (!Process32First(hProcessSnapshot, entry))
-  {
-    CloseHandle(hProcessSnapshot);
-    return "getProcessInfo: could not walk process tree";
-  }
-  else
-  {
-    while (entry->th32ProcessID != processId)
-    {
-      if (!Process32Next(hProcessSnapshot, entry))
-      {
-        CloseHandle(hProcessSnapshot);
-        return "getProcessInfo: could not find process!";
-      }
-    }
-  }
-
-  return NULL;
-}
-
 char* InjectSetEnvironmentVariable(DWORD, LPCWSTR, LPCWSTR);
 
 /* Actual primitives from here */
@@ -545,6 +506,13 @@ CAMLprim value OPAMW_HasGlyph(value checker, value scalar)
   return Val_bool(index != 0xffff);
 }
 
+/*
+ * OPAMW_process_putenv is implemented using Process Injection.
+ * Idea inspired by Bill Stewart's editvar
+ *   (see http://www.westmesatech.com/editv.html)
+ * Full technical details at http://www.codeproject.com/Articles/4610/Three-Ways-to-Inject-Your-Code-into-Another-Proces#section_3
+ */
+
 CAMLprim value OPAMW_process_putenv(value pid, value key, value val)
 {
   char* result;
@@ -650,41 +618,111 @@ CAMLprim value OPAMW_SendMessageTimeout_byte(value * v, int n)
   return OPAMW_SendMessageTimeout(v[0], v[1], v[2], v[3], v[4], v[5]);
 }
 
-CAMLprim value OPAMW_GetParentProcessID(value processId)
+CAMLprim value OPAMW_GetProcessAncestry(value unit)
 {
+  CAMLparam0();
+  CAMLlocal3(result, tail, info);
   PROCESSENTRY32 entry;
-  char* msg;
-  /*
-   * Create a Toolhelp Snapshot of running processes
-   */
-  HANDLE hProcessSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
-
-  if ((msg = getProcessInfo(hProcessSnapshot, Int32_val(processId), &entry)))
-    caml_failwith(msg);
-
-  /*
-   * Finished with the snapshot
-   */
-  CloseHandle(hProcessSnapshot);
-
-  return caml_copy_int32(entry.th32ParentProcessID);
-}
-
-CAMLprim value OPAMW_GetProcessName(value processId)
-{
-  CAMLparam1(processId);
-
-  PROCESSENTRY32 entry;
-  DWORD parent_pid;
-  char* msg;
-  HANDLE hProcessSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
+  HANDLE hProcessSnapshot, hProcess;
+  value cell;
+  ULARGE_INTEGER *processes, *cur;
+  int capacity = 512;
+  int length = 0;
+  DWORD target = GetCurrentProcessId();
+  BOOL read_entry = TRUE;
+  WCHAR ExeName[MAX_PATH + 1];
+  DWORD dwSize;
 
-  if ((msg = getProcessInfo(hProcessSnapshot, Int32_val(processId), &entry)))
-    caml_failwith(msg);
+  result = caml_alloc_small(2, 0);
+  Field(result, 0) = Val_int(0);
+  Field(result, 1) = Val_int(0);
+  tail = result;
+
+  /* Snapshot running processes */
+  hProcessSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
+  if (hProcessSnapshot != INVALID_HANDLE_VALUE) {
+    entry.dwSize = sizeof(PROCESSENTRY32);
+    /* Read the first entry (just because it's a special function) */
+    if (Process32First(hProcessSnapshot, &entry)) {
+      if ((processes = (ULARGE_INTEGER *)malloc(capacity * sizeof(ULARGE_INTEGER)))) {
+        /* Initialise the processes array */
+        if (entry.th32ProcessID == 0) {
+          processes->QuadPart = 0LL;
+        } else {
+          length = 1;
+          processes->LowPart = entry.th32ProcessID;
+          processes->HighPart = entry.th32ParentProcessID;
+          processes[1].QuadPart = 0LL;
+        }
+
+        /* Build the process tree, starting with the current process */
+        do {
+          /* First search through processes we've already read */
+          for (cur = processes; cur->QuadPart != 0; cur++) {
+            if (cur->LowPart == target)
+              break;
+          }
+
+          if (cur->QuadPart == 0LL) {
+            /* Keep reading process entries until we reach the end of the list */
+            while ((read_entry = Process32Next(hProcessSnapshot, &entry))) {
+              if (entry.th32ProcessID != 0) {
+                if (++length >= capacity) {
+                  ULARGE_INTEGER *ptr;
+                  capacity += 512;
+                  ptr = (ULARGE_INTEGER *)realloc(processes, capacity * sizeof(ULARGE_INTEGER));
+                  if (ptr == NULL) {
+                    read_entry = FALSE;
+                    break;
+                  } else {
+                    processes = ptr;
+                  }
+                }
+                cur->LowPart = entry.th32ProcessID;
+                cur->HighPart = entry.th32ParentProcessID;
+                if (cur->LowPart == target) {
+                  cur[1].QuadPart = 0LL;
+                  break;
+                } else {
+                  cur++;
+                }
+              }
+            }
+            if (!read_entry)
+              break;
+          }
+
+          /* Found it - construct the list entry */
+          hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, target);
+          if (hProcess != NULL) {
+            dwSize = MAX_PATH + 1;
+            if (!QueryFullProcessImageName(hProcess, 0, ExeName, &dwSize))
+              ExeName[0] = L'\0';
+            CloseHandle(hProcess);
+          } else {
+            ExeName[0] = L'\0';
+          }
+          info = caml_alloc_tuple(2);
+          Store_field(info, 0, caml_copy_int32(target));
+          Store_field(info, 1, caml_copy_string_of_utf16(ExeName));
+          cell = caml_alloc_small(2, 0);
+          Field(cell, 0) = info;
+          Field(cell, 1) = Val_int(0);
+          Store_field(tail, 1, cell);
+          tail = cell;
+          /* Search for this process's parent on the next round */
+          target = cur->HighPart;
+          /* Guard against looping by zeroing out the parent */
+          cur->HighPart = 0;
+        } while (1);
+      }
+      free(processes);
+    }
 
-  CloseHandle(hProcessSnapshot);
+    CloseHandle(hProcessSnapshot);
+  }
 
-  CAMLreturn(caml_copy_string_of_utf16(entry.szExeFile));
+  CAMLreturn(Field(result, 1));
 }
 
 CAMLprim value OPAMW_GetConsoleAlias(value alias, value exe_name)