Beef/BeefTools/RunAndWait/RunAndWait.cpp

#define _CRT_SECURE_NO_WARNINGS
#include <windows.h>
#include <string>
#include <psapi.h>
#include <shlobj.h>
#include <algorithm>
#include <stdio.h>
#include <time.h>

static bool CreatePipeWithSecurityAttributes(HANDLE& hReadPipe, HANDLE& hWritePipe, SECURITY_ATTRIBUTES* lpPipeAttributes, int nSize)
{
	hReadPipe = 0;
	hWritePipe = 0;
	bool ret = ::CreatePipe(&hReadPipe, &hWritePipe, lpPipeAttributes, nSize);
	if (!ret || (hReadPipe == INVALID_HANDLE_VALUE) || (hWritePipe == INVALID_HANDLE_VALUE))
		return false;
	return true;
}

// Using synchronous Anonymous pipes for process input/output redirection means we would end up 
// wasting a worker threadpool thread per pipe instance. Overlapped pipe IO is desirable, since 
// it will take advantage of the NT IO completion port infrastructure. But we can't really use 
// Overlapped I/O for process input/output as it would break Console apps (managed Console class 
// methods such as WriteLine as well as native CRT functions like printf) which are making an
// assumption that the console standard handles (obtained via GetStdHandle()) are opened
// for synchronous I/O and hence they can work fine with ReadFile/WriteFile synchrnously!
bool CreatePipe(HANDLE& parentHandle, HANDLE& childHandle, bool parentInputs)
{
	SECURITY_ATTRIBUTES securityAttributesParent = { 0 };
	securityAttributesParent.bInheritHandle = 1;

	HANDLE hTmp = INVALID_HANDLE_VALUE;
	if (parentInputs)
		CreatePipeWithSecurityAttributes(childHandle, hTmp, &securityAttributesParent, 0);
	else
		CreatePipeWithSecurityAttributes(hTmp, childHandle, &securityAttributesParent, 0);

	HANDLE dupHandle = 0;

	// Duplicate the parent handle to be non-inheritable so that the child process 
	// doesn't have access. This is done for correctness sake, exact reason is unclear.
	// One potential theory is that child process can do something brain dead like 
	// closing the parent end of the pipe and there by getting into a blocking situation
	// as parent will not be draining the pipe at the other end anymore. 
	if (!::DuplicateHandle(GetCurrentProcess(), hTmp,
		GetCurrentProcess(), &dupHandle,
		0, false, DUPLICATE_SAME_ACCESS))
	{
		return false;
	}

	parentHandle = dupHandle;

	if (hTmp != INVALID_HANDLE_VALUE)
		::CloseHandle(hTmp);

	return true;
}

struct ProcParams
{
	HANDLE mReadHandle;
	HANDLE mWriteHandle;
};

DWORD WINAPI ReadProc(void* lpThreadParameter)
{
	ProcParams* procParams = (ProcParams*)lpThreadParameter;
	while (true)
	{
		char buffer[2048];
		DWORD bytesRead = 0;
		if (::ReadFile(procParams->mReadHandle, buffer, (DWORD)2048, &bytesRead, NULL))
		{
			DWORD bytesWritten;
			::WriteFile(procParams->mWriteHandle, buffer, bytesRead, &bytesWritten, NULL);
		}
		else
		{
			int err = GetLastError();
			break;
		}
	}

	return 0;
}

int main()
{	
	char* cmdLineStr = ::GetCommandLineA();
	
	DWORD flags = CREATE_DEFAULT_ERROR_MODE;
	void* envPtr = NULL;
		
	char* useCmdLineStr = cmdLineStr;

	if (cmdLineStr[0] != 0)
	{
		bool nameQuoted = cmdLineStr[0] == '\"';

		std::string passedName;
		int i;
		for (i = (nameQuoted ? 1 : 0); cmdLineStr[i] != 0; i++)
		{
			wchar_t c = cmdLineStr[i];

			if (((nameQuoted) && (c == '"')) ||
				((!nameQuoted) && (c == ' ')))
			{
				i++;
				break;
			}
			passedName += cmdLineStr[i];
		}

		useCmdLineStr += i;
		while (*useCmdLineStr == L' ')
			useCmdLineStr++;
	}
	
	std::string cmdLine = useCmdLineStr;	
	
	PROCESS_INFORMATION processInfo;

	STARTUPINFOA si;
	ZeroMemory(&si, sizeof(si));
	si.cb = sizeof(si);
	memset(&processInfo, 0, sizeof(processInfo));

	si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);

 	HANDLE stdOut;
 	CreatePipe(stdOut, si.hStdOutput, false);
	
	HANDLE stdErr;
	CreatePipe(stdErr, si.hStdError, false);	

	si.dwFlags = STARTF_USESTDHANDLES;

	DWORD startTick = GetTickCount();
	BOOL worked = CreateProcessA(NULL, (char*)cmdLine.c_str(), NULL, NULL, TRUE,
		flags, envPtr, NULL, &si, &processInfo);

	::CloseHandle(si.hStdOutput);
	::CloseHandle(si.hStdError);

	if (!worked)
		return 1;

	DWORD threadId;
 	ProcParams stdOutParams = { stdOut, GetStdHandle(STD_OUTPUT_HANDLE) };
 	HANDLE stdOutThread = ::CreateThread(NULL, (SIZE_T)128*1024, (LPTHREAD_START_ROUTINE)ReadProc, (void*)&stdOutParams, 0, &threadId);

	ProcParams stdErrParams = { stdErr, GetStdHandle(STD_ERROR_HANDLE) };
	HANDLE stdErrThread = ::CreateThread(NULL, (SIZE_T)128 * 1024, (LPTHREAD_START_ROUTINE)ReadProc, (void*)&stdErrParams, 0, &threadId);

	while (true)
	{		
		if (::WaitForSingleObject(processInfo.hProcess, 20) == WAIT_OBJECT_0)
			break;
	}
	::WaitForSingleObject(stdOutThread, INFINITE);
	::WaitForSingleObject(stdErrThread, INFINITE);

	DWORD exitCode = 0;
	::GetExitCodeProcess(processInfo.hProcess, &exitCode);	
	
	return exitCode;
}