mirror of
https://github.com/nillerusr/source-engine.git
synced 2024-12-23 06:36:54 +00:00
474 lines
14 KiB
C++
474 lines
14 KiB
C++
//========= 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;
|
|
}
|
|
|
|
|
|
|