lib/Support/Windows/Program.inc

   1 //===- Win32/Program.cpp - Win32 Program Implementation ------- -*- C++ -*-===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file provides the Win32 specific implementation of the Program class.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "Windows.h"
  15 #include <cstdio>
  16 #include <malloc.h>
  17 #include <io.h>
  18 #include <fcntl.h>
  19
  20 //===----------------------------------------------------------------------===//
  21 //=== WARNING: Implementation here must contain only Win32 specific code
  22 //===          and must not be UNIX code
  23 //===----------------------------------------------------------------------===//
  24
  25 namespace {
  26   struct Win32ProcessInfo {
  27     HANDLE hProcess;
  28     DWORD  dwProcessId;
  29   };
  30 }
  31
  32 namespace llvm {
  33 using namespace sys;
  34
  35 Program::Program() : Data_(0) {}
  36
  37 Program::~Program() {
  38   if (Data_) {
  39     Win32ProcessInfo* wpi = reinterpret_cast<Win32ProcessInfo*>(Data_);
  40     CloseHandle(wpi->hProcess);
  41     delete wpi;
  42     Data_ = 0;
  43   }
  44 }
  45
  46 unsigned Program::GetPid() const {
  47   Win32ProcessInfo* wpi = reinterpret_cast<Win32ProcessInfo*>(Data_);
  48   return wpi->dwProcessId;
  49 }
  50
  51 // This function just uses the PATH environment variable to find the program.
  52 Path
  53 Program::FindProgramByName(const std::string& progName) {
  54
  55   // Check some degenerate cases
  56   if (progName.length() == 0) // no program
  57     return Path();
  58   Path temp;
  59   if (!temp.set(progName)) // invalid name
  60     return Path();
  61   // Return paths with slashes verbatim.
  62   if (progName.find('\\') != std::string::npos ||
  63       progName.find('/') != std::string::npos)
  64     return temp;
  65
  66   // At this point, the file name is valid and does not contain slashes.
  67   // Let Windows search for it.
  68   char buffer[MAX_PATH];
  69   char *dummy = NULL;
  70   DWORD len = SearchPath(NULL, progName.c_str(), ".exe", MAX_PATH,
  71                          buffer, &dummy);
  72
  73   // See if it wasn't found.
  74   if (len == 0)
  75     return Path();
  76
  77   // See if we got the entire path.
  78   if (len < MAX_PATH)
  79     return Path(buffer);
  80
  81   // Buffer was too small; grow and retry.
  82   while (true) {
  83     char *b = reinterpret_cast<char *>(_alloca(len+1));
  84     DWORD len2 = SearchPath(NULL, progName.c_str(), ".exe", len+1, b, &dummy);
  85
  86     // It is unlikely the search failed, but it's always possible some file
  87     // was added or removed since the last search, so be paranoid...
  88     if (len2 == 0)
  89       return Path();
  90     else if (len2 <= len)
  91       return Path(b);
  92
  93     len = len2;
  94   }
  95 }
  96
  97 static HANDLE RedirectIO(const Path *path, int fd, std::string* ErrMsg) {
  98   HANDLE h;
  99   if (path == 0) {
 100     DuplicateHandle(GetCurrentProcess(), (HANDLE)_get_osfhandle(fd),
 101                     GetCurrentProcess(), &h,
 102                     0, TRUE, DUPLICATE_SAME_ACCESS);
 103     return h;
 104   }
 105
 106   const char *fname;
 107   if (path->isEmpty())
 108     fname = "NUL";
 109   else
 110     fname = path->c_str();
 111
 112   SECURITY_ATTRIBUTES sa;
 113   sa.nLength = sizeof(sa);
 114   sa.lpSecurityDescriptor = 0;
 115   sa.bInheritHandle = TRUE;
 116
 117   h = CreateFile(fname, fd ? GENERIC_WRITE : GENERIC_READ, FILE_SHARE_READ,
 118                  &sa, fd == 0 ? OPEN_EXISTING : CREATE_ALWAYS,
 119                  FILE_ATTRIBUTE_NORMAL, NULL);
 120   if (h == INVALID_HANDLE_VALUE) {
 121     MakeErrMsg(ErrMsg, std::string(fname) + ": Can't open file for " +
 122         (fd ? "input: " : "output: "));
 123   }
 124
 125   return h;
 126 }
 127
 128 /// ArgNeedsQuotes - Check whether argument needs to be quoted when calling
 129 /// CreateProcess.
 130 static bool ArgNeedsQuotes(const char *Str) {
 131   return Str[0] == '\0' || strpbrk(Str, "\t \"&\'()*<>\\`^|") != 0;
 132 }
 133
 134
 135 /// ArgLenWithQuotes - Check whether argument needs to be quoted when calling
 136 /// CreateProcess and returns length of quoted arg with escaped quotes
 137 static unsigned int ArgLenWithQuotes(const char *Str) {
 138   unsigned int len = ArgNeedsQuotes(Str) ? 2 : 0;
 139
 140   while (*Str != '\0') {
 141     if (*Str == '\"')
 142       ++len;
 143
 144     ++len;
 145     ++Str;
 146   }
 147
 148   return len;
 149 }
 150
 151
 152 bool
 153 Program::Execute(const Path& path,
 154                  const char** args,
 155                  const char** envp,
 156                  const Path** redirects,
 157                  unsigned memoryLimit,
 158                  std::string* ErrMsg) {
 159   if (Data_) {
 160     Win32ProcessInfo* wpi = reinterpret_cast<Win32ProcessInfo*>(Data_);
 161     CloseHandle(wpi->hProcess);
 162     delete wpi;
 163     Data_ = 0;
 164   }
 165
 166   if (!path.canExecute()) {
 167     if (ErrMsg)
 168       *ErrMsg = "program not executable";
 169     return false;
 170   }
 171
 172   // Windows wants a command line, not an array of args, to pass to the new
 173   // process.  We have to concatenate them all, while quoting the args that
 174   // have embedded spaces (or are empty).
 175
 176   // First, determine the length of the command line.
 177   unsigned len = 0;
 178   for (unsigned i = 0; args[i]; i++) {
 179     len += ArgLenWithQuotes(args[i]) + 1;
 180   }
 181
 182   // Now build the command line.
 183   char *command = reinterpret_cast<char *>(_alloca(len+1));
 184   char *p = command;
 185
 186   for (unsigned i = 0; args[i]; i++) {
 187     const char *arg = args[i];
 188
 189     bool needsQuoting = ArgNeedsQuotes(arg);
 190     if (needsQuoting)
 191       *p++ = '"';
 192
 193     while (*arg != '\0') {
 194       if (*arg == '\"')
 195         *p++ = '\\';
 196
 197       *p++ = *arg++;
 198     }
 199
 200     if (needsQuoting)
 201       *p++ = '"';
 202     *p++ = ' ';
 203   }
 204
 205   *p = 0;
 206
 207   // The pointer to the environment block for the new process.
 208   char *envblock = 0;
 209
 210   if (envp) {
 211     // An environment block consists of a null-terminated block of
 212     // null-terminated strings. Convert the array of environment variables to
 213     // an environment block by concatenating them.
 214
 215     // First, determine the length of the environment block.
 216     len = 0;
 217     for (unsigned i = 0; envp[i]; i++)
 218       len += strlen(envp[i]) + 1;
 219
 220     // Now build the environment block.
 221     envblock = reinterpret_cast<char *>(_alloca(len+1));
 222     p = envblock;
 223
 224     for (unsigned i = 0; envp[i]; i++) {
 225       const char *ev = envp[i];
 226       size_t len = strlen(ev) + 1;
 227       memcpy(p, ev, len);
 228       p += len;
 229     }
 230
 231     *p = 0;
 232   }
 233
 234   // Create a child process.
 235   STARTUPINFO si;
 236   memset(&si, 0, sizeof(si));
 237   si.cb = sizeof(si);
 238   si.hStdInput = INVALID_HANDLE_VALUE;
 239   si.hStdOutput = INVALID_HANDLE_VALUE;
 240   si.hStdError = INVALID_HANDLE_VALUE;
 241
 242   if (redirects) {
 243     si.dwFlags = STARTF_USESTDHANDLES;
 244
 245     si.hStdInput = RedirectIO(redirects[0], 0, ErrMsg);
 246     if (si.hStdInput == INVALID_HANDLE_VALUE) {
 247       MakeErrMsg(ErrMsg, "can't redirect stdin");
 248       return false;
 249     }
 250     si.hStdOutput = RedirectIO(redirects[1], 1, ErrMsg);
 251     if (si.hStdOutput == INVALID_HANDLE_VALUE) {
 252       CloseHandle(si.hStdInput);
 253       MakeErrMsg(ErrMsg, "can't redirect stdout");
 254       return false;
 255     }
 256     if (redirects[1] && redirects[2] && *(redirects[1]) == *(redirects[2])) {
 257       // If stdout and stderr should go to the same place, redirect stderr
 258       // to the handle already open for stdout.
 259       DuplicateHandle(GetCurrentProcess(), si.hStdOutput,
 260                       GetCurrentProcess(), &si.hStdError,
 261                       0, TRUE, DUPLICATE_SAME_ACCESS);
 262     } else {
 263       // Just redirect stderr
 264       si.hStdError = RedirectIO(redirects[2], 2, ErrMsg);
 265       if (si.hStdError == INVALID_HANDLE_VALUE) {
 266         CloseHandle(si.hStdInput);
 267         CloseHandle(si.hStdOutput);
 268         MakeErrMsg(ErrMsg, "can't redirect stderr");
 269         return false;
 270       }
 271     }
 272   }
 273
 274   PROCESS_INFORMATION pi;
 275   memset(&pi, 0, sizeof(pi));
 276
 277   fflush(stdout);
 278   fflush(stderr);
 279   BOOL rc = CreateProcess(path.c_str(), command, NULL, NULL, TRUE, 0,
 280                           envblock, NULL, &si, &pi);
 281   DWORD err = GetLastError();
 282
 283   // Regardless of whether the process got created or not, we are done with
 284   // the handles we created for it to inherit.
 285   CloseHandle(si.hStdInput);
 286   CloseHandle(si.hStdOutput);
 287   CloseHandle(si.hStdError);
 288
 289   // Now return an error if the process didn't get created.
 290   if (!rc) {
 291     SetLastError(err);
 292     MakeErrMsg(ErrMsg, std::string("Couldn't execute program '") +
 293                path.str() + "'");
 294     return false;
 295   }
 296   Win32ProcessInfo* wpi = new Win32ProcessInfo;
 297   wpi->hProcess = pi.hProcess;
 298   wpi->dwProcessId = pi.dwProcessId;
 299   Data_ = wpi;
 300
 301   // Make sure these get closed no matter what.
 302   AutoHandle hThread(pi.hThread);
 303
 304   // Assign the process to a job if a memory limit is defined.
 305   AutoHandle hJob(0);
 306   if (memoryLimit != 0) {
 307     hJob = CreateJobObject(0, 0);
 308     bool success = false;
 309     if (hJob != 0) {
 310       JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli;
 311       memset(&jeli, 0, sizeof(jeli));
 312       jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
 313       jeli.ProcessMemoryLimit = uintptr_t(memoryLimit) * 1048576;
 314       if (SetInformationJobObject(hJob, JobObjectExtendedLimitInformation,
 315                                   &jeli, sizeof(jeli))) {
 316         if (AssignProcessToJobObject(hJob, pi.hProcess))
 317           success = true;
 318       }
 319     }
 320     if (!success) {
 321       SetLastError(GetLastError());
 322       MakeErrMsg(ErrMsg, std::string("Unable to set memory limit"));
 323       TerminateProcess(pi.hProcess, 1);
 324       WaitForSingleObject(pi.hProcess, INFINITE);
 325       return false;
 326     }
 327   }
 328
 329   return true;
 330 }
 331
 332 int
 333 Program::Wait(const Path &path,
 334               unsigned secondsToWait,
 335               std::string* ErrMsg) {
 336   if (Data_ == 0) {
 337     MakeErrMsg(ErrMsg, "Process not started!");
 338     return -1;
 339   }
 340
 341   Win32ProcessInfo* wpi = reinterpret_cast<Win32ProcessInfo*>(Data_);
 342   HANDLE hProcess = wpi->hProcess;
 343
 344   // Wait for the process to terminate.
 345   DWORD millisecondsToWait = INFINITE;
 346   if (secondsToWait > 0)
 347     millisecondsToWait = secondsToWait * 1000;
 348
 349   if (WaitForSingleObject(hProcess, millisecondsToWait) == WAIT_TIMEOUT) {
 350     if (!TerminateProcess(hProcess, 1)) {
 351       MakeErrMsg(ErrMsg, "Failed to terminate timed-out program.");
 352       return -1;
 353     }
 354     WaitForSingleObject(hProcess, INFINITE);
 355   }
 356
 357   // Get its exit status.
 358   DWORD status;
 359   BOOL rc = GetExitCodeProcess(hProcess, &status);
 360   DWORD err = GetLastError();
 361
 362   if (!rc) {
 363     SetLastError(err);
 364     MakeErrMsg(ErrMsg, "Failed getting status for program.");
 365     return -1;
 366   }
 367
 368   return status;
 369 }
 370
 371 bool
 372 Program::Kill(std::string* ErrMsg) {
 373   if (Data_ == 0) {
 374     MakeErrMsg(ErrMsg, "Process not started!");
 375     return true;
 376   }
 377
 378   Win32ProcessInfo* wpi = reinterpret_cast<Win32ProcessInfo*>(Data_);
 379   HANDLE hProcess = wpi->hProcess;
 380   if (TerminateProcess(hProcess, 1) == 0) {
 381     MakeErrMsg(ErrMsg, "The process couldn't be killed!");
 382     return true;
 383   }
 384
 385   return false;
 386 }
 387
 388 bool Program::ChangeStdinToBinary(){
 389   int result = _setmode( _fileno(stdin), _O_BINARY );
 390   return result == -1;
 391 }
 392
 393 bool Program::ChangeStdoutToBinary(){
 394   int result = _setmode( _fileno(stdout), _O_BINARY );
 395   return result == -1;
 396 }
 397
 398 bool Program::ChangeStderrToBinary(){
 399   int result = _setmode( _fileno(stderr), _O_BINARY );
 400   return result == -1;
 401 }
 402
 403 }