Implement memoryLimit on Windows.
[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 was developed by Jeff Cohen and is distributed under the 
6 // University of Illinois Open Source 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 = path->toString().c_str();
82   if (*fname == 0)
83     fname = "NUL";
84
85   SECURITY_ATTRIBUTES sa;
86   sa.nLength = sizeof(sa);
87   sa.lpSecurityDescriptor = 0;
88   sa.bInheritHandle = TRUE;
89
90   h = CreateFile(fname, fd ? GENERIC_WRITE : GENERIC_READ, FILE_SHARE_READ,
91                  &sa, fd == 0 ? OPEN_EXISTING : CREATE_ALWAYS,
92                  FILE_ATTRIBUTE_NORMAL, NULL);
93   if (h == INVALID_HANDLE_VALUE) {
94     MakeErrMsg(ErrMsg, std::string(fname) + ": Can't open file for " +
95         (fd ? "input: " : "output: "));
96   }
97
98   return h;
99 }
100
101 int 
102 Program::ExecuteAndWait(const Path& path, 
103                         const char** args,
104                         const char** envp,
105                         const Path** redirects,
106                         unsigned secondsToWait,
107                         unsigned memoryLimit,
108                         std::string* ErrMsg) {
109   if (!path.canExecute()) {
110     if (ErrMsg)
111       *ErrMsg = "program not executable";
112     return -1;
113   }
114
115   // Windows wants a command line, not an array of args, to pass to the new
116   // process.  We have to concatenate them all, while quoting the args that
117   // have embedded spaces.
118
119   // First, determine the length of the command line.
120   unsigned len = 0;
121   for (unsigned i = 0; args[i]; i++) {
122     len += strlen(args[i]) + 1;
123     if (strchr(args[i], ' '))
124       len += 2;
125   }
126
127   // Now build the command line.
128   char *command = reinterpret_cast<char *>(_alloca(len));
129   char *p = command;
130
131   for (unsigned i = 0; args[i]; i++) {
132     const char *arg = args[i];
133     size_t len = strlen(arg);
134     bool needsQuoting = strchr(arg, ' ') != 0;
135     if (needsQuoting)
136       *p++ = '"';
137     memcpy(p, arg, len);
138     p += len;
139     if (needsQuoting)
140       *p++ = '"';
141     *p++ = ' ';
142   }
143
144   *p = 0;
145
146   // Create a child process.
147   STARTUPINFO si;
148   memset(&si, 0, sizeof(si));
149   si.cb = sizeof(si);
150   si.hStdInput = INVALID_HANDLE_VALUE;
151   si.hStdOutput = INVALID_HANDLE_VALUE;
152   si.hStdError = INVALID_HANDLE_VALUE;
153
154   if (redirects) {
155     si.dwFlags = STARTF_USESTDHANDLES;
156     
157     si.hStdInput = RedirectIO(redirects[0], 0, ErrMsg);
158     if (si.hStdInput == INVALID_HANDLE_VALUE) {
159       MakeErrMsg(ErrMsg, "can't redirect stdin");
160       return -1;
161     }
162     si.hStdOutput = RedirectIO(redirects[1], 1, ErrMsg);
163     if (si.hStdOutput == INVALID_HANDLE_VALUE) {
164       CloseHandle(si.hStdInput);
165       MakeErrMsg(ErrMsg, "can't redirect stdout");
166       return -1;
167     }
168     if (redirects[1] && redirects[2] && *(redirects[1]) != *(redirects[2])) {
169       si.hStdError = RedirectIO(redirects[2], 2, ErrMsg);
170       if (si.hStdError == INVALID_HANDLE_VALUE) {
171         CloseHandle(si.hStdInput);
172         CloseHandle(si.hStdOutput);
173         MakeErrMsg(ErrMsg, "can't redirect stderr");
174         return -1;
175       }
176     } else {
177       DuplicateHandle(GetCurrentProcess(), si.hStdOutput,
178                       GetCurrentProcess(), &si.hStdError,
179                       0, TRUE, DUPLICATE_SAME_ACCESS);
180     }
181   }
182   
183   PROCESS_INFORMATION pi;
184   memset(&pi, 0, sizeof(pi));
185
186   fflush(stdout);
187   fflush(stderr);
188   BOOL rc = CreateProcess(path.c_str(), command, NULL, NULL, FALSE, 0,
189                           envp, NULL, &si, &pi);
190   DWORD err = GetLastError();
191
192   // Regardless of whether the process got created or not, we are done with
193   // the handles we created for it to inherit.
194   CloseHandle(si.hStdInput);
195   CloseHandle(si.hStdOutput);
196   CloseHandle(si.hStdError);
197
198   // Now return an error if the process didn't get created.
199   if (!rc)
200   {
201     SetLastError(err);
202     MakeErrMsg(ErrMsg, std::string("Couldn't execute program '") + 
203                path.toString() + "'");
204     return -1;
205   }
206
207   // Make sure these get closed no matter what.
208   AutoHandle hProcess(pi.hProcess);
209   AutoHandle hThread(pi.hThread);
210
211   // Assign the process to a job if a memory limit is defined.
212   AutoHandle hJob(0);
213   if (memoryLimit != 0) {
214     hJob = CreateJobObject(0, 0);
215     bool success = false;
216     if (hJob != 0) {
217       JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli;
218       memset(&jeli, 0, sizeof(jeli));
219       jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
220       jeli.ProcessMemoryLimit = memoryLimit * 1048576;
221       if (SetInformationJobObject(hJob, JobObjectExtendedLimitInformation,
222                                   &jeli, sizeof(jeli))) {
223         if (AssignProcessToJobObject(hJob, pi.hProcess))
224           success = true;
225       }
226     }
227     if (!success) {
228       SetLastError(GetLastError());
229       MakeErrMsg(ErrMsg, std::string("Unable to set memory limit"));
230       TerminateProcess(pi.hProcess, 1);
231       WaitForSingleObject(pi.hProcess, INFINITE);
232       return -1;
233     }
234   }
235
236   // Wait for it to terminate.
237   DWORD millisecondsToWait = INFINITE;
238   if (secondsToWait > 0)
239     millisecondsToWait = secondsToWait * 1000;
240
241   if (WaitForSingleObject(pi.hProcess, millisecondsToWait) == WAIT_TIMEOUT) {
242     if (!TerminateProcess(pi.hProcess, 1)) {
243       MakeErrMsg(ErrMsg, std::string("Failed to terminate timed-out program '")
244           + path.toString() + "'");
245       return -1;
246     }
247     WaitForSingleObject(pi.hProcess, INFINITE);
248   }
249   
250   // Get its exit status.
251   DWORD status;
252   rc = GetExitCodeProcess(pi.hProcess, &status);
253   err = GetLastError();
254
255   if (!rc) {
256     SetLastError(err);
257     MakeErrMsg(ErrMsg, std::string("Failed getting status for program '") + 
258                path.toString() + "'");
259     return -1;
260   }
261
262   return status;
263 }
264
265 bool Program::ChangeStdinToBinary(){
266   int result = _setmode( _fileno(stdin), _O_BINARY );
267   return result == -1;
268 }
269
270 bool Program::ChangeStdoutToBinary(){
271   int result = _setmode( _fileno(stdout), _O_BINARY );
272   return result == -1;
273 }
274
275 }