1HEADmaster

1 files changed, 473 insertions, 0 deletions

diff --git a/vstdlib/processutils.cpp b/vstdlib/processutils.cpp
new file mode 100644
index 0000000..dc4ee59
--- /dev/null
+++ b/vstdlib/processutils.cpp
@@ -0,0 +1,473 @@
+//========= Copyright Valve Corporation, All rights reserved. ============//
+//
+// Purpose: 
+//
+//===========================================================================//
+
+#if !defined( _X360 )
+#include <windows.h>
+#endif
+#include "vstdlib/iprocessutils.h"
+#include "tier1/utllinkedlist.h"
+#include "tier1/utlstring.h"
+#include "tier1/utlbuffer.h"
+#include "tier1/tier1.h"
+
+//-----------------------------------------------------------------------------
+// At the moment, we can only run one process at a time 
+//-----------------------------------------------------------------------------
+class CProcessUtils : public CTier1AppSystem< IProcessUtils >
+{
+	typedef CTier1AppSystem< IProcessUtils > BaseClass;
+
+public:
+	CProcessUtils() : BaseClass( false ) {}
+
+	// Inherited from IAppSystem
+	virtual InitReturnVal_t Init();
+	virtual void Shutdown();
+
+	// Inherited from IProcessUtils
+	virtual ProcessHandle_t StartProcess( const char *pCommandLine, bool bConnectStdPipes );
+	virtual ProcessHandle_t StartProcess( int argc, const char **argv, bool bConnectStdPipes );
+	virtual void CloseProcess( ProcessHandle_t hProcess );
+	virtual void AbortProcess( ProcessHandle_t hProcess );
+	virtual bool IsProcessComplete( ProcessHandle_t hProcess );
+	virtual void WaitUntilProcessCompletes( ProcessHandle_t hProcess );
+	virtual int SendProcessInput( ProcessHandle_t hProcess, char *pBuf, int nBufLen );
+	virtual int GetProcessOutputSize( ProcessHandle_t hProcess );
+	virtual int GetProcessOutput( ProcessHandle_t hProcess, char *pBuf, int nBufLen );
+	virtual int GetProcessExitCode( ProcessHandle_t hProcess );
+
+private:
+	struct ProcessInfo_t
+	{
+		HANDLE m_hChildStdinRd;
+		HANDLE m_hChildStdinWr;
+		HANDLE m_hChildStdoutRd;
+		HANDLE m_hChildStdoutWr;
+		HANDLE m_hChildStderrWr;
+		HANDLE m_hProcess;
+		CUtlString m_CommandLine;
+		CUtlBuffer m_ProcessOutput;
+	};
+
+	// Returns the last error that occurred
+	char *GetErrorString( char *pBuf, int nBufLen );
+
+	// creates the process, adds it to the list and writes the windows HANDLE into info.m_hProcess
+	ProcessHandle_t CreateProcess( ProcessInfo_t &info, bool bConnectStdPipes );
+
+	// Shuts down the process handle
+	void ShutdownProcess( ProcessHandle_t hProcess );
+
+	// Methods used to read	output back from a process
+	int GetActualProcessOutputSize( ProcessHandle_t hProcess );
+	int GetActualProcessOutput( ProcessHandle_t hProcess, char *pBuf, int nBufLen );
+
+	CUtlFixedLinkedList< ProcessInfo_t >	m_Processes;
+	ProcessHandle_t m_hCurrentProcess;
+	bool m_bInitialized;
+};
+
+
+//-----------------------------------------------------------------------------
+// Purpose: singleton accessor
+//-----------------------------------------------------------------------------
+static CProcessUtils s_ProcessUtils;
+EXPOSE_SINGLE_INTERFACE_GLOBALVAR( CProcessUtils, IProcessUtils, PROCESS_UTILS_INTERFACE_VERSION, s_ProcessUtils );
+
+
+//-----------------------------------------------------------------------------
+// Initialize, shutdown process system
+//-----------------------------------------------------------------------------
+InitReturnVal_t CProcessUtils::Init()
+{
+	InitReturnVal_t nRetVal = BaseClass::Init();
+	if ( nRetVal != INIT_OK )
+		return nRetVal;
+
+	m_bInitialized = true;
+	m_hCurrentProcess = PROCESS_HANDLE_INVALID;
+	return INIT_OK;
+}
+
+void CProcessUtils::Shutdown()
+{
+	Assert( m_bInitialized );
+	Assert( m_Processes.Count() == 0 );
+	if ( m_Processes.Count() != 0 )
+	{
+		AbortProcess( m_hCurrentProcess );
+	}
+	m_bInitialized = false;
+	return BaseClass::Shutdown();
+}
+
+
+//-----------------------------------------------------------------------------
+// Returns the last error that occurred
+//-----------------------------------------------------------------------------
+char *CProcessUtils::GetErrorString( char *pBuf, int nBufLen )
+{
+	FormatMessage( FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), 0, pBuf, nBufLen, NULL );
+	char *p = strchr(pBuf, '\r');	// get rid of \r\n
+	if(p) 
+	{
+		p[0] = 0;
+	}
+	return pBuf;
+}
+
+
+ProcessHandle_t CProcessUtils::CreateProcess( ProcessInfo_t &info, bool bConnectStdPipes )
+{
+	STARTUPINFO si;
+	memset(&si, 0, sizeof si);
+	si.cb = sizeof(si);
+	if ( bConnectStdPipes )
+	{
+		si.dwFlags = STARTF_USESTDHANDLES;
+		si.hStdInput = info.m_hChildStdinRd;
+		si.hStdError = info.m_hChildStderrWr;
+		si.hStdOutput = info.m_hChildStdoutWr;
+	}
+
+	PROCESS_INFORMATION pi;
+	if ( ::CreateProcess( NULL, info.m_CommandLine.GetForModify(), NULL, NULL, TRUE, DETACHED_PROCESS, NULL, NULL, &si, &pi ) )
+	{
+		info.m_hProcess = pi.hProcess;
+		m_hCurrentProcess = m_Processes.AddToTail( info );
+		return m_hCurrentProcess;
+	}
+
+	char buf[ 512 ];
+	Warning( "Could not execute the command:\n   %s\n"
+		"Windows gave the error message:\n   \"%s\"\n",
+		info.m_CommandLine.Get(), GetErrorString( buf, sizeof(buf) ) );
+
+	return PROCESS_HANDLE_INVALID;
+}
+
+//-----------------------------------------------------------------------------
+// Options for compilation
+//-----------------------------------------------------------------------------
+ProcessHandle_t CProcessUtils::StartProcess( const char *pCommandLine, bool bConnectStdPipes )
+{
+	Assert( m_bInitialized );
+
+	// NOTE: For the moment, we can only run one process at a time
+	// although in the future, I expect to have a process queue.
+	if ( m_hCurrentProcess != PROCESS_HANDLE_INVALID )
+	{
+		WaitUntilProcessCompletes( m_hCurrentProcess ); 
+	}
+
+	ProcessInfo_t info;
+	info.m_CommandLine = pCommandLine;
+
+	if ( !bConnectStdPipes )
+	{
+		info.m_hChildStderrWr = INVALID_HANDLE_VALUE;
+		info.m_hChildStdinRd = info.m_hChildStdinWr = INVALID_HANDLE_VALUE;
+		info.m_hChildStdoutRd = info.m_hChildStdoutWr = INVALID_HANDLE_VALUE;
+
+		return CreateProcess( info, false );
+	}
+
+    SECURITY_ATTRIBUTES saAttr; 
+
+    // Set the bInheritHandle flag so pipe handles are inherited.
+	saAttr.nLength = sizeof(SECURITY_ATTRIBUTES); 
+    saAttr.bInheritHandle = TRUE; 
+    saAttr.lpSecurityDescriptor = NULL; 
+ 
+    // Create a pipe for the child's STDOUT. 
+    if ( CreatePipe( &info.m_hChildStdoutRd, &info.m_hChildStdoutWr, &saAttr, 0 ) )
+	{
+		if ( CreatePipe( &info.m_hChildStdinRd, &info.m_hChildStdinWr, &saAttr, 0 ) )
+		{
+			if ( DuplicateHandle( GetCurrentProcess(), info.m_hChildStdoutWr, GetCurrentProcess(), 
+				&info.m_hChildStderrWr, 0, TRUE, DUPLICATE_SAME_ACCESS ) )
+			{
+//				_setmode( info.m_hChildStdoutRd, _O_TEXT );
+//				_setmode( info.m_hChildStdoutWr, _O_TEXT );
+//				_setmode( info.m_hChildStderrWr, _O_TEXT );
+
+				ProcessHandle_t hProcess = CreateProcess( info, true );
+				if ( hProcess != PROCESS_HANDLE_INVALID )
+					return hProcess;
+
+				CloseHandle( info.m_hChildStderrWr );
+			}
+			CloseHandle( info.m_hChildStdinRd );
+			CloseHandle( info.m_hChildStdinWr );
+		}
+		CloseHandle( info.m_hChildStdoutRd );
+		CloseHandle( info.m_hChildStdoutWr );
+	}
+	return PROCESS_HANDLE_INVALID;
+}
+
+
+//-----------------------------------------------------------------------------
+// Start up a process
+//-----------------------------------------------------------------------------
+ProcessHandle_t CProcessUtils::StartProcess( int argc, const char **argv, bool bConnectStdPipes )
+{
+	CUtlString commandLine;
+	for ( int i = 0; i < argc; ++i )
+	{
+		commandLine += argv[i];
+		if ( i != argc-1 )
+		{
+			commandLine += " ";
+		}
+	}
+	return StartProcess( commandLine.Get(), bConnectStdPipes );
+}
+
+
+//-----------------------------------------------------------------------------
+// Shuts down the process handle
+//-----------------------------------------------------------------------------
+void CProcessUtils::ShutdownProcess( ProcessHandle_t hProcess )
+{
+	ProcessInfo_t& info = m_Processes[hProcess];
+	CloseHandle( info.m_hChildStderrWr );
+	CloseHandle( info.m_hChildStdinRd );
+	CloseHandle( info.m_hChildStdinWr );
+	CloseHandle( info.m_hChildStdoutRd );
+	CloseHandle( info.m_hChildStdoutWr );
+
+	m_Processes.Remove( hProcess );
+}
+
+
+//-----------------------------------------------------------------------------
+// Closes the process
+//-----------------------------------------------------------------------------
+void CProcessUtils::CloseProcess( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+	if ( hProcess != PROCESS_HANDLE_INVALID )
+	{
+		WaitUntilProcessCompletes( hProcess );
+		ShutdownProcess( hProcess );
+	}
+}
+
+
+//-----------------------------------------------------------------------------
+// Aborts the process
+//-----------------------------------------------------------------------------
+void CProcessUtils::AbortProcess( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+	if ( hProcess != PROCESS_HANDLE_INVALID )
+	{
+		if ( !IsProcessComplete( hProcess ) )
+		{
+			ProcessInfo_t& info = m_Processes[hProcess];
+			TerminateProcess( info.m_hProcess, 1 );
+		}
+		ShutdownProcess( hProcess );
+	}
+}
+
+
+//-----------------------------------------------------------------------------
+// Returns true if the process is complete
+//-----------------------------------------------------------------------------
+bool CProcessUtils::IsProcessComplete( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+	Assert( hProcess != PROCESS_HANDLE_INVALID );
+	if ( m_hCurrentProcess != hProcess )
+		return true;
+
+	HANDLE h = m_Processes[hProcess].m_hProcess;
+	return ( WaitForSingleObject( h, 0 ) != WAIT_TIMEOUT );
+}
+
+
+//-----------------------------------------------------------------------------
+// Methods used to write input into a process
+//-----------------------------------------------------------------------------
+int CProcessUtils::SendProcessInput( ProcessHandle_t hProcess, char *pBuf, int nBufLen )
+{
+	// Unimplemented yet
+	Assert( 0 );
+	return 0;
+}
+
+
+//-----------------------------------------------------------------------------
+// Methods used to read	output back from a process
+//-----------------------------------------------------------------------------
+int CProcessUtils::GetActualProcessOutputSize( ProcessHandle_t hProcess )
+{
+	Assert( hProcess != PROCESS_HANDLE_INVALID );
+
+	ProcessInfo_t& info = m_Processes[ hProcess ];
+	if ( info.m_hChildStdoutRd == INVALID_HANDLE_VALUE )
+		return 0;
+
+	DWORD dwCount = 0;
+	if ( !PeekNamedPipe( info.m_hChildStdoutRd, NULL, NULL, NULL, &dwCount, NULL ) )
+	{
+		char buf[ 512 ];
+		Warning( "Could not read from pipe associated with command %s\n"
+			"Windows gave the error message:\n   \"%s\"\n",
+			info.m_CommandLine.Get(), GetErrorString( buf, sizeof(buf) ) );
+		return 0;
+	}
+
+	// Add 1 for auto-NULL termination
+	return ( dwCount > 0 ) ? (int)dwCount + 1 : 0;
+}
+
+int CProcessUtils::GetActualProcessOutput( ProcessHandle_t hProcess, char *pBuf, int nBufLen )
+{
+	ProcessInfo_t& info = m_Processes[ hProcess ];
+	if ( info.m_hChildStdoutRd == INVALID_HANDLE_VALUE )
+		return 0;
+
+	DWORD dwCount = 0;
+	DWORD dwRead = 0;
+
+	// FIXME: Is there a way of making pipes be text mode so we don't get /n/rs back?
+	char *pTempBuf = (char*)_alloca( nBufLen );
+	if ( !PeekNamedPipe( info.m_hChildStdoutRd, NULL, NULL, NULL, &dwCount, NULL ) )
+	{
+		char buf[ 512 ];
+		Warning( "Could not read from pipe associated with command %s\n"
+			"Windows gave the error message:\n   \"%s\"\n",
+			info.m_CommandLine.Get(), GetErrorString( buf, sizeof(buf) ) );
+		return 0;
+	}
+
+	dwCount = min( dwCount, (DWORD)nBufLen - 1 );
+	ReadFile( info.m_hChildStdoutRd, pTempBuf, dwCount, &dwRead, NULL);
+	
+	// Convert /n/r -> /n
+	int nActualCountRead = 0;
+	for ( unsigned int i = 0; i < dwRead; ++i )
+	{
+		char c = pTempBuf[i];
+		if ( c == '\r' )
+		{
+			if ( ( i+1 < dwRead ) && ( pTempBuf[i+1] == '\n' ) )
+			{
+				pBuf[nActualCountRead++] = '\n';
+				++i;
+				continue;
+			}
+		}
+
+		pBuf[nActualCountRead++] = c;
+	}
+
+	return nActualCountRead;
+}
+
+
+int CProcessUtils::GetProcessOutputSize( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+	if ( hProcess == PROCESS_HANDLE_INVALID )
+		return 0;
+
+	return GetActualProcessOutputSize( hProcess ) + m_Processes[hProcess].m_ProcessOutput.TellPut();
+}
+
+
+int CProcessUtils::GetProcessOutput( ProcessHandle_t hProcess, char *pBuf, int nBufLen )
+{
+	Assert( m_bInitialized );
+
+	if ( hProcess == PROCESS_HANDLE_INVALID )
+		return 0;
+
+	ProcessInfo_t &info = m_Processes[hProcess];
+	int nCachedBytes = info.m_ProcessOutput.TellPut();
+	int nBytesRead = 0;
+	if ( nCachedBytes )
+	{
+		nBytesRead = min( nBufLen-1, nCachedBytes );
+		info.m_ProcessOutput.Get( pBuf, nBytesRead );
+		pBuf[ nBytesRead ] = 0;
+		nBufLen -= nBytesRead;
+		pBuf += nBytesRead;
+		if ( info.m_ProcessOutput.GetBytesRemaining() == 0 )
+		{
+			info.m_ProcessOutput.Purge();
+		}
+
+		if ( nBufLen <= 1 )
+			return nBytesRead;
+	}
+
+	// Auto-NULL terminate
+	int nActualCountRead = GetActualProcessOutput( hProcess, pBuf, nBufLen );
+	pBuf[nActualCountRead] = 0;
+	return nActualCountRead + nBytesRead + 1;
+}
+
+
+//-----------------------------------------------------------------------------
+// Returns the exit code for the process. Doesn't work unless the process is complete
+//-----------------------------------------------------------------------------
+int CProcessUtils::GetProcessExitCode( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+	ProcessInfo_t &info = m_Processes[hProcess];
+	DWORD nExitCode;
+	BOOL bOk = GetExitCodeProcess( info.m_hProcess, &nExitCode );
+	if ( !bOk || nExitCode == STILL_ACTIVE )
+		return -1;
+	return nExitCode;
+}
+
+
+//-----------------------------------------------------------------------------
+// Waits until a process is complete
+//-----------------------------------------------------------------------------
+void CProcessUtils::WaitUntilProcessCompletes( ProcessHandle_t hProcess )
+{
+	Assert( m_bInitialized );
+
+	// For the moment, we can only run one process at a time
+	if ( ( hProcess == PROCESS_HANDLE_INVALID ) || ( m_hCurrentProcess != hProcess ) )
+		return;
+
+	ProcessInfo_t &info = m_Processes[ hProcess ];
+
+	if ( info.m_hChildStdoutRd == INVALID_HANDLE_VALUE )
+	{
+		WaitForSingleObject( info.m_hProcess, INFINITE );
+	}
+	else
+	{
+		// NOTE: The called process can block during writes to stderr + stdout
+		// if the pipe buffer is empty. Therefore, waiting INFINITE is not
+		// possible here. We must queue up messages received to allow the 
+		// process to continue
+		while ( WaitForSingleObject( info.m_hProcess, 100 ) == WAIT_TIMEOUT )
+		{
+			int nLen = GetActualProcessOutputSize( hProcess );
+			if ( nLen > 0 )
+			{
+				int nPut = info.m_ProcessOutput.TellPut();
+				info.m_ProcessOutput.EnsureCapacity( nPut + nLen );
+				int nBytesRead = GetActualProcessOutput( hProcess, (char*)info.m_ProcessOutput.PeekPut(), nLen );
+				info.m_ProcessOutput.SeekPut( CUtlBuffer::SEEK_HEAD, nPut + nBytesRead );
+			}
+		}
+	}
+
+	m_hCurrentProcess = PROCESS_HANDLE_INVALID;
+}
+
+
+