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