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