Subversion Repositories DevTools

Rev

Blame | Last modification | View Log | RSS feed

//==============================================================================
//
//     Copyright : (C) 2008 Videlli Limited
//     This file belongs to Videlli Limited
//
//==============================================================================
//
//     Purpose :
//     This file contains the win32 "buildtool attend" service source code.
//     This is a buildtool helper service, designed to "attend" to rogue processes
//     that insist on presenting a dialog on the console of unattended build machine
//     through termination.
//     This has the effect of holding a build thread to ransom.
//     The rules of engagement are:
//     1 a configured rogue process is the descendant of a configured ancestor process (the build daemon win32 service).
//     2 a configured rogue process has been running longer than a configured maximum run time.
//     Only when 1 and 2 are true will the rogue process be terminated.
//     The service is designed to be as lightweight as possible, checking for such a
//     scenario once every 60 seconds, or responding to a stop or shutdown request immediately.
//     The rogue tools are limited to a suite of borland tools at the time of writing,
//     though this list is configurable should the need arise to extend the list.
//     They are configured in the following registry key:
//     HKEY_LOCAL_MACHINE\\SYSTEM\\CurrentControlSet\\Services\\buildtool attend\\Parameters
//     AppParameters String Value.
//     This takes the form:
//     <fully qualified ancestor path> <max run time secs> <rogue child 1 process name> <rogue child 2 process name> ...
//     For example:
//     D:\\Work\\buildtool\\srvany.exe 60 bcc32.exe bpr2mak.exe brc32.exe brcc32.exe coff2omf.exe dcc32.exe ilink32.exe impdef.exe implib.exe make.exe tlib.exe tlibimp.exe
//     
//==============================================================================

#include <windows.h>
#include <iostream.h>
#include <vector>
#include <sstream>
#include "Psapi.h"
#include "Tlhelp32.h"
#include "Process.h"

void WINAPI ServiceMain( DWORD argc, LPTSTR * argv );
void WINAPI ControlHandler( DWORD dwControl );
void Start( DWORD argc, LPSTR * argv );
void Stop();
HANDLE hWaitEvent = 0L;
SERVICE_STATUS_HANDLE stat_hand = {0};
SERVICE_STATUS servstat = {0};
char *serviceName = "buildtool attend\0";

//------------------------------------------------------------------------------
//  log
//  Logs a message to the Application Event Log
//------------------------------------------------------------------------------
//  Parameters:
//      message IN - the message to log
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void log( const char * message )
{
        HANDLE hEventLog = RegisterEventSource( 0L, "Application" );

        if ( hEventLog != 0L )
        {
                ReportEvent(  hEventLog, EVENTLOG_INFORMATION_TYPE, 0, 0, 0L, 1, 0, &message, 0L );
                DeregisterEventSource( hEventLog );
        }
}

//------------------------------------------------------------------------------
//  isDescendedFrom
//  Determines if a child process is descended from an ancestor process
//------------------------------------------------------------------------------
//  Parameters:
//      ancestor IN - the ancestor process
//      child IN - the child process
//      processCollection IN - the collection of running processes
//------------------------------------------------------------------------------
//  Returns:
//      true if the child is descended from the ancestor
//------------------------------------------------------------------------------
bool isDescendedFrom( const std::string &ancestor, const Process &child, const std::vector< Process * > processCollection )
{
        bool retVal = false;

        if ( child.mPath == ancestor )
        {
                // child is descended from ancestor
                retVal = true;
        }
        else
        {
                // find the childs parent
                Process *parent = 0L;
                std::vector< Process * >::const_iterator processCollectionIterator = processCollection.begin();

                while ( processCollectionIterator != processCollection.end() )
                {
                        Process *p = *processCollectionIterator;
                        if ( child.mPpid == p->mPid  )
                        {
                                parent = p;
                                break;
                        }
                
                        processCollectionIterator ++;
                }

                if ( parent != 0L )
                {
                        retVal = isDescendedFrom( ancestor, *parent, processCollection );
                }
        }

        return retVal;
}

//------------------------------------------------------------------------------
//  main
//  This is the entry point for the process.
//  It sets up a service table entry and then calls StartServiceCtrlDispatcher.
//  This call does not return, but uses the thread which called it as a
//  control dispatcher for all the services implemented by this process.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments on the command line
//      argv IN - the arguments on the command line
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void __cdecl main( DWORD argc, LPSTR * argv )
{
        // start the control dispatcher
        // this call gives the SCManager this thread for the entire period that this service is running,
        // so that it can call us back with service controls
        // it will spawn a new thread to run the service itself, starting at entrypoint ServiceMain
        SERVICE_TABLE_ENTRY stab[] = {
                { serviceName, ServiceMain },
                { NULL, NULL }
        };

        StartServiceCtrlDispatcher( stab );
}

//------------------------------------------------------------------------------
//  ServiceMain
//  This is the main entry point for the service.
//  The service control dispatcher creates a thread to start at this routine.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments on the command line
//      argv IN - the arguments on the command line
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void WINAPI ServiceMain( DWORD argc, LPTSTR * argv )
{
        // register a service control handler for the service
        // this will be called by the control dispatcher when it has instructions for the service
        stat_hand = RegisterServiceCtrlHandler( serviceName, ControlHandler );

        if ( stat_hand == 0 )
        {
                std::string message = "ServiceMain RegisterServiceCrtlHandler failed with last error " + GetLastError();
                log( message.c_str() );
        }

        // create the handle used by the ControlHandler to signal to this method to return
        hWaitEvent = CreateEvent( 0L, FALSE, FALSE, 0L );

        if ( hWaitEvent == 0L )
        {
                std::string message = "ServiceMain CreateEvent failed with last error " + GetLastError();
                log( message.c_str() );
        }

        // inform the service control manager that the service is running
        servstat.dwServiceType = SERVICE_WIN32;
        servstat.dwCurrentState = SERVICE_RUNNING;
        servstat.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
        servstat.dwWin32ExitCode = NO_ERROR;
        servstat.dwServiceSpecificExitCode = 0;
        servstat.dwCheckPoint = 0;
        servstat.dwWaitHint = 0;

        if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
        {
                std::string message = "ServiceMain SetServiceStatus 1 failed with last error " + GetLastError();
                log( message.c_str() );
        }

        HKEY hKey;
        char appParams[255];
        DWORD dwBufLen;

        RegOpenKeyEx( HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\buildtool attend\\Parameters", 0, KEY_QUERY_VALUE, &hKey );
        RegQueryValueEx( hKey, "AppParameters", 0L, 0L, (LPBYTE) appParams, &dwBufLen);
        RegCloseKey( hKey );

        // convert appParams to args
        // calculate ac
        int ac = 0;

        // increment ac if a non space is read and either no previous character existed or previous character was space
        for ( int i = 0; i < strlen( appParams ); i++ )
        {
                if ( !isspace( appParams[ i ] ) && ( i == 0 || isspace( appParams[ i - 1] ) ) )
                {
                        ac++;
                }
        }

        char ** args = new char * [ ac ];
        int arg = 0;

        for ( i = 0; i < strlen( appParams ); i++ )
        {
                if ( !isspace( appParams[ i ] ) && ( i == 0 || isspace( appParams[ i - 1] ) ) )
                {
                        args[ arg ] = &appParams[ i ];
                        arg++;
                }
        }

        int length = strlen( appParams );

        for ( i = 0; i < length; i++ )
        {
                if ( isspace( appParams[ i ] ) )
                {
                        appParams[ i ] = '\0';
                }
        }

        for ( arg = 0; arg < ac; arg++ )
        {
                std::string message = "ServiceMain args[ ";
                std::string param;
                std::stringstream ss;
                ss << arg;
                ss >> param;
                message += param;
                message += "] = ";
                message += args[ arg ];
                log( message.c_str() );
        }

        // start the service
        Start( ac, args );

        // Start has returned because of a command
        // stop the service
        Stop();

        // inform the service control manager that the service is stopped
        servstat.dwCurrentState = SERVICE_STOPPED;

        if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
        {
                std::string message = "ServiceMain SetServiceStatus 2 failed with last error " + GetLastError();
                log( message.c_str() );
        }

        return;
}

//------------------------------------------------------------------------------
//  ControlHandler
//  This is called by the control dispatcher when it has instructions for the service
//------------------------------------------------------------------------------
//  Parameters:
//      dwControl IN - the instruction
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void WINAPI ControlHandler( DWORD dwControl )
{
        switch( dwControl )
        {
        case SERVICE_CONTROL_PAUSE:
        case SERVICE_CONTROL_CONTINUE:
        case SERVICE_CONTROL_INTERROGATE:
                servstat.dwWaitHint = 0;
                break;
        case SERVICE_CONTROL_SHUTDOWN:
        case SERVICE_CONTROL_STOP:
                servstat.dwCurrentState = SERVICE_STOP_PENDING;
                servstat.dwWaitHint = 0;
                break;
        }

        // inform the service control manager of the service state
        if ( SetServiceStatus( stat_hand, &servstat ) == 0 )
        {
                std::string message = "ControlHandler SetServiceStatus failed with last error " + GetLastError();
                log( message.c_str() );
        }

        if ( ( dwControl == SERVICE_CONTROL_STOP ) || ( dwControl == SERVICE_CONTROL_SHUTDOWN ) )
        {
                log( "ControlHandler detected stop request" );

                // signal ServiceMain to return 
                if ( SetEvent( hWaitEvent ) == 0 )
                {
                        std::string message = "ControlHandler SetEvent failed with last error " + GetLastError();
                        log( message.c_str() );
                }

        }
}

//------------------------------------------------------------------------------
//  Start
//  The service logic, run when the service starts.
//------------------------------------------------------------------------------
//  Parameters:
//      argc IN - the number of arguments configured in the registry
//      argv IN - the arguments configured in the registry
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void Start( DWORD argc, LPSTR * argv )
{
        boolean doSomethingUseful = true;

        if ( argc < 3 )
        {
                log( "Start 1 configuration issue" );
                doSomethingUseful = false;
        }

        int maxRunTimeSecs = atoi( argv[ 1 ] );

        if ( maxRunTimeSecs == 0 )
        {
                log( "Start 2 configuration issue" );
                doSomethingUseful = false;
        }

        if ( doSomethingUseful )
        {
                std::vector< Process * > processCollection;
                std::string ancestor = argv[ 0 ];
                PROCESSENTRY32 entry;
                std::vector< Process * >::iterator processCollectionIterator = processCollection.begin();
                DWORD run = WAIT_TIMEOUT;

                // set the SeDebugPrivilege to prevent OpenProcess failing with access is denied
                LUID luid;

                if ( LookupPrivilegeValue( 0L, SE_DEBUG_NAME, &luid ) )
                {
                        HANDLE hToken;
                        if ( OpenProcessToken( GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken ) )
                        {
                                TOKEN_PRIVILEGES tp;
                                tp.PrivilegeCount = 1;
                                tp.Privileges[ 0 ].Luid = luid;
                                tp.Privileges[ 0 ].Attributes = SE_PRIVILEGE_ENABLED;
                                AdjustTokenPrivileges( hToken, FALSE, &tp, sizeof( TOKEN_PRIVILEGES ), 0L, 0L );
                                CloseHandle( hToken );
                        }

                }

                do
                {
                        // clear processCollection
                        processCollectionIterator = processCollection.begin();

                        while ( processCollectionIterator != processCollection.end() )
                        {
                                Process *p = *processCollectionIterator;
                                delete p;
                                p = 0L;
                                processCollectionIterator ++;
                        }

                        processCollection.clear();

                        HANDLE snapshot = CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 );

                        if ( snapshot )
                        {
                                memset( &entry, 0, sizeof( entry ) );
                                entry.dwSize = sizeof( entry );

                                if ( Process32First( snapshot, &entry ) == TRUE )
                                {
                                        do
                                        {
                                                // get a handle to the process
                                                HANDLE hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID );

                                                if ( hProcess )
                                                {
                                                        // get processes fully qualified path
                                                        HMODULE hMod;
                                                        DWORD cbNeeded = 0;

                                                        char path[MAX_PATH] = "";
                                                        if ( EnumProcessModules( hProcess, &hMod, sizeof( hMod ), &cbNeeded ) )
                                                        {
                                                                FILETIME creationTime;
                                                                FILETIME exitTime;
                                                                FILETIME kernelTime;
                                                                FILETIME userTime;
                                                                GetModuleFileNameEx( hProcess, hMod, path, sizeof( path) );
                                                                GetProcessTimes( hProcess, &creationTime, &exitTime, &kernelTime, &userTime );
                                                                ULARGE_INTEGER nano;
                                                                nano.LowPart = creationTime.dwLowDateTime;
                                                                nano.HighPart = creationTime.dwHighDateTime;
                                                                Process *p = new Process( path, entry.szExeFile, entry.th32ProcessID, entry.th32ParentProcessID, nano.QuadPart );
                                                                processCollection.push_back( p );
                                                        }

                                                        CloseHandle( hProcess );
                                                }
                                                
                                        } while ( Process32Next( snapshot, &entry ) );
                                }
                                CloseHandle( snapshot );
                        }

                        // iterate through the processCollection
                        // if its name matches a child from argv[2] to argv[argc-1], and it is a child of ancestor, and its runtime has expired, then terminate the process
                        processCollectionIterator = processCollection.begin();

                        while ( processCollectionIterator != processCollection.end() )
                        {
                                Process *p = *processCollectionIterator;
                                int child = 2;

                                while( child < argc )
                                {
                                        if ( p->mName == argv[ child ] )
                                        {
                                                std::string message = "Start found configured process running ";
                                                message += p->mName;
                                                log( message.c_str() );
                                                // its name matches a child
                                                if ( isDescendedFrom( ancestor, *p, processCollection ) )
                                                {
                                                        std::string message = "Start found configured process running and descendant ";
                                                        message += p->mName;
                                                        log( message.c_str() );
                                                        // it is a child, directly or indirectly, of ancestor
                                                        SYSTEMTIME st;
                                                        FILETIME ft;
                                                        // get current time
                                                        GetSystemTime(&st);
                                                        // convert to file time format
                                                        SystemTimeToFileTime(&st, &ft);
                                                        ULARGE_INTEGER nano;
                                                        nano.LowPart = ft.dwLowDateTime;
                                                        nano.HighPart = ft.dwHighDateTime;
                                                        ULARGE_INTEGER runTime;
                                                        runTime.QuadPart = nano.QuadPart - p->mTime;

                                                        int runTimeSecs = runTime.QuadPart/10000000;

                                                        if ( runTimeSecs > maxRunTimeSecs )
                                                        {
                                                                std::string message = "Start found configured process running too long ";
                                                                message += p->mName;
                                                                log( message.c_str() );
                                                                // it has been running for too long, terminate
                                                                HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_TERMINATE, FALSE, p->mPid );

                                                                if ( hProcess )
                                                                {
                                                                        std::string message = "Start terminating ";
                                                                        message += p->mName;
                                                                        log( message.c_str() );
                                                                        TerminateProcess( hProcess, CONTROL_C_EXIT );
                                                                        CloseHandle( hProcess );
                                                                }

                                                        }

                                                }
                                                break;
                                        }
                                        child++;
                                }
                        
                                processCollectionIterator ++;
                        }

                        // this is effectively the ServiceMain routine
                        // do not exit this loop and return until signalled
                        // wait for 60 secs and continue to run if not signalled
                        run = WaitForSingleObject( hWaitEvent, 60000 );

                        if ( run == WAIT_FAILED )
                        {
                                std::string message = "Start WaitForSingleObject 1 failed with last error " + GetLastError();
                                log( message.c_str() );
                        }

                } while ( run == WAIT_TIMEOUT );

                // clear processCollection
                processCollectionIterator = processCollection.begin();

                while ( processCollectionIterator != processCollection.end() )
                {
                        Process *p = *processCollectionIterator;
                        delete p;
                        p = 0L;
                        processCollectionIterator ++;
                }

                processCollection.clear();

        }
        else
        {
                log( "Start calling WaitForSingleObject" );
                // this is effectively the ServiceMain routine
                // do not return until signalled
                if ( WaitForSingleObject( hWaitEvent, INFINITE ) == WAIT_FAILED )
                {
                        std::string message = "Start WaitForSingleObject 2 failed with last error " + GetLastError();
                        log( message.c_str() );
                }
                log( "Start called WaitForSingleObject" );
        }

}

//------------------------------------------------------------------------------
//  Stop
//  The service clean up, run when the service is to be stopped.
//------------------------------------------------------------------------------
//  Parameters:
//------------------------------------------------------------------------------
//  Returns:
//      None
//------------------------------------------------------------------------------
void Stop()
{
        // nothing to do
}