win: port multiprocess_exec_test_child.cc

util/test/multiprocess_exec_test_child.cc

Index: util/test/multiprocess_exec_test_child.cc
diff --git a/util/test/multiprocess_exec_test_child.cc b/util/test/multiprocess_exec_test_child.cc
index e1a06122ed1779d9cf1d33b125505053c1e1979c..bc266b2277780ae7cbc9e53cd4eece03f6ca84a1 100644
--- a/util/test/multiprocess_exec_test_child.cc
+++ b/util/test/multiprocess_exec_test_child.cc
@@ -17,11 +17,115 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <sys/types.h>
-#include <unistd.h>
 
 #include <algorithm>
 
+#if defined(__APPLE__) || defined(__linux__)
+#define OS_POSIX 1
+#elif defined(_WIN32)
+#define OS_WIN 1
+#endif
+
+#if defined(OS_POSIX)
+#include <unistd.h>
+#elif defined(OS_WIN)
+#include <windows.h>
+#endif
+
+#if defined(OS_WIN)
+
+// Various semi-documented NT internals to retrieve open handles.
+
+typedef enum _SYSTEM_INFORMATION_CLASS {
+  SystemHandleInformation = 16
+} SYSTEM_INFORMATION_CLASS;
+
+typedef struct _SYSTEM_HANDLE_INFORMATION {
+  USHORT ProcessId;
+  USHORT CreatorBackTraceIndex;
+  UCHAR ObjectTypeNumber;
+  UCHAR Flags;
+  USHORT Handle;
+  PVOID Object;
+  ACCESS_MASK GrantedAccess;
+} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;
+
+typedef struct _SYSTEM_HANDLE_INFORMATION_EX {
+  ULONG NumberOfHandles;
+  SYSTEM_HANDLE_INFORMATION Information[1];
+} SYSTEM_HANDLE_INFORMATION_EX, *PSYSTEM_HANDLE_INFORMATION_EX;
+
+typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
+    SYSTEM_INFORMATION_CLASS SystemInformationClass,
+    PVOID SystemInformation,
+    ULONG SystemInformationLength,
+    PULONG ReturnLength);
+
+void EnsureOnlyStdioHandlesOpen() {

[Mark Mentovai, 2015/01/08 18:58:15]

Unnamed namespace for this.

[scottmg, 2015/01/08 19:07:43]

Done.

+  // Initialize the NTAPI functions we need.
+  HMODULE ntdll_handle = GetModuleHandle(L"ntdll.dll");
+  if (!ntdll_handle) {
+    fprintf(stderr, "GetModuleHandle ntdll.dll failed.\n");
+    abort();
+  }
+
+  NTQUERYSYSTEMINFORMATION NtQuerySystemInformation;
+  NtQuerySystemInformation = reinterpret_cast<NTQUERYSYSTEMINFORMATION>(
+                      GetProcAddress(ntdll_handle, "NtQuerySystemInformation"));
+  if (!NtQuerySystemInformation) {
+    fprintf(stderr, "GetProcAddress NtQuerySystemInformation failed.\n");
+    abort();
+  }
+
+  // Get the number of handles on the system.
+  DWORD buffer_size = 0;
+  SYSTEM_HANDLE_INFORMATION_EX temp_info;
+  NTSTATUS status = NtQuerySystemInformation(
+      SystemHandleInformation, &temp_info, sizeof(temp_info), &buffer_size);
+  if (!buffer_size) {
+    fprintf(stderr,
+            "NtQuerySystemInformation for number of handles failed: 0x%lX\n",
+            status);
+    abort();
+  }
+
+  SYSTEM_HANDLE_INFORMATION_EX *system_handles =
+      reinterpret_cast<SYSTEM_HANDLE_INFORMATION_EX*>(new BYTE[buffer_size]);
+
+  // This is likely flaky as we're racing with other handles being created on
+  // the system between the size query above, and the actual retrieval here.

[Mark Mentovai, 2015/01/08 18:58:15]

Well, hopefully we’re single-threaded here.

[scottmg, 2015/01/08 19:07:43]

Unfortunately, NtQuerySystemInformation has no way to filter by process at
retrieval time so we're retrieving handles for the entire system. If it's a
problem in practice, it might be sufficient to make the buffer arbitrarily
larger.

+  status = NtQuerySystemInformation(SystemHandleInformation, system_handles,
+                                    buffer_size, &buffer_size);
+  if (status != 0) {
+    fprintf(stderr, "Failed to get the handle list: 0x%lX\n", status);
+    delete[] system_handles;
+    abort();
+  }
+
+  for (ULONG i = 0; i < system_handles->NumberOfHandles; ++i) {
+    USHORT h = system_handles->Information[i].Handle;
+    if (system_handles->Information[i].ProcessId != GetCurrentProcessId())
+      continue;
+
+    // TODO(scottmg): This is probably insufficient, we'll need to allow having
+    // a few other standard handles open (for example, to the window station),
+    // or only check for handles of certain types.
+    HANDLE handle = reinterpret_cast<HANDLE>(h);
+    if (handle != GetStdHandle(STD_INPUT_HANDLE) &&
+        handle != GetStdHandle(STD_OUTPUT_HANDLE) &&
+        handle != GetStdHandle(STD_ERROR_HANDLE)) {
+      fprintf(stderr, "Handle 0x%lX is not stdio handle\n", handle);
+      abort();
+    }
+  }
+
+  delete [] system_handles;
+}
+
+#endif  // OS_WIN
+
 int main(int argc, char* argv[]) {
+#if defined(OS_POSIX)
   // Make sure that there’s nothing open at any FD higher than 3. All FDs other
   // than stdin, stdout, and stderr should have been closed prior to or at
   // exec().
@@ -46,6 +150,27 @@ int main(int argc, char* argv[]) {
   if (rv != 1) {
     abort();
   }
+#elif defined(OS_WIN)
+  // Make sure there's nothing open other than stdin, stdout, and stderr.
+  EnsureOnlyStdioHandlesOpen();
+
+  // Read a byte from stdin, expecting it to be a specific value.
+  char c;
+  DWORD bytes_read;
+  if (!ReadFile(GetStdHandle(STD_INPUT_HANDLE), &c, 1, &bytes_read, nullptr) ||
+      bytes_read != 1 || c != 'z') {
+    abort();
+  }
+
+  // Write a byte to stdout.
+  c = 'Z';
+  DWORD bytes_written;
+  if (!WriteFile(
+          GetStdHandle(STD_OUTPUT_HANDLE), &c, 1, &bytes_written, nullptr) ||
+      bytes_written != 1) {
+    abort();
+  }
+#endif  // OS_POSIX
 
   return 0;
 }