Fix error in the Win32 implementation pointed out by Howard Su.
[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 int
113 Program::ExecuteAndWait(const Path& path,
114                         const char** args,
115                         const char** envp,
116                         const Path** redirects,
117                         unsigned secondsToWait,
118                         unsigned memoryLimit,
119                         std::string* ErrMsg) {
120   if (!path.canExecute()) {
121     if (ErrMsg)
122       *ErrMsg = "program not executable";
123     return -1;
124   }
125
126   // Windows wants a command line, not an array of args, to pass to the new
127   // process.  We have to concatenate them all, while quoting the args that
128   // have embedded spaces.
129
130   // First, determine the length of the command line.
131   unsigned len = 0;
132   for (unsigned i = 0; args[i]; i++) {
133     len += strlen(args[i]) + 1;
134     if (strchr(args[i], ' '))
135       len += 2;
136   }
137
138   // Now build the command line.
139   char *command = reinterpret_cast<char *>(_alloca(len+1));
140   char *p = command;
141
142   for (unsigned i = 0; args[i]; i++) {
143     const char *arg = args[i];
144     size_t len = strlen(arg);
145     bool needsQuoting = strchr(arg, ' ') != 0;
146     if (needsQuoting)
147       *p++ = '"';
148     memcpy(p, arg, len);
149     p += len;
150     if (needsQuoting)
151       *p++ = '"';
152     *p++ = ' ';
153   }
154
155   *p = 0;
156
157   // The pointer to the environment block for the new process.
158   char *envblock = 0;
159
160   if (envp) {
161     // An environment block consists of a null-terminated block of
162     // null-terminated strings. Convert the array of environment variables to
163     // an environment block by concatenating them.
164
165     // First, determine the length of the environment block.
166     len = 0;
167     for (unsigned i = 0; envp[i]; i++)
168       len += strlen(envp[i]) + 1;
169
170     // Now build the environment block.
171     envblock = reinterpret_cast<char *>(_alloca(len+1));
172     p = envblock;
173
174     for (unsigned i = 0; envp[i]; i++) {
175       const char *ev = envp[i];
176       size_t len = strlen(ev) + 1;
177       memcpy(p, ev, len);
178       p += len;
179     }
180
181     *p = 0;
182   }
183
184   // Create a child process.
185   STARTUPINFO si;
186   memset(&si, 0, sizeof(si));
187   si.cb = sizeof(si);
188   si.hStdInput = INVALID_HANDLE_VALUE;
189   si.hStdOutput = INVALID_HANDLE_VALUE;
190   si.hStdError = INVALID_HANDLE_VALUE;
191
192   if (redirects) {
193     si.dwFlags = STARTF_USESTDHANDLES;
194
195     si.hStdInput = RedirectIO(redirects[0], 0, ErrMsg);
196     if (si.hStdInput == INVALID_HANDLE_VALUE) {
197       MakeErrMsg(ErrMsg, "can't redirect stdin");
198       return -1;
199     }
200     si.hStdOutput = RedirectIO(redirects[1], 1, ErrMsg);
201     if (si.hStdOutput == INVALID_HANDLE_VALUE) {
202       CloseHandle(si.hStdInput);
203       MakeErrMsg(ErrMsg, "can't redirect stdout");
204       return -1;
205     }
206     if (redirects[1] && redirects[2] && *(redirects[1]) == *(redirects[2])) {
207       // If stdout and stderr should go to the same place, redirect stderr
208       // to the handle already open for stdout.
209       DuplicateHandle(GetCurrentProcess(), si.hStdOutput,
210                       GetCurrentProcess(), &si.hStdError,
211                       0, TRUE, DUPLICATE_SAME_ACCESS);
212     } else {
213       // Just redirect stderr
214       si.hStdError = RedirectIO(redirects[2], 2, ErrMsg);
215       if (si.hStdError == INVALID_HANDLE_VALUE) {
216         CloseHandle(si.hStdInput);
217         CloseHandle(si.hStdOutput);
218         MakeErrMsg(ErrMsg, "can't redirect stderr");
219         return -1;
220       }
221     }
222   }
223
224   PROCESS_INFORMATION pi;
225   memset(&pi, 0, sizeof(pi));
226
227   fflush(stdout);
228   fflush(stderr);
229   BOOL rc = CreateProcess(path.c_str(), command, NULL, NULL, TRUE, 0,
230                           envblock, NULL, &si, &pi);
231   DWORD err = GetLastError();
232
233   // Regardless of whether the process got created or not, we are done with
234   // the handles we created for it to inherit.
235   CloseHandle(si.hStdInput);
236   CloseHandle(si.hStdOutput);
237   CloseHandle(si.hStdError);
238
239   // Now return an error if the process didn't get created.
240   if (!rc)
241   {
242     SetLastError(err);
243     MakeErrMsg(ErrMsg, std::string("Couldn't execute program '") +
244                path.toString() + "'");
245     return -1;
246   }
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 -1;
274     }
275   }
276
277   // Wait for it to terminate.
278   DWORD millisecondsToWait = INFINITE;
279   if (secondsToWait > 0)
280     millisecondsToWait = secondsToWait * 1000;
281
282   if (WaitForSingleObject(pi.hProcess, millisecondsToWait) == WAIT_TIMEOUT) {
283     if (!TerminateProcess(pi.hProcess, 1)) {
284       MakeErrMsg(ErrMsg, std::string("Failed to terminate timed-out program '")
285           + path.toString() + "'");
286       return -1;
287     }
288     WaitForSingleObject(pi.hProcess, INFINITE);
289   }
290
291   // Get its exit status.
292   DWORD status;
293   rc = GetExitCodeProcess(pi.hProcess, &status);
294   err = GetLastError();
295
296   if (!rc) {
297     SetLastError(err);
298     MakeErrMsg(ErrMsg, std::string("Failed getting status for program '") +
299                path.toString() + "'");
300     return -1;
301   }
302
303   return status;
304 }
305
306 bool Program::ChangeStdinToBinary(){
307   int result = _setmode( _fileno(stdin), _O_BINARY );
308   return result == -1;
309 }
310
311 bool Program::ChangeStdoutToBinary(){
312   int result = _setmode( _fileno(stdout), _O_BINARY );
313   return result == -1;
314 }
315
316 }