/*============================================================================
**
**    ERG TRANSIT SYSTEMS      Licensed software
**    (C) 2008                 All rights reserved
**
**============================================================================
**
**  Project/Product : 
**  Filename        : pipe.c
**  Author(s)       : DDP
**
**  Description     : Unix 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         : UNIX ONLY
**
**
**==========================================================================*/

#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/wait.h>
#include "string.h"
#include "strings.h"

#define MAX_LINE    1024                /* Limit of user lines */

/*
**  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;               /* Sise of above */
} proc;

char buf[MAX_LINE];
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 );

int main (int argc, char **argv)
{
    int     filedes1[2], filedes2[2], filedes3[2];
    int     pid;
    fd_set  fds;
    int     maxfd;
    int     nready;
    int     stderr_isopen;
    int     stdout_isopen;
    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 ( strncasecmp( argv[argp], "-OUT:", 5 ) == 0)
        {
            line_out.name = argv[argp] + 5;
        }
        else if ( strncasecmp( argv[argp], "-ERR:", 5 ) == 0)
        {
            line_err.name = argv[argp] + 5;
        }
        else
        {
            fprintf(stderr, "Unknown option: %s\n", argv[argp] );
            exit(-1);
        }
    }

    /*
    **  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 )
    {
      fputs("Insufficient number of arguments given\n", stderr);
      fputs("Need at least the name of the program\n", stderr);
      exit(-1);
    }

  
  /* Argument checking. */

  if (argc <= 1)
    {
      fputs("Insufficient number of arguments given\n", stderr);
      fputs("Need at least the name of the program\n", stderr);
      exit(-1);
    }

  /* Make our pipes. */

  if (pipe(filedes1) == -1)
    {
      perror ("pipe");
      exit(-1);
    }

  if (pipe(filedes2) == -1)
    {
      perror ("pipe");
      exit(-1);
    }

  if (pipe(filedes3) == -1)
    {
      perror ("pipe");
      exit(-1);
    }
    

  /* Fork a child */

    if ((pid = fork()) == 0)
    {
        /*
        **    Child
        */
        dup2(filedes1[0], fileno(stdin));     /* Copy the reading end of the pipe. */
        dup2(filedes2[1], fileno(stdout));    /* Copy the writing end of the pipe */
        dup2(filedes3[1], fileno(stderr));    /* Copy the writing end of the pipe */

        close( filedes1[1] );
        close( filedes2[0] );
        close( filedes3[0] );

         /* If execv() returns at all, there was an error. */

        if (execv(argv[argp], &argv[argp]) == -1)
        {
            perror("execv");
            exit(128);
        }
    }
    else if (pid == -1)
    {
        perror("fork");
        exit(128);
    }
    else
    {
        /*
        **    Parent
        */
        
        int     stat_loc;
        pid_t   result;
        close( filedes1[0] );
        close( filedes2[1] );
        close( filedes3[1] );

        /*
        **  Read from childs STDERR and STDOUT as the data becomes
        **  available
        **
        **  Prefex the data with an indication of the stream that
        **  it came from
        */

        /*
        ** Set up the select information
        */
        maxfd = (filedes2[0] > filedes3[0] ? filedes2[0] : filedes3[0]) + 1;

        stdout_isopen = 1;
        stderr_isopen = 1;
        while(stdout_isopen || stderr_isopen)
        {

	        /* Set up polling using select. */
		    FD_ZERO(&fds);
		    if (stdout_isopen)
                FD_SET(filedes2[0],&fds);
		    if (stderr_isopen)
                FD_SET(filedes3[0],&fds);
        
		    /* Wait for some input. */
		    nready = select(maxfd, &fds, (fd_set *) 0, (fd_set *) 0, (struct timeval *) 0);
            if ( nready < 0 )
            {
                break;
            }
		    /* If either descriptor has some input, read it */

		    if( FD_ISSET(filedes2[0], &fds))
		    {
			    int nread = read(filedes2[0], buf, sizeof(buf));
			    /* If error or eof, terminate. */
			    if(nread < 1){
        		    stdout_isopen = 0;
			    }
                proc_data( buf, nread, &line_out);
		    }

		    if( FD_ISSET(filedes3[0], &fds))
		    {
			    int nread = read(filedes3[0], buf, sizeof(buf));
			    /* If error or eof, terminate. */
			    if(nread < 1){
        		    stderr_isopen = 0;
			    }
                proc_data( buf, nread, &line_err );
		    }
        }

        flush_data( &line_out);
        flush_data( &line_err);

        /*
        **  Determine child exit status
        */
        result = waitpid(pid, &stat_loc, 0 );
        if ( result == pid )
        {
            if ( WIFEXITED(stat_loc)  )
                exit( WEXITSTATUS(stat_loc));
        }
    }

    /*
    **  Bad pid (and keep compiler happy)
    */
    return (-1);
}

/*----------------------------------------------------------------------------
** 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 - 2  )
        {
            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 )
{
    if ( p->ptr != p->buf )
    {
        *p->ptr++ = '\n';
        *p->ptr = 0;

        if ( p->name_len )
            fputs(p->name, stdout);
        fputs(p->buf, stdout);
        p->ptr = p->buf;
    }
}