Subversion Repositories DevTools

/*============================================================================
**  Copyright (c) VIX TECHNOLOGY (AUST) LTD
**============================================================================
**
**  Project/Product : 
**  Filename        : timeout_win32.c
**  Author(s)       : DDP
**
**  Description     : Windows program to limit the wall-clock time that an program can run
** 
**                      Create an other process
**                      Kill it after it has run for a specified length of time
**
**                    Used to help JATS to:
**                      Limit compilation times
**                          Handle silly compilers that pop up error messages with an OK
**                      Limit Unit Test run times
**
** Usage            : timeout <options>* command <command args>*
**                    Where options are:
**                      -time:nnn[.nnn][s]  - Timeout in seconds
**                                            Floating point number with suffix
**                                            Suffix - s - seconds
**                                                     m - minutes
**                                                     h - hours
**                                                     d - days
**                    Command               - Users command
**                    command args          - Arguments passed to the command
**
** Note:            This code MUST be kept in sync with code for all platforms supported
**                  under JATS
** 
**  Information     :
**   Compiler       : ANSI C
**   Target         : WIN32 ONLY
**
**==========================================================================*/

#define _CRT_SECURE_NO_WARNINGS 1       // Supress warnings about use of strcpy()
#define _WIN32_WINNT 0x0500             // Access Job Control Functions

#include <stdio.h>
#include <windows.h>
#include <limits.h>

/*
** Exit Codes
*/
#define     EXIT_TIMEDOUT      124      /* job timed out*/
#define     EXIT_CANCELED      125      /* internal error*/
#define     EXIT_CANNOT_INVOKE 126      /* error executing job */
#define     EXIT_ENOENT        127      /* couldn't find job to exec */

/*
**  Globals
*/
 
CRITICAL_SECTION CriticalSection;
HANDLE hJobObject;
DWORD dwThreadIdTimer;
HANDLE hThreadTimer;
DWORD   dwChildExitCode;
HANDLE  hChildProcess;
unsigned int timeOut = 10;                   // Seconds

/*
**  Prototypes (forward declarations)
*/
HANDLE CreateChildProcess(LPTSTR);
VOID ErrorExit(LPTSTR);
void ErrorExitMsg(UINT ecode, LPTSTR lpszFunction) ;
DWORD WINAPI ThreadTimer( LPVOID lpParam );
BOOL WINAPI KillJob(DWORD dwCtrlType);
unsigned int parseDuration( char *ptr);

/*----------------------------------------------------------------------------
** FUNCTION           : main
**
** DESCRIPTION        : Mainentry point
**
**
** INPUTS             : argc            - Number of args
**                      argv            - Address of args
**
** RETURNS            :
**
----------------------------------------------------------------------------*/

int main(int argc, char *argv[])
{
   SECURITY_ATTRIBUTES saAttr; 
   LPTSTR   command;
   int      argp;

    /*
    **  Process command line arguments
    **  Options for this program will be located first
    */
    for ( argp = 1; argp < argc ; argp++ )
    {
        /*
        **  Scan until first non option argument
        */
        if ( *argv[argp] != '-' )
            break;

        if ( _strnicmp( argv[argp], "-TIME:", 6 ) == 0)
        {
            timeOut = parseDuration(argv[argp] + 6);
        }
        else
        {
            fprintf(stderr, "Unknown option: %s", argv[argp] );
            ErrorExit( "Bad Option");
        }
    }

    /*
    **  Need at least one more argument - the program to run
    */
    if ( argp >= argc )
    {
      ErrorExit("Insufficient arguments");
    }

    /*
    **  Process the remainder of the command line and create a string
    **  For the command line. Some of the arguments will need to be quoted
    */
    {
        size_t length = 0;
        int start;
        char join = 0;

        /*
        **  Calc required size of the command line
        **  Allow for two " and 1 space per argument
        */
        for ( start = argp; start < argc ; start++ )
        {
            length += 3 + strlen ( argv[start] );
        }
        command = (LPTSTR) malloc ( length + 1);
        if ( !command )
            ErrorExit("Cannot allocate memory");
        *command = 0;

        for ( start = argp; start < argc ; start++ )
        {
            char *quote =  strrchr(argv[start], ' ' );

            if ( join )
                strcat (command, " " );
            join = 1;

            if ( quote )
                strcat (command, "\"" );
            strcat (command, argv[start] );
            if ( quote )
                strcat (command, "\"" );
        }
    }

    /*
    ** Create a JOB object to contain our job
    */
    hJobObject = CreateJobObject(NULL, NULL);
    if (! hJobObject)
            ErrorExitMsg(EXIT_CANCELED, "CreateJobObject");

    /*
    ** Register a handler for control-C
    */
    if ( !SetConsoleCtrlHandler( KillJob, TRUE))
    {
        ErrorExitMsg(EXIT_CANCELED, "SetConsoleCtrlHandler");
    }
  

    /*
    **  Create a security attribute
    **  Set the bInheritHandle flag so pipe handles are inherited
    */
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE; 
    saAttr.lpSecurityDescriptor = NULL; 
 
    /*
    **  Now create the child process.
    */
    hChildProcess = CreateChildProcess(command);
    if (! hChildProcess)
        ErrorExitMsg(EXIT_CANNOT_INVOKE,"CreateChildProcess");
    if (!AssignProcessToJobObject(hJobObject, hChildProcess)) {
        TerminateProcess(hChildProcess, 101);
        ErrorExitMsg(EXIT_ENOENT, "AssignProcessToJobObject");
    };


    /*
    **  Create a critical section to be used by pipe writers
    */
    InitializeCriticalSection(&CriticalSection);

    /*
    **  Create a thread to timeout the main process
    */
    hThreadTimer = CreateThread(
        NULL,                        // default security attributes 
        0,                           // use default stack size  
        ThreadTimer,                 // thread function
        NULL,                        // argument to thread function
        0,                           // use default creation flags 
        &dwThreadIdTimer);           // returns the thread identifier

    if (hThreadTimer == NULL)
        ErrorExitMsg(EXIT_CANCELED, "Create Timer Thread");

    /*
    **  What for the process to complete
    */
    WaitForSingleObject( hChildProcess, INFINITE );
    
    /*
    ** If the one process I started terminates, then kill the entire JOB
    */
    TerminateThread(hThreadTimer, 0);
    TerminateJobObject(hJobObject, 100);
    

    /*
    **  Determine the exist status of the child process
    */
    if (!GetExitCodeProcess( hChildProcess, &dwChildExitCode ))
        ErrorExit("Error getting child exist code");

    /*
    **  Release resources used by the critical section object.
    */
    DeleteCriticalSection(&CriticalSection);

    /*
    **  Exist with the childs exist code
    */
    ExitProcess(dwChildExitCode);

    /*
    ** Should not happen
    */
    return EXIT_FAILURE;
}

/*----------------------------------------------------------------------------
** FUNCTION           : apply_time_suffix
**
** DESCRIPTION        : Helper function to process time suffix
**                      Supports s,h,m,d
**
** INPUTS             : x           - Ref to value to modify
**                      suffix_char - Possible suffix char
**
** RETURNS            : True - suffix was OK
**
----------------------------------------------------------------------------*/

static int apply_time_suffix (double *x, char suffix_char)
{
    int multiplier;

    switch (suffix_char)
    {
    case 0:
    case 's':
        multiplier = 1;
        break;
    case 'm':
        multiplier = 60;"Invalid duration\n"
        break;
    case 'h':
        multiplier = 60 * 60;
        break;
    case 'd':
        multiplier = 60 * 60 * 24;
        break;
    default:
        return 0;
    }
    *x *= multiplier;
    return 1;
}

/*----------------------------------------------------------------------------
** FUNCTION           : parseDuration 
**
** DESCRIPTION        : Parse the timeout duration
**                      Allowed - number + 1char suffix
**                          1 Char suffix may be:
**                          'm' - Minutes
**                          'h' - Hours
**                          'd' - Days
**                      The number is an unsigned integer
**
** INPUTS             : ptr     - Input string (Floating point)
**
** RETURNS            : Unsigned integer
**                      May not return. May exit with error message
**
----------------------------------------------------------------------------*/

static unsigned int parseDuration (const char* str)
{
    double duration;
    unsigned int intDuration;
    unsigned int duration_floor;
    char *ep;
    int ok = 1;

    errno = 0;
    duration = strtod(str, &ep);
    if (duration < 0 )
        ok = 0;
    if ( ep == str )
        ok = 0;
    if (errno != 0)
        ok = 0;
    if (*ep && *(ep + 1))
        ok = 0;

    if(ok)
        ok = apply_time_suffix (&duration, *ep);

    if (UINT_MAX <= duration)
        ok = 0;

    if(!ok)
    {
        ErrorExit( "Invalid duration");
    }

    /*
    ** Round up - unless exact
    */
    duration_floor = (unsigned int)duration;
    intDuration = duration_floor + (duration_floor < duration);

    /*fprintf(stderr, "parseDuration: %f (%u)\n", duration,  intDuration);*/

  return intDuration;
}

/*----------------------------------------------------------------------------
** FUNCTION           : CreateChildProcess
**
** DESCRIPTION        : Create the child process
**                      It will be created in its own process group , so that all members
**                      of the same process can be killed.
**
** INPUTS             : command         - Command string
**
** RETURNS            : False if can't create
**
----------------------------------------------------------------------------*/

HANDLE CreateChildProcess( LPTSTR command)
{ 
   PROCESS_INFORMATION piProcInfo; 
   STARTUPINFO siStartInfo; 
 
    /*
    **  Set up members of STARTUPINFO structure.
    */
    ZeroMemory( &siStartInfo, sizeof(STARTUPINFO) );
    siStartInfo.cb = sizeof(STARTUPINFO);
 
    /*
    **  Create the child process.
    */
    if (! CreateProcess(
            NULL,                       // Use from command line
            command,                    // command line
            NULL,                       // process security attributes
            NULL,                       // primary thread security attributes 
            TRUE,                       // handles are inherited 
            CREATE_NEW_PROCESS_GROUP ,  // creation flags 
            NULL,                       // use parent's environment 
            NULL,                       // use parent's current directory 
            &siStartInfo,               // STARTUPINFO pointer 
            &piProcInfo                 // receives PROCESS_INFORMATION
            )) return (NULL);

    return piProcInfo.hProcess;
}


/*----------------------------------------------------------------------------
** FUNCTION           : ErrorExit
**                      ErrorExitMsg - report system messae
**
** DESCRIPTION        : Error processing
**                      REport an error and terminate process
**
**
** INPUTS             : lpszMessage     - Message to display
**
** RETURNS            : Does't return
**                      Will exit with bad code
**
----------------------------------------------------------------------------*/

VOID ErrorExit (LPTSTR lpszMessage)
{ 
   fprintf(stderr, "%s\n", lpszMessage);
   fflush(stderr) ;
   ExitProcess(EXIT_CANCELED);
}

void ErrorExitMsg(UINT ecode, LPTSTR lpszFunction) 
{ 
    // Retrieve the system error message for the last-error code

    LPVOID lpMsgBuf;
    DWORD dw = GetLastError(); 

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | 
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dw,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR) &lpMsgBuf,
        0, NULL );

    // Display the error message and exit the process

    fprintf(stderr, "%s failed with error %d: %s\n", lpszFunction, dw, lpMsgBuf);
    fflush(stderr) ;
    LocalFree(lpMsgBuf);

    ExitProcess(ecode); 
}

/*----------------------------------------------------------------------------
** FUNCTION           : ThreadTimer
**
** DESCRIPTION        : Thread to timeout the created process
**                      Done on a thread so that we can sleep
**
** INPUTS             : lpParam                 - Not used
**
** RETURNS            : 0
**
----------------------------------------------------------------------------*/

DWORD WINAPI ThreadTimer( LPVOID lpParam )
{
    char lBuf[30];
    unsigned int secs;

    _snprintf(lBuf, sizeof(lBuf), "%u", timeOut);
    SetEnvironmentVariable( "GBE_JOBLIMIT", lBuf );

    for (secs = timeOut; secs > 0; secs--) 
    { 
        
        /*
        ** Update ENV with time remaining
        ** Simply for diagnostic purposes
        */
        _snprintf(lBuf, sizeof(lBuf), "%u", secs);
        SetEnvironmentVariable( "GBE_JOBTIME", lBuf );

        /*
        **  Wait one second
        */
        Sleep(1000);
    }

    fprintf(stderr, "Process exceeded time limit (%u Secs)\n", timeOut);
    fflush(stderr) ;

    TerminateJobObject(hJobObject, EXIT_TIMEDOUT);
    return 0;
}

/*----------------------------------------------------------------------------
** FUNCTION           : KillJob
**
** DESCRIPTION        : Kill the one job controlled by this program
**
** INPUTS             : dwCtrlType                 - Not used
**
** RETURNS            : TRUE - Signal has been processed
**
----------------------------------------------------------------------------*/

BOOL WINAPI KillJob(DWORD dwCtrlType)
{
    fprintf(stderr, "Process killed with Control-C\n");
    fflush(stderr) ;

    TerminateJobObject(hJobObject, EXIT_TIMEDOUT);
    return TRUE;
}