#include "stdafx.h"
#include "ProcessExecute.h"
#include "WcharMbcsConverter.h"
/* The Console Redirection is taken from the TagsView plugin from Vitaliy Dovgan.
* My thanks to him for pointing me in the right direction. :)
* And of course for NppExec, without which, Notepad++ would only
* be half as powerful.
*/
#define DEFAULT_PIPE_SIZE 1
#define PIPE_READBUFSIZE 4096
const char *ProcessExecute::STREAM_NAME_STDOUT = "OUT";
const char *ProcessExecute::STREAM_NAME_STDERR = "ERR";
ProcessExecute::ProcessExecute()
{
}
ProcessExecute::~ProcessExecute()
{
}
bool ProcessExecute::isWindowsNT()
{
OSVERSIONINFO osv;
osv.dwOSVersionInfoSize = sizeof(osv);
::GetVersionEx(&osv);
return (osv.dwPlatformId >= VER_PLATFORM_WIN32_NT);
}
DWORD ProcessExecute::execute(const TCHAR *commandLine, boost::python::object pyStdout, boost::python::object pyStderr, boost::python::object /*pyStdin*/, bool spoolToFile /* = false */)
{
DWORD returnValue = 0;
if (pyStdout.is_none())
throw process_start_exception("stdout cannot be None");
if (pyStderr.is_none())
throw process_start_exception("stderr cannot be None");
// Create out, err, and in pipes (ignore in, initially)
SECURITY_DESCRIPTOR sd;
SECURITY_ATTRIBUTES sa;
process_start_exception exceptionThrown("");
bool thrown = false;
PipeReaderArgs stdoutReaderArgs;
PipeReaderArgs stderrReaderArgs;
// Only used if spooling, but we need to delete it later.
TCHAR tmpFilename[MAX_PATH] = {0};
HANDLE hStdOutReadPipe, hStdOutWritePipe;
HANDLE hStdErrReadPipe, hStdErrWritePipe;
Py_BEGIN_ALLOW_THREADS
try
{
if (isWindowsNT())
{
::InitializeSecurityDescriptor( &sd, SECURITY_DESCRIPTOR_REVISION );
::SetSecurityDescriptorDacl( &sd, TRUE, NULL, FALSE );
sa.lpSecurityDescriptor = &sd;
}
else
{
sa.lpSecurityDescriptor = NULL;
}
sa.nLength = sizeof(SECURITY_ATTRIBUTES);
sa.bInheritHandle = TRUE;
if (!::CreatePipe(&hStdOutReadPipe, &hStdOutWritePipe, &sa, DEFAULT_PIPE_SIZE))
{
throw process_start_exception("Error creating pipe for stdout");
}
if (!::CreatePipe(&hStdErrReadPipe, &hStdErrWritePipe, &sa, DEFAULT_PIPE_SIZE))
{
throw process_start_exception("Error creating pipe for stderr");
}
HANDLE stopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
DWORD dwThreadId;
stdoutReaderArgs.processExecute = this;
stdoutReaderArgs.hPipeRead = hStdOutReadPipe;
stdoutReaderArgs.hPipeWrite = hStdOutWritePipe;
stdoutReaderArgs.pythonFile = pyStdout;
stdoutReaderArgs.stopEvent = stopEvent;
stdoutReaderArgs.completedEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
stdoutReaderArgs.streamName = STREAM_NAME_STDOUT;
stdoutReaderArgs.toFile = false;
stderrReaderArgs.processExecute = this;
stderrReaderArgs.hPipeRead = hStdErrReadPipe;
stderrReaderArgs.hPipeWrite = hStdErrWritePipe;
stderrReaderArgs.pythonFile = pyStderr;
stderrReaderArgs.stopEvent = stopEvent;
stderrReaderArgs.completedEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
stderrReaderArgs.streamName = STREAM_NAME_STDERR;
stderrReaderArgs.toFile = false;
/* If we're in an event, we need to spool the output to a file
* first, before we write it to python - so that the python writing is
* done in *this* thread
*/
if (spoolToFile)
{
stdoutReaderArgs.toFile = true;
stderrReaderArgs.toFile = true;
/// Create the mutex for writing to the file
stdoutReaderArgs.fileMutex = CreateMutex(NULL, FALSE, NULL);
stderrReaderArgs.fileMutex = stdoutReaderArgs.fileMutex;
/// Create the temp file
TCHAR tmpPath[MAX_PATH];
GetTempPath(MAX_PATH, tmpPath);
if(!GetTempFileName(tmpPath, _T("py"), 0, tmpFilename))
{
throw process_start_exception("Error creating temporary filename for output spooling");
}
stdoutReaderArgs.file = new std::fstream(tmpFilename, std::fstream::binary | std::fstream::in | std::fstream::out);
stderrReaderArgs.file = stdoutReaderArgs.file;
/*
stdoutReaderArgs.fileHandle = CreateFile(tmpFilename, GENERIC_READ | GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
stderrReaderArgs.fileHandle = stdoutReaderArgs.fileHandle;
if (INVALID_HANDLE_VALUE == stdoutReaderArgs.fileHandle)
{
throw process_start_exception("Error opening temporary file for output spooling");
}
*/
}
// start thread functions for stdout and stderr
HANDLE hStdoutThread = CreateThread(
NULL, // no security attribute
0, // default stack size
pipeReader, // thread proc
(LPVOID) &stdoutReaderArgs, // thread parameter
0, // not suspended
&dwThreadId); // returns thread ID
HANDLE hStderrThread = CreateThread(
NULL, // no security attribute
0, // default stack size
pipeReader, // thread proc
(LPVOID) &stderrReaderArgs, // thread parameter
0, // not suspended
&dwThreadId); // returns thread ID
// start process
PROCESS_INFORMATION pi;
STARTUPINFO si;
::SetHandleInformation(hStdOutReadPipe, HANDLE_FLAG_INHERIT, 0);
::SetHandleInformation(hStdErrReadPipe, HANDLE_FLAG_INHERIT, 0);
// initialize STARTUPINFO struct
::ZeroMemory( &si, sizeof(STARTUPINFO) );
si.cb = sizeof(STARTUPINFO);
si.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
si.wShowWindow = SW_HIDE;
si.hStdInput = NULL;
si.hStdOutput = hStdOutWritePipe;
si.hStdError = hStdErrWritePipe;
::ZeroMemory( &pi, sizeof(PROCESS_INFORMATION) );
size_t commandLineLength = _tcslen(commandLine) + 1;
// We use an auto_ptr here because of a potential early out due to an exception thrown.
std::auto_ptr<TCHAR> cmdLine(new TCHAR[commandLineLength]);
_tcscpy_s(cmdLine.get(), commandLineLength, commandLine);
bool processStartSuccess;
if ( ::CreateProcess(
NULL,
cmdLine.get(),
NULL, // security
NULL, // security
TRUE, // inherits handles
CREATE_NEW_PROCESS_GROUP, // creation flags
NULL, // environment
NULL, // current directory
&si, // startup info
&pi // process info
))
{
// wait for process to exit
::WaitForSingleObject(pi.hProcess, INFINITE);
CloseHandle(pi.hThread);
processStartSuccess = true;
}
else
{
processStartSuccess = false;
}
// Stop the reader threads
SetEvent(stopEvent);
// Wait for the pipe reader threads to complete
HANDLE handles[] = { stdoutReaderArgs.completedEvent, stderrReaderArgs.completedEvent };
WaitForMultipleObjects(2, handles, TRUE, INFINITE);
CloseHandle(stdoutReaderArgs.completedEvent);
CloseHandle(stderrReaderArgs.completedEvent);
CloseHandle(hStdoutThread);
CloseHandle(hStderrThread);
CloseHandle(stopEvent);
CloseHandle(hStdOutReadPipe);
CloseHandle(hStdOutWritePipe);
CloseHandle(hStdErrReadPipe);
CloseHandle(hStdErrWritePipe);
if (processStartSuccess)
{
GetExitCodeProcess(pi.hProcess, &returnValue);
CloseHandle(pi.hProcess);
}
else
{
DWORD errorNo = ::GetLastError();
TCHAR *buffer;
::FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_ALLOCATE_BUFFER, NULL, errorNo, 0, reinterpret_cast<LPTSTR>(&buffer), 0, NULL);
std::shared_ptr<char> message = WcharMbcsConverter::tchar2char(buffer);
process_start_exception ex(message.get());
::LocalFree(buffer);
throw ex;
}
}
catch(process_start_exception& ex)
{
exceptionThrown = ex;
thrown = true;
}
Py_END_ALLOW_THREADS
if (spoolToFile)
{
CloseHandle(stdoutReaderArgs.fileMutex);
spoolFile(stdoutReaderArgs.file, pyStdout, pyStderr);
stdoutReaderArgs.file->close();
if (tmpFilename[0] != 0)
{
DeleteFile(tmpFilename);
}
}
if (thrown)
{
throw exceptionThrown;
}
return returnValue;
}
DWORD WINAPI ProcessExecute::pipeReader(void *args)
{
PipeReaderArgs* pipeReaderArgs = reinterpret_cast<PipeReaderArgs*>(args);
DWORD bytesRead;
char buffer[PIPE_READBUFSIZE];
BOOL processFinished = FALSE;
for(;;)
{
::PeekNamedPipe(pipeReaderArgs->hPipeRead, NULL, 0, NULL, &bytesRead, NULL);
if (processFinished && 0 == bytesRead)
{
break;
}
if (bytesRead > 0)
{
if (ReadFile(pipeReaderArgs->hPipeRead, buffer, PIPE_READBUFSIZE - 1, &bytesRead, NULL))
{
if (pipeReaderArgs->toFile)
{
pipeReaderArgs->processExecute->writeToFile(pipeReaderArgs, bytesRead, buffer);
}
else
{
pipeReaderArgs->processExecute->writeToPython(pipeReaderArgs, bytesRead, buffer);
}
}
}
else
{
DWORD handleIndex = WaitForSingleObject(pipeReaderArgs->stopEvent, 100);
if (WAIT_OBJECT_0 == handleIndex)
{
processFinished = TRUE;
}
}
}
SetEvent(pipeReaderArgs->completedEvent);
return 0;
}
void ProcessExecute::writeToPython(PipeReaderArgs *pipeReaderArgs, DWORD bytesRead, char *buffer)
{
buffer[bytesRead] = '\0';
PyGILState_STATE gstate = PyGILState_Ensure();
try
{
pipeReaderArgs->pythonFile.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
}
catch(...)
{
PyErr_Print();
}
PyGILState_Release(gstate);
}
void ProcessExecute::writeToFile(PipeReaderArgs *pipeReaderArgs, DWORD bytesRead, char *buffer)
{
WaitForSingleObject(pipeReaderArgs->fileMutex, INFINITE);
pipeReaderArgs->file->write(pipeReaderArgs->streamName, ProcessExecute::STREAM_NAME_LENGTH);
if (pipeReaderArgs->file->bad())
throw process_start_exception("Error writing to spool file");
char byteCount[20];
sprintf_s(byteCount, 20, "%ul\n", bytesRead);
pipeReaderArgs->file->write(byteCount, strlen(byteCount));
pipeReaderArgs->file->write(buffer, bytesRead);
if (pipeReaderArgs->file->bad())
throw process_start_exception("Error writing buffer to spool file");
ReleaseMutex(pipeReaderArgs->fileMutex);
}
void ProcessExecute::spoolFile(std::fstream* file, boost::python::object pyStdout, boost::python::object pyStderr)
{
// Rewind
file->seekg(0);
char infoBuffer[30];
char *buffer = NULL;
size_t bufferSize = 0;
size_t bytesToRead;
while (!file->eof())
{
file->getline(infoBuffer, 30, '\n');
if (file->eof())
break;
bytesToRead = strtoul(infoBuffer + STREAM_NAME_LENGTH, NULL, 0);
if (bufferSize < bytesToRead || !buffer)
{
if (buffer)
delete[] buffer;
bufferSize = bytesToRead + 1;
buffer = new char[bufferSize];
}
file->read(buffer, bytesToRead);
buffer[bytesToRead] = '\0';
if (0 == strncmp(infoBuffer, STREAM_NAME_STDOUT, STREAM_NAME_LENGTH))
{
// Is a stdout message
pyStdout.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
}
else
{
// is a stderr message
pyStderr.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
}
}
if (buffer)
{
delete[] buffer;
}
}
