#define _WIN32_WINNT 0x0500 /*Win2k*/
#define STRICT

#include <windows.h>
#include <stdio.h>
#include <getopt.h>
#include <sys/cygwin.h>
#include <sys/queue.h>
#include <process.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <errno.h>
#include <Tlhelp32.h>


#define PRGNAME "injob"
#define PRGVER "1.0"
#define PRGAUTHOR "Daniel Colascione <dan.colascione@gmail.com>"
#define PRGCOPY "Copyright (C) 2011 " PRGAUTHOR
#define PRGLICENSE "GPLv2 or later <http://www.gnu.org/licenses/gpl-2.0.html>"

/**
 * Small utility to run an arbitrary set of processes within a job
 * object.  We reach to Cygwin job control signals by appropriately
 * manipulating the job object, providing a crude form of job control
 * for Win32 applications being run from Cygwin programs.
 *
 * It works like this:
 *
 *   - Startup.
 *
 *   - Create a pipe with ends P_R and P_W.
 *
 *   - Block signals.
 *
 *   - Fork
 *     
 *     * Child closes P_W, blocks reading P_R.  If it gets EOF, child
 *       knows parent died for some reason and exits without doing
 *       anything else.
 *
 *     * Child reads 1 byte from pipe, indicating all-clear.
 *
 *     * Child execs target program.
 *     
 *   - Meanwhile parent closes P_R and knows child is blocked on pipe.
 *     
 *   - Parent creates job object and puts the child
 *     into it (child is still blocked).
 *
 *   - Parent gives all-clear signal to child by writing one byte to
 *     P_W and closing it.
 *
 *   - Parent waits for SIGINT, SIGTERM, SIGCHLD, etc.
 *
 */

static BOOL WINAPI
(*XIsProcessInJob)(
    HANDLE ProcessHandle,
    HANDLE JobHandle,
    PBOOL Result
    );

#define CHK(op)                                                 \
    ({                                                          \
        int chk_ret;                                            \
                                                                \
        do {                                                    \
            chk_ret = (op);                                     \
        } while (chk_ret == -1 && errno == EINTR);              \
                                                                \
        if (chk_ret == -1) {                                    \
            fprintf (stderr, PRGNAME ": " #op ": %s\n",         \
                     strerror (errno));                         \
            goto out;                                           \
        }                                                       \
        chk_ret;                                                \
    })

#define CHK_W32_HANDLE(op)                                      \
    ({                                                          \
        HANDLE chk_ret = (op);                                  \
        if (chk_ret == NULL ||                                  \
            chk_ret == INVALID_HANDLE_VALUE)                    \
        {                                                       \
            fprintf (stderr, PRGNAME ": " #op ": %s\n",         \
                     errmsg (GetLastError ()));                 \
            goto out;                                           \
        }                                                       \
                                                                \
        chk_ret;                                                \
    })

#define CHK_W32_BOOL(op)                                        \
    ({                                                          \
        BOOL chk_ret = (op);                                    \
        if (chk_ret == FALSE) {                                 \
            fprintf (stderr, PRGNAME ": " #op ": %s\n",         \
                     errmsg (GetLastError ()));                 \
            goto out;                                           \
        }                                                       \
                                                                \
        chk_ret;                                                \
    })

    
#define PIPE_READ  0
#define PIPE_WRITE 1

struct suspend
{
    DWORD   thread_id;
    HANDLE  thread;
    
    SLIST_ENTRY (suspend) entries;
};

static void
usage()
{
    fprintf(
        stdout, 
        PRGNAME " PROGRAM ARG1 ARG2...: Run PROGRAM in a job object\n"
        "\n"
        "  PROGRAM will be run in a job object.  A SIGTERM or SIGINT sent to\n"
        "  this proess will terminate PROGRAM and all its children.\n"
        "\n"
        "  Both Cygwin and non-Cygwin children will be terminated.\n"
        "\n"
        PRGNAME " -h\n"
        PRGNAME " --help\n"
        "\n"
        "  Display this help message.\n"
        "\n"
        PRGNAME " -V\n"
        PRGNAME " --version\n"
        "\n"
        "  Display version information.\n"
	);
}

static void
versinfo ()
{
    fprintf(stdout,
            PRGNAME " " PRGVER "\n"
            PRGCOPY "\n"
            PRGLICENSE "\n"
        );
}

/* Decode a Win32 error code to a localized string. Return
   a malloc()ed string. */
static char*
errmsg(DWORD errorcode)
{
	char* msg = NULL;
        char* msg_fmt;
        
	FormatMessageA(
		FORMAT_MESSAGE_FROM_SYSTEM|
		FORMAT_MESSAGE_ALLOCATE_BUFFER,
		NULL,
		errorcode,
		0,
		(LPTSTR)&msg_fmt,
		0,
		NULL);

	if(msg_fmt == NULL) {
            msg = strdup("[unknown error]");
	} else {
            msg = strdup (msg_fmt);
            LocalFree (msg_fmt);
        }

	if (msg[strlen(msg) - 1] == '\n') {
            msg[strlen(msg) - 1] = '\0';
	}

	return msg;
}

/* Holds a list of processes we've suspended. */

typedef SLIST_HEAD (suspend_list_head, suspend) suspend_list_head_t;
static suspend_list_head_t suspend_list_head = SLIST_HEAD_INITIALIZER (suspend_list_head);

/* Resume everything we remember we suspended. */
static void
resume_all ()
{
    struct suspend* susp;
    fprintf (stderr, "RESUMING\n");

    while (!SLIST_EMPTY (&suspend_list_head)) {
        susp = SLIST_FIRST (&suspend_list_head);
        SLIST_REMOVE_HEAD (&suspend_list_head, entries);

        /* Don't care about failures here. Best effort. */
        ResumeThread (susp->thread);
        CloseHandle (susp->thread);
        free (susp);
    }
}

static BOOL
is_thread_already_suspended (DWORD thread_id)
{
    struct suspend* s;
    SLIST_FOREACH (s, &suspend_list_head, entries) {
        if (s->thread_id == thread_id) {
            return TRUE;
        }
    }

    return FALSE;
}

/* Brute force. Suspend until we can't suspend anymore. */
static void
suspend_all_in_job (HANDLE job)
{
    unsigned nr_suspended;
    HANDLE snap;
    THREADENTRY32 thent;
    HANDLE proc;
    HANDLE thread;
    BOOL process_in_job;
    struct suspend* new_susp = NULL;
    unsigned try_count = 1000;

    fprintf (stderr, "tryin to suspend all in job\n");

    do {
        nr_suspended = 0;
        snap = CreateToolhelp32Snapshot (TH32CS_SNAPTHREAD, 0);

        if (snap == NULL) {
            goto next_snapshot;
        }
        
        memset (&thent, 0, sizeof(thent));
        thent.dwSize = sizeof(thent);
        
        if (Thread32First (snap, &thent) == FALSE) {
            goto next_snapshot;
        }
        
        do {
            // Thread32{First,Next} is allowed to return a struct
            // smaller that the one we asked for, so many sure the
            // fields we used are included in this structure.
            
            if (thent.dwSize < sizeof(DWORD)*4) {
                goto next_thread;
            }

            proc = NULL;
            thread = NULL;

            proc = OpenProcess (PROCESS_ALL_ACCESS,
                                FALSE /*do not inherit*/,
                                thent.th32OwnerProcessID);

            if (proc == NULL) {
                goto next_thread;
            }

            if (XIsProcessInJob (proc, job, &process_in_job) == FALSE) {
                goto next_thread;
            }
            
            if (process_in_job == FALSE) {
                goto next_thread;
            }
            
            /* We found a thread in a process that's in our job.  Now
             * here's the O(N^2) part where we examine our entire list
             * so far to make sure we haven't already noticed this
             * thread. */
            
            if (is_thread_already_suspended (thent.th32ThreadID)) {
                goto next_thread;
            }
            
            thread = OpenThread (THREAD_SUSPEND_RESUME,
                                 FALSE /*do not inherit*/,
                                 thent.th32ThreadID);
            
            if (thread == NULL) {
                goto next_thread;
            }
            
            /* We found a match we didn't notice before. */
            
            new_susp = malloc (sizeof(*new_susp));
            if (new_susp == NULL) {
                goto next_thread;
            }
            
            new_susp->thread = thread;
            new_susp->thread_id = thent.th32ThreadID;
            
            /* Try to suspend the thread */

            if (SuspendThread (thread) == (DWORD)-1) {
                goto next_thread;
            }
            
            SLIST_INSERT_HEAD (&suspend_list_head, new_susp, entries);
            ++nr_suspended;

            fprintf (stderr, "  suspended threads nr_suspended=%u\n",
                     (unsigned)nr_suspended);
            
            thread = NULL;
            new_susp = NULL;
            
          next_thread:

            free (new_susp);

            if (thread != NULL) {
                CloseHandle (thread);
            }

            if (proc != NULL) {
                CloseHandle (proc);
            }
            
            memset (&thent, 0, sizeof(thent));
            thent.dwSize = sizeof(thent);
        } while (Thread32Next (snap, &thent));
        
      next_snapshot:
        if (snap != NULL) {
            CloseHandle (snap);
            snap = NULL;
        }
    } while (try_count-- && nr_suspended > 0);
}

static void
dummy_sighandler (int dummy)
{}

static const struct option
longopts[] = 
{
	{ "help",      0, 0, 'h' },
        { "version",   0, 0, 'V' },
	{ 0 }
};

static int
child_main (int argc, char** argv, int* child_pipe)
{
    int ret = 1;
    ssize_t rret;
    char buf[1];
    int child_status;

    CHK (close (child_pipe[PIPE_WRITE]));

    do {
        rret = read (child_pipe[PIPE_READ], &buf, 1);
    } while (rret == -1 && errno == EINTR);

    if (rret == 0) {
        /* Parent died before it readied us, so die along with it. */
        goto out;
    }

    CHK (close (child_pipe[PIPE_READ]));

    do {
        ret = spawnvp (_P_NOWAIT, argv[0], (const char**)argv);
    } while (ret == -1 && errno == EINTR);

    if (ret == -1) {
        fprintf (stderr, PRGNAME ": could not spawn \"%s\": %s\n",
                 argv[0], strerror (errno));
        ret = 128;
        goto out;
    }

    CHK (wait (&child_status));
        
    ret = ( WIFEXITED (child_status)
            ? WEXITSTATUS (child_status)
            : 128 + WTERMSIG (child_status) );

  out:
    return ret;
}

int
main (int argc, char** argv)
{
    int c;
    int ret = 1;
    HANDLE job;
    int child_status;
    sigset_t waitmask;
    sigset_t origmask;
    pid_t child_pid;
    int child_pipe[2];
    int sig;
    DWORD child_w32_pid;
    HANDLE child_proc_handle;
    JOBOBJECT_BASIC_LIMIT_INFORMATION job_basic_limits;
    HANDLE kernel32dll;

    /* Initialize */

    while ((c = getopt_long(argc, argv, "Vh", longopts, 0)) != -1) {
        switch(c) {
            case 'h':
                usage();
                ret  = 0;
                goto out;
            case 'V':
                versinfo ();
                ret = 0;
                goto out;
            default:
                fprintf(stderr, PRGNAME ": use --help for usage\n");
                goto out;
        }
    }

    argc -= optind;
    argv += optind;

    if (argc == 0) {
        fprintf (stderr, PRGNAME ": missing PROGRAM argument\n");
        fprintf (stderr, PRGNAME ": use --help for usage\n");
        goto out;
    }

    kernel32dll = CHK_W32_HANDLE (LoadLibrary ("kernel32.dll"));
    XIsProcessInJob = GetProcAddress (kernel32dll, "IsProcessInJob");
    if (XIsProcessInJob == NULL) {
        fprintf (stderr, PRGNAME ": could not find IsProcessInJob: OS too old?\n");
        goto out;
    }
    FreeLibrary (kernel32dll);
    
    
    CHK (pipe (child_pipe));
    fflush (NULL);

    /* Signals blocked below, except while waiting. */

    CHK (sigemptyset (&waitmask));
    CHK (sigaddset (&waitmask, SIGCHLD));
    CHK (sigaddset (&waitmask, SIGTERM));
    CHK (sigaddset (&waitmask, SIGINT));
    CHK (sigaddset (&waitmask, SIGTSTP));
    CHK (sigaddset (&waitmask, SIGALRM));
    CHK (sigprocmask (SIG_BLOCK, &waitmask, &origmask));

    signal (SIGCHLD, dummy_sighandler);
    signal (SIGTERM, dummy_sighandler);
    signal (SIGINT,  dummy_sighandler);
    signal (SIGTSTP, dummy_sighandler);
    signal (SIGALRM, dummy_sighandler);
    
    /* Create child.  Child will just wait for us at first. */
    
    child_pid = CHK (fork ());
    if (child_pid == 0) {
        CHK (sigprocmask (SIG_SETMASK, &origmask, NULL));
        return child_main (argc, argv, child_pipe);
    }

    CHK (close (child_pipe[PIPE_READ]));
    
    /* Child is alive and blocked. Set up its job object. */

    child_w32_pid = (DWORD)cygwin_internal (
        CW_CYGWIN_PID_TO_WINPID, child_pid);

    if (child_w32_pid == 0) {
        fprintf (stderr, PRGNAME ": could not get child W32 PID\n");
        goto out;
    }

    
    child_proc_handle = CHK_W32_HANDLE (
        OpenProcess (PROCESS_ALL_ACCESS,
                     FALSE,
                     child_w32_pid));
    
    job = CHK_W32_HANDLE (CreateJobObject (NULL, NULL));
    CHK_W32_BOOL (AssignProcessToJobObject (job, child_proc_handle));
    CHK_W32_BOOL (CloseHandle (child_proc_handle));

    /* Child is now in the job object. Send the all-clear signal. */

    CHK (write (child_pipe[PIPE_WRITE], &ret /* arbitary byte */, 1));
    CHK (close (child_pipe[PIPE_WRITE]));

    /* Begin processing signals as they come in. */

    for (;;) {
        switch ((sig = sigwaitinfo (&waitmask, NULL))) {
            case SIGCHLD:
                /* Child exited.  We get its exist status below. */
                goto done;
            case SIGTERM:
            case SIGINT:
                /* We were asked to quit.  Kill everything in the job. */
                CHK_W32_BOOL (TerminateJobObject (job, 1));
                goto done;
            case SIGTSTP:
                suspend_all_in_job (job);
                raise (SIGSTOP);
                resume_all ();
                break;
            default:
                fprintf (stderr, PRGNAME ": unexpected signal %d: %s\n",
                         sig, strerror (errno));
                goto out;
        }
    }

  done:

    CHK (wait (&child_status));
        
    ret = ( WIFEXITED (child_status)
            ? WEXITSTATUS (child_status)
            : 128 + WTERMSIG (child_status) );
    
    out:
    return ret;
}