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