path: root/ProcessUtilities.win.cc

#include "ProcessUtilities.h"

#include <vector>

#include <windows.h>

class MutexGrabber
{
public:
	MutexGrabber(HANDLE ahMutex)
		: hMutex(ahMutex)
	{
		WaitForSingleObject(hMutex,INFINITE);
	}
	~MutexGrabber()
	{
		ReleaseMutex(hMutex);
	}
private:
	HANDLE hMutex;
};

class ProcessInfo
{
	public:
		ProcessInfo()
			: AccessMutex(NULL), SignalEvent(NULL), StdOutputFileHandle(NULL), StdErrorFileHandle(NULL), Process(NULL),
			ProcessFinished(false)
		{
		}
		~ProcessInfo()
		{
			CloseHandle(AccessMutex);
			CloseHandle(SignalEvent);
			CloseHandle(StdOutputFileHandle);
			CloseHandle(StdErrorFileHandle);
			CloseHandle(Process);
		}

		HANDLE AccessMutex;
		HANDLE SignalEvent;
		HANDLE StdOutputFileHandle;
		HANDLE StdErrorFileHandle;
		HANDLE Process;
		bool ProcessFinished;

		ProcessController* Master;
		std::string OutputBuffer;
		std::string ErrorBuffer;

		void monitorStdOutput();
		void monitorStdError();
};


void ProcessInfo::monitorStdOutput()
{
	while (true)
	{
		char Buffer[1025];
		DWORD ActuallyRead;
		if (!ReadFile(StdOutputFileHandle,Buffer,1024,&ActuallyRead,NULL))
			break;
		MutexGrabber Lock(AccessMutex);
		OutputBuffer += std::string(Buffer,ActuallyRead);
		if (OutputBuffer.length())
			SetEvent(SignalEvent);
	}
	MutexGrabber Lock(AccessMutex);
	CloseHandle(StdOutputFileHandle);
	StdOutputFileHandle = NULL;
	SetEvent(SignalEvent);
}

void ProcessInfo::monitorStdError()
{
	while (true)
	{
		char Buffer[1025];
		DWORD ActuallyRead;
		if (!ReadFile(StdErrorFileHandle,Buffer,1024,&ActuallyRead,NULL))
			break;
		MutexGrabber Lock(AccessMutex);
		ErrorBuffer += std::string(Buffer,ActuallyRead);
		if (ErrorBuffer.length())
			SetEvent(SignalEvent);
	}
	MutexGrabber Lock(AccessMutex);
	CloseHandle(StdErrorFileHandle);
	StdErrorFileHandle = NULL;
	SetEvent(SignalEvent);
}


DWORD WINAPI MonitorStdOutput(LPVOID lpvThreadParam)
{
	ProcessInfo* PI = static_cast<ProcessInfo*>(lpvThreadParam);
	PI->monitorStdOutput();
	return 0;
}

DWORD WINAPI MonitorStdError(LPVOID lpvThreadParam)
{
	ProcessInfo* PI = static_cast<ProcessInfo*>(lpvThreadParam);
	PI->monitorStdError();
	return 0;
}


std::vector<ProcessInfo*> Processes;

bool launchProcess(ProcessController* theController,
	const std::string& Executable, const std::vector<std::string>& Args)
{
/* 	std::cout << Executable;
	for (unsigned int i=0; i<Args.size(); ++i)
		std::cout << " " << Args[i];
	std::cout << std::endl; */


	SECURITY_ATTRIBUTES saAttr; 
	// Set the bInheritHandle flag so pipe handles are inherited. 
	saAttr.nLength = sizeof(SECURITY_ATTRIBUTES); 
	saAttr.bInheritHandle = TRUE; 
	saAttr.lpSecurityDescriptor = NULL; 

	HANDLE StdOutRd, StdOutWr;
	// Create a pipe for the child process's STDOUT. 
	if ( !CreatePipe(&StdOutRd, &StdOutWr, &saAttr, 0) ) 
	{
		std::cerr << "Can't pipe" << std::endl;
		return false;
	}
	// Ensure the read handle to the pipe for STDOUT is not inherited.
	if ( ! SetHandleInformation(StdOutRd, HANDLE_FLAG_INHERIT, 0) )
	{
		std::cerr << "Can't handle handle" << std::endl;
		return false;
	}
	HANDLE StdErrRd, StdErrWr;
	// Create a pipe for the child process's STDErr. 
	if ( !CreatePipe(&StdErrRd, &StdErrWr, &saAttr, 0) ) 
	{
		std::cerr << "Can't pipe" << std::endl;
		return false;
	}
	// Ensure the read handle to the pipe for STDErr is not inherited.
	if ( ! SetHandleInformation(StdErrRd, HANDLE_FLAG_INHERIT, 0) )
	{
		std::cerr << "Can't handle handle" << std::endl;
		return false;
	}

	TCHAR szProcName[10241];
	int x = MultiByteToWideChar(CP_ACP,0,Executable.c_str(),Executable.length(),szProcName,10240);
	szProcName[x] = 0;


	std::string CommandLine("\""+Executable+"\"");
	for (unsigned int i=0; i<Args.size(); ++i)
	{
		CommandLine += " ";
		CommandLine += Args[i];
	}

	TCHAR szCmdline[10241];
	x = MultiByteToWideChar(CP_ACP,0,CommandLine.c_str(),CommandLine.length(),szCmdline,10240);
	szCmdline[x] = 0;
	PROCESS_INFORMATION piProcInfo; 
	STARTUPINFO siStartInfo;
	BOOL bSuccess = FALSE; 
 
	// Set up members of the PROCESS_INFORMATION structure. 
	ZeroMemory( &piProcInfo, sizeof(PROCESS_INFORMATION) );
 
	// Set up members of the STARTUPINFO structure. 
	// This structure specifies the STDIN and STDOUT handles 
	// for redirection.
	ZeroMemory( &siStartInfo, sizeof(STARTUPINFO) );
	siStartInfo.cb = sizeof(STARTUPINFO); 
	siStartInfo.hStdError = StdErrWr;
	siStartInfo.hStdOutput = StdOutWr;
	siStartInfo.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
	siStartInfo.dwFlags |= STARTF_USESTDHANDLES;
 
	// Create the child process. 
	bSuccess = CreateProcess(szProcName, 
		szCmdline,     // command line 
		NULL,          // process security attributes 
		NULL,          // primary thread security attributes 
		TRUE,          // handles are inherited 
		0,             // creation flags 
		NULL,          // use parent's environment 
		NULL,          // use parent's current directory 
		&siStartInfo,  // STARTUPINFO pointer 
		&piProcInfo);  // receives PROCESS_INFORMATION 
   
	// If an error occurs, exit the application. 
	if ( ! bSuccess ) 
	{
		std::cerr << "Can't execute" << std::endl;
		return false;
	}
	else 
	{
		// Close handles to the child process' primary thread.
		CloseHandle(piProcInfo.hThread);
	}
	if (!CloseHandle(StdOutWr))
	{
		std::cerr << "Can't closure" << std::endl;
	}
	if (!CloseHandle(StdErrWr))
	{
		std::cerr << "Can't closure" << std::endl;
	}
	ProcessInfo* PI = new ProcessInfo();
	PI->Process = piProcInfo.hProcess;
	PI->StdOutputFileHandle = StdOutRd;
	PI->StdErrorFileHandle = StdErrRd;
	PI->Master = theController;
	PI->AccessMutex = CreateMutex(NULL,false,NULL);
	if (!PI->AccessMutex)
	{
		std::cerr << "Can't mutex" << std::endl;
		return false;
	}
	PI->SignalEvent = CreateEvent(NULL,FALSE,FALSE,NULL);
	if (!PI->SignalEvent)
	{
		std::cerr << "Can't organize event" << std::endl;
		return false;
	}
	HANDLE StdOutputThread = CreateThread(NULL,0,MonitorStdOutput,PI,0,NULL);
	if (!StdOutputThread)
	{
		std::cerr << "Can't thread a needle" << std::endl;
		return false;
	}
	CloseHandle(StdOutputThread);
	HANDLE StdErrorThread = CreateThread(NULL,0,MonitorStdError,PI,0,NULL);
	if (!StdOutputThread)
	{
		std::cerr << "Can't thread a needle" << std::endl;
		return false;
	}
	CloseHandle(StdErrorThread);
	Processes.push_back(PI);
	//
	return true;
}

unsigned int waitForProcessEvent()
{
	if (!Processes.size()) return 0;
	// TODO wait for some process shit to happen
	std::vector<HANDLE> ReadSignals;
	for (unsigned int i=0; i<Processes.size(); ++i)
		ReadSignals.push_back(Processes[i]->SignalEvent);
	for (unsigned int i=0; i<Processes.size(); ++i)
		if (!Processes[i]->ProcessFinished)
			ReadSignals.push_back(Processes[i]->Process);
	DWORD Res = WaitForMultipleObjects(
		ReadSignals.size(),
		&ReadSignals[0],
		false,
		3000);	
	if (Res == WAIT_TIMEOUT)
	{
		std::cout << "Waiting..." << std::endl;
		return Processes.size();
	}
	else if (Res == WAIT_FAILED)
	{
		std::cout << "WaitForMultipleObjects failed" << std::endl;
		return Processes.size();
	}
	Res -= WAIT_OBJECT_0;
	bool ProcessFinished = false;
	if (Res >= Processes.size())
	{
		Res -= Processes.size();
		for (unsigned int i=0; i<Processes.size(); ++i)
		{
			if (!Processes[i]->ProcessFinished)
			{
				if (Res == 0)
				{
					ProcessFinished = true;
					Res = i;
					break;
				}
				else
					--Res;
			}
		}
	}
	ProcessInfo* PI = Processes[Res];
	std::string Output,Error;
	bool BrokenPipe = false;
	{
		MutexGrabber Lock(PI->AccessMutex);
		if (ProcessFinished)
			PI->ProcessFinished = true;
		Output = PI->OutputBuffer;
		PI->OutputBuffer.clear();
		Error = PI->ErrorBuffer;
		PI->ErrorBuffer.clear();
		BrokenPipe = (PI->StdOutputFileHandle == NULL) && (PI->StdErrorFileHandle == NULL) && PI->ProcessFinished;
	}
	if (Output.length())
		PI->Master->standardOutput(Output);
	if (Error.length())
		PI->Master->standardError(Error);
	if (BrokenPipe)
	{
		DWORD ExitCode = 10;
		GetExitCodeProcess(PI->Process,&ExitCode);
		PI->Master->finished(ExitCode);
		PI->Master->signalFinish(ExitCode);
		if (ExitCode)
			PI->Master->failed();
//		std::cout << "exit code=" << ExitCode << std::endl;
		delete PI;
		Processes.erase(Processes.begin()+Res);
	}
	return Processes.size(); 
}

bool isShellRedirective(const std::string& S)
{
	return (S=="|") || (S=="<") || (S==">");
}

bool shellLaunchProcess(ProcessController* theController,
	const std::string& Executable, const std::vector<std::string>& Args)
{
	unsigned int i=0;
	for (i=0; i<Args.size(); ++i)
		if (isShellRedirective(Args[i]))
			break;
	if (i==Args.size())
		return launchProcess(theController, Executable, Args);
	std::string ShellCommand(Executable);
	for (unsigned int i=0; i<Args.size(); ++i)
	{
		ShellCommand += " ";
		ShellCommand += Args[i];
	}
	std::vector<std::string> NewArgs;
	NewArgs.push_back("/c");
	NewArgs.push_back(ShellCommand);
	return launchProcess(theController, "cmd.exe", NewArgs);
}