/*============================================================================
**
**    ERG TRANSIT SYSTEMS      Licensed software
**    (C) 2008                 All rights reserved
**
**============================================================================
**
**  Project/Product : 
**  Filename        : pipe_win32.c
**  Author(s)       : DDP
**
**  Description     : Windows program to
**                      Create an other process
**                      Monitor its stdout/stderr
**                      Tag data with STDOUT or STDERR
**                      Write it out my STDOUT
**
**                    Used to help Perl monitor the status of
**                    complex processes. Current Perl IPC has
**                    problems monitoring both STDOUT and STDERR
**                    in cases where the output s non-trivial
**
**                    Only a problem for WIN32 under Perl as Open3
**                    cannot be used to read both streams - with complete
**                    sucess
**
** Usage            : stdmux <options>* command <command args>*
**                    Where options are:
**                      -out:<text>     - Text to prefix stdout. Default is "STDOUT:"
**                      -err:<text>     - Text to prefix stderr. Default is "STDERR:"
**                    Command           - Users command
**                    command args      - Arguments passed to the command
**
**  Information     :
**   Compiler       : ANSI C
**   Target         : WIN32 ONLY
**
**
**==========================================================================*/

#define _CRT_SECURE_NO_WARNINGS 1       // Supress warnings about use of strcpy()

#include <stdio.h>
#include <windows.h> 
 
#define BUFSIZE     4096
#define IBUFSIZE    4096
#define MAX_LINE    1024                // Limit of user lines
 
HANDLE  hChildStdinRd, hChildStdinWr, hChildStdinWrDup,
        hChildStdoutRd, hChildStdoutWr, hChildStdoutRdDup,
        hChildStderrRd, hChildStderrWr, hChildStderrRdDup,
        hInputFile, hSaveStdin, hSaveStdout, hSaveStderr;
CRITICAL_SECTION CriticalSection;
DWORD dwThreadIdOut, dwThreadIdErr;
HANDLE hThreadOut, hThreadErr;
DWORD   dwChildExitCode;
HANDLE  hChildProcess;

/*
**  Prototypes (forward declarations)
*/
HANDLE CreateChildProcess(LPTSTR);
VOID WriteToPipe(VOID); 
VOID ErrorExit(LPTSTR);
VOID ErrMsg(LPTSTR, BOOL); 
DWORD WINAPI ThreadReadStdOut( LPVOID lpParam );
DWORD WINAPI ThreadReadStdErr( LPVOID lpParam );

/*
**  A structure to contain the data from stdin and stdout before
**  it is printed to stdout, with a prefix
*/
typedef struct proc {
    char    *name;              /* Prefix to add */
    size_t  name_len;           /* Length of *name */
    char    *ptr;               /* Current pointer */
    char    *buf;               /* Buffer */
    size_t  size;               /* Size of above */
} proc;

char buf1[MAX_LINE];
char buf2[MAX_LINE];

proc    line_out = {"STDOUT:", 7, buf1, buf1, sizeof(buf1) } ;
proc    line_err = {"STDERR:", 7, buf2, buf2, sizeof(buf2) } ;

void proc_data( char *buf, size_t nread, proc *p );
void flush_data( proc *p );

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

int main(int argc, char *argv[])
{
   SECURITY_ATTRIBUTES saAttr; 
   BOOL     fSuccess;
   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], "-OUT:", 5 ) == 0)
        {
            line_out.name = argv[argp] + 5;
        }
        else if ( _strnicmp( argv[argp], "-ERR:", 5 ) == 0)
        {
            line_err.name = argv[argp] + 5;
        }
        else
        {
            fprintf(stderr, "Unknown option: %s\n", argv[argp] );
            ErrorExit( "Bad Option\n");
        }
    }

    /*
    **  Fixup the string lenghts
    */
    line_out.name_len = strlen( line_out.name);
    line_err.name_len = strlen( line_err.name);

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

    /*
    **  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 = malloc ( length + 1);
        if ( !command )
            ErrorExit("Cannot allocate memory\n");
        *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 security attribute
    **  Set the bInheritHandle flag so pipe handles are inherited
    */
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE; 
    saAttr.lpSecurityDescriptor = NULL; 
 
   // The steps for redirecting child process's STDOUT and STDERR:
   //     1. Save current STDOUT, to be restored later. 
   //     2. Create anonymous pipe to be STDOUT for child process. 
   //     3. Set STDOUT of the parent process to be write handle to 
   //        the pipe, so it is inherited by the child process. 
   //     4. Create a noninheritable duplicate of the read handle and
   //        close the inheritable read handle. 
 
// Save the handle to the current STDOUT. 
 
   hSaveStdout = GetStdHandle(STD_OUTPUT_HANDLE);
   hSaveStderr = GetStdHandle(STD_ERROR_HANDLE);
 
// Create a pipe for the child process's STDOUT. 
 
   if (! CreatePipe(&hChildStdoutRd, &hChildStdoutWr, &saAttr, 0)) 
      ErrorExit("Stdout pipe creation failed\n");

   if (! CreatePipe(&hChildStderrRd, &hChildStderrWr, &saAttr, 0))
      ErrorExit("Stderr pipe creation failed\n");
     
 
// Set a write handle to the pipe to be STDOUT. 
 
   if (! SetStdHandle(STD_OUTPUT_HANDLE, hChildStdoutWr)) 
      ErrorExit("Redirecting STDOUT failed"); 

   if (! SetStdHandle(STD_ERROR_HANDLE, hChildStderrWr))
      ErrorExit("Redirecting STDERR failed");
 
// Create noninheritable read handle and close the inheritable read 
// handle. 

    fSuccess = DuplicateHandle(GetCurrentProcess(), hChildStdoutRd,
        GetCurrentProcess(), &hChildStdoutRdDup , 0,
        FALSE,
        DUPLICATE_SAME_ACCESS);
    if( !fSuccess )
        ErrorExit("DuplicateHandle failed");
    CloseHandle(hChildStdoutRd);

    fSuccess = DuplicateHandle(GetCurrentProcess(), hChildStderrRd,
        GetCurrentProcess(), &hChildStderrRdDup , 0,
        FALSE,
        DUPLICATE_SAME_ACCESS);
    if( !fSuccess )
        ErrorExit("DuplicateHandle failed");
    CloseHandle(hChildStderrRd);
    
#if 0
   // The steps for redirecting child process's STDIN: 
   //     1.  Save current STDIN, to be restored later. 
   //     2.  Create anonymous pipe to be STDIN for child process. 
   //     3.  Set STDIN of the parent to be the read handle to the 
   //         pipe, so it is inherited by the child process. 
   //     4.  Create a noninheritable duplicate of the write handle, 
   //         and close the inheritable write handle. 
 
// Save the handle to the current STDIN. 
 
   hSaveStdin = GetStdHandle(STD_INPUT_HANDLE); 
 
// Create a pipe for the child process's STDIN. 
 
   if (! CreatePipe(&hChildStdinRd, &hChildStdinWr, &saAttr, 0)) 
      ErrorExit("Stdin pipe creation failed\n"); 
 
// Set a read handle to the pipe to be STDIN. 
 
   if (! SetStdHandle(STD_INPUT_HANDLE, hChildStdinRd)) 
      ErrorExit("Redirecting Stdin failed"); 
 
// Duplicate the write handle to the pipe so it is not inherited. 
 
   fSuccess = DuplicateHandle(GetCurrentProcess(), hChildStdinWr, 
      GetCurrentProcess(), &hChildStdinWrDup, 0, 
      FALSE,                  // not inherited 
      DUPLICATE_SAME_ACCESS); 
   if (! fSuccess) 
      ErrorExit("DuplicateHandle failed"); 
 
   CloseHandle(hChildStdinWr);
#endif
 
    /*
    **  Now create the child process.
    */
    hChildProcess = CreateChildProcess(command);
    if (! hChildProcess)
        ErrorExit("Create process failed");

    /*
    **  Restore the saved STDIN and STDOUT and STDERR.
    */
    if (! SetStdHandle(STD_INPUT_HANDLE, hSaveStdin))
        ErrorExit("Re-redirecting Stdin failed\n");
 
   if (! SetStdHandle(STD_OUTPUT_HANDLE, hSaveStdout)) 
      ErrorExit("Re-redirecting Stdout failed\n");

   if (! SetStdHandle(STD_ERROR_HANDLE, hSaveStderr))
      ErrorExit("Re-redirecting Stderr failed\n");
    
 
//// Get a handle to the parent's input file.
//
//   if (argc > 1)
//      hInputFile = CreateFile(argv[1], GENERIC_READ, 0, NULL,
//         OPEN_EXISTING, FILE_ATTRIBUTE_READONLY, NULL);
//   else
//      hInputFile = hSaveStdin;
//
//   if (hInputFile == INVALID_HANDLE_VALUE)
//      ErrorExit("no input file\n");
//
//// Write to pipe that is the standard input for a child process.
//
//   WriteToPipe();
 
// Read from pipe that is the standard output for child process. 


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

    /*
    **  Create two thread to read the processes stdout and stderr
    **  pipes
    */
    hThreadOut = CreateThread(
        NULL,                        // default security attributes 
        0,                           // use default stack size  
        ThreadReadStdOut,            // thread function
        NULL,                        // argument to thread function
        0,                           // use default creation flags 
        &dwThreadIdOut);             // returns the thread identifier

    if (hThreadOut == NULL)
        ErrorExit("Error Creating thread ReadOut\n");

    hThreadErr = CreateThread(
        NULL,                        // default security attributes 
        0,                           // use default stack size  
        ThreadReadStdErr,            // thread function
        NULL,                        // argument to thread function
        0,                           // use default creation flags 
        &dwThreadIdErr);                // returns the thread identifier

    if (hThreadErr == NULL)
        ErrorExit("Error Creating thread ReadErr\n");

    /*
    **  What for the threads to complete
    */
    WaitForSingleObject ( hThreadOut, INFINITE );
    WaitForSingleObject ( hThreadErr, INFINITE );

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

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

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

/*----------------------------------------------------------------------------
** FUNCTION           : CreateChildProcess
**
** DESCRIPTION        : Craete the child process
**
**
** 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 
            0,             // 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           : WriteToPipe
**
** DESCRIPTION        : Write data to processes stdin
**                      Not used (yet)
**
**
** INPUTS             :
**
** RETURNS            :
**
----------------------------------------------------------------------------*/

VOID WriteToPipe(VOID)
{ 
   DWORD dwRead, dwWritten; 
   CHAR chBuf[BUFSIZE]; 
 
// Read from a file and write its contents to a pipe. 
 
   for (;;) 
   { 
      if (! ReadFile(hInputFile, chBuf, BUFSIZE, &dwRead, NULL) || 
         dwRead == 0) break; 
      if (! WriteFile(hChildStdinWrDup, chBuf, dwRead, 
         &dwWritten, NULL)) break; 
   } 
 
// Close the pipe handle so the child process stops reading. 
 
   if (! CloseHandle(hChildStdinWrDup)) 
      ErrorExit("Close pipe failed\n"); 
} 
 

/*----------------------------------------------------------------------------
** FUNCTION           : ErrorExit
**
** 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(-1);
} 

/*----------------------------------------------------------------------------
** FUNCTION           : ThreadReadStdOut
**
** DESCRIPTION        : Thread to read child's stdout and re-badge the data
**                      Done on a thread so that we can handle both stdout
**                      and stderr
**
** INPUTS             : lpParam                 - Not used
**
** RETURNS            : 0
**
----------------------------------------------------------------------------*/

DWORD WINAPI ThreadReadStdOut( LPVOID lpParam )
{
   DWORD dwRead;
   CHAR chBuf[IBUFSIZE];
   HANDLE hStdout = GetStdHandle(STD_OUTPUT_HANDLE); 

   /*
   **   Close the write end of the pipe before reading from the
   **   read end of the pipe.
   */
   if (!CloseHandle(hChildStdoutWr))
      ErrorExit("Closing handle failed"); 
 
    for (;;) 
    { 
        if( !ReadFile( hChildStdoutRdDup, chBuf, IBUFSIZE, &dwRead,
            NULL) || dwRead == 0)
                break;

        proc_data( chBuf, dwRead, &line_out);
   }
    flush_data( &line_out);
    return 0;
}

DWORD WINAPI ThreadReadStdErr( LPVOID lpParam )
{
   DWORD dwRead;
   CHAR chBuf[IBUFSIZE];
   HANDLE hStderr = GetStdHandle(STD_ERROR_HANDLE);

   /*
   **   Close the write end of the pipe before reading from the
   **   read end of the pipe.
   */
   if (!CloseHandle(hChildStderrWr))
      ErrorExit("Closing err handle failed");
 
   for (;;)
   { 
      if( !ReadFile( hChildStderrRdDup, chBuf, IBUFSIZE, &dwRead,
         NULL) || dwRead == 0) break;
        proc_data( chBuf, dwRead, &line_err);
   } 
    flush_data( &line_err);
    return 0;
}

/*----------------------------------------------------------------------------
** FUNCTION           : proc_data
**
** DESCRIPTION        : Multiplex data to stdout
**                      Prefix each line with a 'TAG' that indicates where
**                      it came from.
**
**
** INPUTS             : buf                     - Data to process
**                      nread                   - Bytes of data
**                      p                       - p structure
**
** RETURNS            : Nothing
**
----------------------------------------------------------------------------*/

void proc_data( char *buf, size_t nread, proc *p )
{
    while ( nread-- )
    {
        char ch = *buf++;

        if ( ch == '\r' )
            continue;

        if ( ch == '\n' )
        {
            flush_data(p);
            continue;
        }

        *p->ptr++ = ch;
        if ( (size_t)(p->ptr - p->buf) >= p->size - 1  )
        {
            flush_data(p);
            continue;
        }
    }
}

/*----------------------------------------------------------------------------
** FUNCTION           : flush_data
**
** DESCRIPTION        : Write a complete line to the stdout
**                      Needs to be atomic write
**
**
** INPUTS             : p               - Ref to proc structure
**
** RETURNS            : Nothing
**
----------------------------------------------------------------------------*/
void flush_data( proc *p )
{
    DWORD dwWritten;
    if ( p->ptr != p->buf )
    {
        *p->ptr++ = '\n';

        EnterCriticalSection(&CriticalSection);
        WriteFile(hSaveStdout, p->name,(DWORD) p->name_len     , &dwWritten, NULL);
        WriteFile(hSaveStdout, p->buf, (DWORD)(p->ptr - p->buf), &dwWritten, NULL);
        LeaveCriticalSection(&CriticalSection);

        p->ptr = p->buf;
    }
}