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(const char* str);

/*----------------------------------------------------------------------------

** 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;

//        #define JOB_OBJECT_LIMIT_BREAKAWAY_OK 0x00000800

        JOBOBJECT_BASIC_LIMIT_INFORMATION basicLimits = {0, 0, JOB_OBJECT_LIMIT_BREAKAWAY_OK};

        JOBOBJECT_EXTENDED_LIMIT_INFORMATION extendedLimits = {basicLimits};

        /*

        **  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");

    /*

    **  Set extended information to allow job breakaway

    **  Required so that we can have a timeout job within a timeout job

    */

    if (!SetInformationJobObject(hJobObject, JobObjectExtendedLimitInformation, &extendedLimits, sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)))

        ErrorExitMsg(EXIT_CANCELED, "Error setting job object information.");

    /*

    ** 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;

        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.

    */

    #define LOCAL_CREATE_BREAKAWAY_FROM_JOB 0x01000000

    #define CREATION_FLAGS LOCAL_CREATE_BREAKAWAY_FROM_JOB | CREATE_NEW_PROCESS_GROUP 

    if (! CreateProcess(

            NULL,                       // Use from command line

            command,                    // command line

            NULL,                       // process security attributes

            NULL,                       // primary thread security attributes 

            TRUE,                       // handles are inherited 

            CREATION_FLAGS ,            // 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;

}