/*
 * Shell wrapper for GNU Make
 *
 * Takes a bash command line as arguments, passes it to shell,
 * reads output into a buffer and when shell finishes, outputs
 * buffer contents to corresponding stdout/stderr streams. 
 *
 * Use by setting SHELL := mwrap
 *
 * If MWRAP_DEBUG environment variable is set, echo all executed
 * commands to stderr.
 *
 */

#include <sys/types.h>
#include <sys/file.h>
#include <sys/wait.h>
#include <sys/select.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>

const char *SHELL = "/bin/bash";
const char *LOCK = ".mwrap.lock";

#define BUFSIZE 65536

/* This is where output is temporarily stored. One struct per line
 * with corresponding output file descriptor number. */
struct linebuf {
    int fd;
    char *buffer;
};

/* Populate shell argument array pointer */
char** setargv(int argc, char *argv[])
{
    int i;
    char **shargv = calloc(1, (argc + 1) * sizeof(char *));
    shargv[0] = (char *) SHELL;
    for(i = 1; i < argc; i++)
        shargv[i] = argv[i];
    return shargv;
}

int main(int argc, char *argv[])
{
    pid_t child;
    int outfd[2], errfd[2], lockfd = -1;
    int status, retval, line = 0, buflen = 128, i = 0;
    char buffer[BUFSIZE];
    ssize_t bytesread;
    struct linebuf *outbuf = calloc(1, buflen * sizeof(struct linebuf));
    fd_set rset;

    char **shargv = setargv(argc, argv);

    pipe(outfd);
    pipe(errfd);

    /* Redirect stdout/stderr and execute */
    child = fork();
    if(child == (pid_t) 0) {
        close(1); close(2);
        close(outfd[0]); close(errfd[0]);
        dup2(outfd[1], 1); dup2(errfd[1], 2);
        execvp(SHELL, shargv);
        return 1; /* should never happen */
    } else if(child == (pid_t) -1)
        return 2; /* fork() failed */

    close(outfd[1]);
    close(errfd[1]);

/*
 * Run select() on child stdout and stederr file descriptors
 */
read:
    retval = 0;
    FD_ZERO(&rset);
    FD_SET(outfd[0], &rset);
    FD_SET(errfd[0], &rset);
    retval = select(20, &rset, NULL, NULL, NULL);
    /* Data ready in stdout */
    if(retval && FD_ISSET(outfd[0], &rset)) {
        if((bytesread = read(outfd[0], buffer, sizeof(buffer))) > 0) {
            if(line >= buflen) {
                buflen += buflen;
                outbuf = realloc(outbuf, sizeof(struct linebuf) * (buflen + line + 1));
            }
            outbuf[line].fd = 1;
            outbuf[line].buffer = calloc(1, bytesread+1);
            strncpy(outbuf[line].buffer, buffer, bytesread);
            line++;
        }
    }
    if(retval && FD_ISSET(errfd[0], &rset)) {
        if((bytesread = read(errfd[0], buffer, sizeof(buffer))) > 0) {
            if(line >= buflen) {
                buflen += buflen;
                outbuf = realloc(outbuf, sizeof(struct linebuf) * (buflen + line + 1));
            }
            outbuf[line].fd = 2;
            outbuf[line].buffer = calloc(1, bytesread+1);
            strncpy(outbuf[line].buffer, buffer, bytesread);
            line++;
        }
    }

    /* If child hasn't exited yet, select() again */
    while(waitpid(child, &status, WNOHANG) == 0) {
        i = 0;
        goto read;
    }
    /* Check once more for any data, just in case */
    if(i == 0) { i = 1; goto read; }

    /* Grab a lock */
    lockfd = open(LOCK, O_WRONLY|O_CREAT, 0600);
    flock(lockfd, LOCK_EX);

    /* Print the original command line if MWRAP_DEBUG env var is set */
    if(getenv("MWRAP_DEBUG") != NULL) {
        for(i = 2; i < argc; i++)
            fprintf(stderr, "%s ", argv[i]);
        fprintf(stderr, "\n");
    }

    /* Print out the buffered lines */
    for(i = 0; i < line; i++) {
        write(outbuf[i].fd, outbuf[i].buffer, strlen(outbuf[i].buffer));
        free(outbuf[i].buffer);
    }

    free(outbuf);
    free(shargv);
    unlink(LOCK);
    flock(lockfd, LOCK_UN);
    close(lockfd);

    if(WIFEXITED(status))
        return WEXITSTATUS(status);
    else
        return 127;
}