hl2_src-leak-2017/src/vstdlib/processutils.cpp

474 lines
14 KiB
C++
Raw Permalink Normal View History

//========= 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;
}