/*============================================================================
**  COPYRIGHT - VIX IP PTY LTD ("VIX"). ALL RIGHTS RESERVED.
**============================================================================
**
**  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
*/
 
HANDLE          hJobObject;             // Job to contain the child process
DWORD           dwChildExitCode;        // Child process exit code
unsigned int    timeOut = 10;           // Seconds

/*
**  Prototypes (forward declarations)
*/
VOID ErrorExit(LPTSTR);
void ErrorExitMsg(UINT ecode, LPTSTR lpszFunction) ;
BOOL WINAPI KillJob(DWORD dwCtrlType);
unsigned int parseDuration(const char* str);

/*----------------------------------------------------------------------------
** FUNCTION           : main
**
** DESCRIPTION        : Main entry point
**                      timeout [-time:nnn] program ......
**
**
** INPUTS             : argc            - Number of args
**                      argv            - Address of args
**
** RETURNS            : 0   - Program exited OK
**                      ... - Program exit code
**                      124 - Timeout
**                      125 - Internal error
**                      126 - Could not find program
**                      127 - Could not create Job
**
----------------------------------------------------------------------------*/

int main(int argc, char *argv[])
{
   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");

    /*
    **  Set extended information to allow job breakaway
    **  Required so that we can have a timeout job within a timeout job
    **  Note: Not sure this works to well on all systems (Windows-7)
    */
     JOBOBJECT_EXTENDED_LIMIT_INFORMATION JBLI;
     if(!QueryInformationJobObject(hJobObject, JobObjectExtendedLimitInformation, &JBLI, sizeof(JBLI), NULL)){
        ErrorExitMsg(EXIT_CANCELED, "QueryInformationJobObject");
     }

     JBLI.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_BREAKAWAY_OK;
     if(!SetInformationJobObject(hJobObject, JobObjectExtendedLimitInformation, &JBLI, sizeof(JBLI))) {
        ErrorExitMsg(EXIT_CANCELED, "SetInformationJobObject");      
     }

    /*
    ** Register a handler for control-C
    */
    if ( !SetConsoleCtrlHandler( KillJob, TRUE)) {
        ErrorExitMsg(EXIT_CANCELED, "SetConsoleCtrlHandler");
    }
  
    /*
    **  Create the child process.
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "CreateChild" );
    PROCESS_INFORMATION piProcInfo;
    
    /*
    **  Set up members of STARTUPINFO structure.
    */
    STARTUPINFO siStartInfo; 
    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_SUSPENDED | CREATE_NEW_PROCESS_GROUP | CREATE_BREAKAWAY_FROM_JOB,  // creation flags 
            //CREATE_BREAKAWAY_FROM_JOB | CREATE_NEW_PROCESS_GROUP ,  // creation flags 
            NULL,                       // use parent's environment 
            NULL,                       // use parent's current directory 
            &siStartInfo,               // STARTUPINFO pointer 
            &piProcInfo                 // receives PROCESS_INFORMATION
            )) {
        ErrorExitMsg(EXIT_CANNOT_INVOKE,"CreateChildProcess");
    }

    if (!AssignProcessToJobObject(hJobObject, piProcInfo.hProcess)) {
        TerminateProcess(piProcInfo.hProcess, 101);
        ErrorExitMsg(EXIT_ENOENT, "AssignProcessToJobObject");
    };

    /*
    **  Start(Resume) the master thread for the comamnd
    **  What for the process to complete
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "ResumeChild" );
    if (ResumeThread (piProcInfo.hThread) == -1 ){
         ErrorExitMsg(EXIT_ENOENT, "ResumeThread - did not resume");
    }

    /*
    ** Body of the wating process
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "WaitChild" );
    {
        char lBuf[30];
        unsigned int secs;
        DWORD dwWaitResult;
        bool  bChildDone = false;

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

        /*
        ** Wait for the Job to complete, one second at a time 
        ** This allows us to indicate progress via an EnvVar 
        */
        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
            */
            dwWaitResult  = WaitForSingleObject( piProcInfo.hProcess, 1000 );
            switch (dwWaitResult) 
            {
                /* 
                **  Event object was signaled
                **      The process has completed 
                */ 
                case WAIT_OBJECT_0: 
                    SetEnvironmentVariable( "GBE_JOBDEBUG", "ChildDone" );
                    bChildDone = true;
                    break;
                    
                /* 
                **  Timeout
                **      Update EnvVar and do again 
                */ 
                case WAIT_TIMEOUT:
                    break;

                /* 
                **  An error occurred 
                */ 
                default: 
                    ErrorExitMsg(EXIT_CANCELED, "WaitForSingleObject");
                    break;
            }
        }

        /*
        ** Timer loop complete 
        ** May be a timeout or the Child has completed 
        */
        if (!bChildDone)
        {
            /*
            ** Timeout 
            **  Terminate the Job 
            */
            fprintf(stderr, "Process exceeded time limit (%u Secs)\n", timeOut); 
            fflush(stderr) ;

            SetEnvironmentVariable( "GBE_JOBDEBUG", "TimeOut" );
            TerminateJobObject(hJobObject, EXIT_TIMEDOUT);
        }
    }
    
    SetEnvironmentVariable( "GBE_JOBDEBUG", "TimerDone" );
    TerminateJobObject(hJobObject, 100);

    /*
    **  Determine the exist status of the child process
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "GetChildExitCode" );
    if (!GetExitCodeProcess( piProcInfo.hProcess, &dwChildExitCode ))
        ErrorExit("Error getting child exist code");

    /*
    **  Exit with the childs exit code
    **      Have tried to use ExitProcess, but it hangs from time to time
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "Exit" );
    ExitProcess(dwChildExitCode);

    /*
    ** Should not happen
    */
    SetEnvironmentVariable( "GBE_JOBDEBUG", "ExitReturn" );
    return EXIT_CANCELED;
}

/*----------------------------------------------------------------------------
** 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           : 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)
{ 
   SetEnvironmentVariable( "GBE_JOBDEBUG", 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();
    char msgBuf[300]; 

    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

    _snprintf(msgBuf, sizeof(msgBuf),  "%s failed with error %d: %s\n", lpszFunction, dw, lpMsgBuf);
    SetEnvironmentVariable( "GBE_JOBDEBUG", msgBuf );
    fprintf(stderr, msgBuf);
    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;
}