Guard dynamic library loading.
[oota-llvm.git] / lib / System / Win32 / Program.inc
index e1ad155bb34d5fbffcf30d556f025914dec925d6..49086b8348e6a07c8cd99dc88ae47a9d1d14d99f 100644 (file)
@@ -1,10 +1,10 @@
 //===- Win32/Program.cpp - Win32 Program Implementation ------- -*- C++ -*-===//
-// 
+//
 //                     The LLVM Compiler Infrastructure
 //
-// This file was developed by Jeff Cohen and is distributed under the 
-// University of Illinois Open Source License. See LICENSE.TXT for details.
-// 
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
 //===----------------------------------------------------------------------===//
 //
 // This file provides the Win32 specific implementation of the Program class.
@@ -18,7 +18,7 @@
 #include <fcntl.h>
 
 //===----------------------------------------------------------------------===//
-//=== WARNING: Implementation here must contain only Win32 specific code 
+//=== WARNING: Implementation here must contain only Win32 specific code
 //===          and must not be UNIX code
 //===----------------------------------------------------------------------===//
 
@@ -78,9 +78,11 @@ static HANDLE RedirectIO(const Path *path, int fd, std::string* ErrMsg) {
     return h;
   }
 
-  const char *fname = path->toString().c_str();
-  if (*fname == 0)
+  const char *fname;
+  if (path->isEmpty())
     fname = "NUL";
+  else
+    fname = path->toString().c_str();
 
   SECURITY_ATTRIBUTES sa;
   sa.nLength = sizeof(sa);
@@ -93,17 +95,27 @@ static HANDLE RedirectIO(const Path *path, int fd, std::string* ErrMsg) {
   if (h == INVALID_HANDLE_VALUE) {
     MakeErrMsg(ErrMsg, std::string(fname) + ": Can't open file for " +
         (fd ? "input: " : "output: "));
-    return h;
   }
+
   return h;
 }
 
-int 
-Program::ExecuteAndWait(const Path& path, 
+#ifdef __MINGW32__
+  // Due to unknown reason, mingw32's w32api doesn't have this declaration.
+  extern "C"
+  BOOL WINAPI SetInformationJobObject(HANDLE hJob,
+                                      JOBOBJECTINFOCLASS JobObjectInfoClass,
+                                      LPVOID lpJobObjectInfo,
+                                      DWORD cbJobObjectInfoLength);
+#endif
+
+int
+Program::ExecuteAndWait(const Path& path,
                         const char** args,
                         const char** envp,
                         const Path** redirects,
                         unsigned secondsToWait,
+                        unsigned memoryLimit,
                         std::string* ErrMsg) {
   if (!path.canExecute()) {
     if (ErrMsg)
@@ -124,7 +136,7 @@ Program::ExecuteAndWait(const Path& path,
   }
 
   // Now build the command line.
-  char *command = reinterpret_cast<char *>(_alloca(len));
+  char *command = reinterpret_cast<char *>(_alloca(len+1));
   char *p = command;
 
   for (unsigned i = 0; args[i]; i++) {
@@ -142,6 +154,33 @@ Program::ExecuteAndWait(const Path& path,
 
   *p = 0;
 
+  // The pointer to the environment block for the new process.
+  char *envblock = 0;
+
+  if (envp) {
+    // An environment block consists of a null-terminated block of
+    // null-terminated strings. Convert the array of environment variables to
+    // an environment block by concatenating them.
+
+    // First, determine the length of the environment block.
+    len = 0;
+    for (unsigned i = 0; envp[i]; i++)
+      len += strlen(envp[i]) + 1;
+
+    // Now build the environment block.
+    envblock = reinterpret_cast<char *>(_alloca(len+1));
+    p = envblock;
+
+    for (unsigned i = 0; envp[i]; i++) {
+      const char *ev = envp[i];
+      size_t len = strlen(ev) + 1;
+      memcpy(p, ev, len);
+      p += len;
+    }
+
+    *p = 0;
+  }
+
   // Create a child process.
   STARTUPINFO si;
   memset(&si, 0, sizeof(si));
@@ -152,7 +191,7 @@ Program::ExecuteAndWait(const Path& path,
 
   if (redirects) {
     si.dwFlags = STARTF_USESTDHANDLES;
-    
+
     si.hStdInput = RedirectIO(redirects[0], 0, ErrMsg);
     if (si.hStdInput == INVALID_HANDLE_VALUE) {
       MakeErrMsg(ErrMsg, "can't redirect stdin");
@@ -164,7 +203,14 @@ Program::ExecuteAndWait(const Path& path,
       MakeErrMsg(ErrMsg, "can't redirect stdout");
       return -1;
     }
-    if (redirects[1] && redirects[2] && *(redirects[1]) != *(redirects[2])) {
+    if (redirects[1] && redirects[2] && *(redirects[1]) == *(redirects[2])) {
+      // If stdout and stderr should go to the same place, redirect stderr
+      // to the handle already open for stdout.
+      DuplicateHandle(GetCurrentProcess(), si.hStdOutput,
+                      GetCurrentProcess(), &si.hStdError,
+                      0, TRUE, DUPLICATE_SAME_ACCESS);
+    } else {
+      // Just redirect stderr
       si.hStdError = RedirectIO(redirects[2], 2, ErrMsg);
       if (si.hStdError == INVALID_HANDLE_VALUE) {
         CloseHandle(si.hStdInput);
@@ -172,10 +218,6 @@ Program::ExecuteAndWait(const Path& path,
         MakeErrMsg(ErrMsg, "can't redirect stderr");
         return -1;
       }
-    } else {
-      DuplicateHandle(GetCurrentProcess(), si.hStdOutput,
-                      GetCurrentProcess(), &si.hStdError,
-                      0, TRUE, DUPLICATE_SAME_ACCESS);
     }
   }
 
@@ -184,8 +226,8 @@ Program::ExecuteAndWait(const Path& path,
 
   fflush(stdout);
   fflush(stderr);
-  BOOL rc = CreateProcess(path.c_str(), command, NULL, NULL, FALSE, 0,
-                          envp, NULL, &si, &pi);
+  BOOL rc = CreateProcess(path.c_str(), command, NULL, NULL, TRUE, 0,
+                          envblock, NULL, &si, &pi);
   DWORD err = GetLastError();
 
   // Regardless of whether the process got created or not, we are done with
@@ -198,11 +240,40 @@ Program::ExecuteAndWait(const Path& path,
   if (!rc)
   {
     SetLastError(err);
-    MakeErrMsg(ErrMsg, std::string("Couldn't execute program '") + 
+    MakeErrMsg(ErrMsg, std::string("Couldn't execute program '") +
                path.toString() + "'");
     return -1;
   }
 
+  // Make sure these get closed no matter what.
+  AutoHandle hProcess(pi.hProcess);
+  AutoHandle hThread(pi.hThread);
+
+  // Assign the process to a job if a memory limit is defined.
+  AutoHandle hJob(0);
+  if (memoryLimit != 0) {
+    hJob = CreateJobObject(0, 0);
+    bool success = false;
+    if (hJob != 0) {
+      JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli;
+      memset(&jeli, 0, sizeof(jeli));
+      jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
+      jeli.ProcessMemoryLimit = uintptr_t(memoryLimit) * 1048576;
+      if (SetInformationJobObject(hJob, JobObjectExtendedLimitInformation,
+                                  &jeli, sizeof(jeli))) {
+        if (AssignProcessToJobObject(hJob, pi.hProcess))
+          success = true;
+      }
+    }
+    if (!success) {
+      SetLastError(GetLastError());
+      MakeErrMsg(ErrMsg, std::string("Unable to set memory limit"));
+      TerminateProcess(pi.hProcess, 1);
+      WaitForSingleObject(pi.hProcess, INFINITE);
+      return -1;
+    }
+  }
+
   // Wait for it to terminate.
   DWORD millisecondsToWait = INFINITE;
   if (secondsToWait > 0)
@@ -216,19 +287,15 @@ Program::ExecuteAndWait(const Path& path,
     }
     WaitForSingleObject(pi.hProcess, INFINITE);
   }
-  
+
   // Get its exit status.
   DWORD status;
   rc = GetExitCodeProcess(pi.hProcess, &status);
   err = GetLastError();
 
-  // Done with the handles; go close them.
-  CloseHandle(pi.hProcess);
-  CloseHandle(pi.hThread);
-
   if (!rc) {
     SetLastError(err);
-    MakeErrMsg(ErrMsg, std::string("Failed getting status for program '") + 
+    MakeErrMsg(ErrMsg, std::string("Failed getting status for program '") +
                path.toString() + "'");
     return -1;
   }