/*
 * socketpipe - zero overhead remote process plumbing
 * Microsoft Windows version
 *
 * (C) Copyright 2003-2005 Diomidis Spinellis
 *
 * Permission to use, copy, and distribute this software and its
 * documentation for any purpose and without fee is hereby granted,
 * provided that the above copyright notice appear in all copies and that
 * both that copyright notice and this permission notice appear in
 * supporting documentation.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
 * MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * $Id: socketpipe-win.c 1.1 2005/09/27 06:00:52 dds Exp $
 *
 */

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <stdarg.h>

#include <limits.h>

#include <io.h>

#include <process.h>

#include <fcntl.h>

#include <ctype.h>

#include <errno.h>

#include <assert.h>

#include <winsock2.h>

#include <windows.h>

#include <mswsock.h>

#include <sys/types.h>


#define MAXHOSTNAMELEN 1024


#define STDIN_FILENO 0

#define STDOUT_FILENO 1


static char *progname;

static void
usage(const char *msg)
{
	fprintf(stderr, "%s: %s\n", progname, msg);
	fprintf(stderr, "usage:\t%s [-b] [-i|o|r|l { command [args ...] }]\n", progname);
	fprintf(stderr, "\t(must specify a -l and a -r command and at least one of -i or -o)\n");
#ifdef DEBUG

	fprintf(stderr, "\t%s -s host port command [args ...]\n", progname);
	fprintf(stderr, "\t(typically automatically executed at the remote end)\n");
#endif

	exit(1);
}


/* Return as a string the error description for err */
static char *
wstrerror(LONG err)
{
	static LPVOID lpMsgBuf;

	if (lpMsgBuf)
		LocalFree(lpMsgBuf);
	FormatMessage(
	    FORMAT_MESSAGE_ALLOCATE_BUFFER |
	    FORMAT_MESSAGE_FROM_SYSTEM |
	    FORMAT_MESSAGE_IGNORE_INSERTS,
	    NULL,
	    err,		// GetLastError() does not seem to work reliably here
	    MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), // Default language
	    (LPTSTR) &lpMsgBuf,
	    0,
	    NULL
	);
	return lpMsgBuf;
}

/*
 * Report failure due to a fatal error as directed
 * by fmt and its arguments and exit the program
 */
static void
fatal(const char *fmt, ...)
{
	va_list marker;

	fprintf(stderr, "%s: ", progname);
	va_start(marker, fmt);
	vfprintf(stderr, fmt, marker);
	va_end(marker);
	fputc('\n', stderr);
	exit(2);
}

/* Checked malloc */
static void *
xmalloc(size_t n)
{
	void *p = malloc(n);

	if (p)
		return p;
	else
		fatal("out of memory");
}

/* Program argument strings */
static char *inputv, *outputv, *remotev, *loginv;
static int batch = 0;

/*
 * Set the input, output, remote, and login strings based on
 * the arguments supplied
 */
static void
parse_arguments(char *argv[])
{
	char **p, **start, **s;
	char **result;
	int nest, len;

	for (p = argv + 1; *p; p++) {
		/* Require a single character option */
		if (p[0][0] != '-' || !p[0][1] || p[0][2])
			usage("single character option expected");
		switch (p[0][1]) {
		case 'i': result = &inputv; break;
		case 'o': result = &outputv; break;
		case 'r': result = &remotev; break;
		case 'l': result = &loginv; break;
		case 'b': batch = 1; continue;
		default: usage("invalid option");
		}
		if (!*++p || strcmp(*p, "{"))
			usage("opening block expected");
		start = ++p;
		nest = 1;
		for (;;) {
			if (!*p)
				usage("unterminated block");
			if (strcmp(*p, "{") == 0)
				nest++;
			if (strcmp(*p, "}") == 0)
				nest--;
			if (nest == 0)
				break;
			p++;
		}
		if (p - start == 0)
			usage("command can not be empty");
		len = 0;
		for (s = start; s != p; s++)
			len += strlen(*s) + 1;
		*result = (char *)xmalloc(len);
		for (s = start; s != p; s++) {
			strcat(*result, *s);
			if (s + 1 != p)
				strcat(*result, " ");
		}
	}
}

/*
 * Client invocation interface
 * Run the remote command on the remote machine connecting it to
 * the local input and/or output processes.
 */
static void
client(char *argv[])
{
	struct sockaddr_in loc_addr;
	SOCKET sockfd, newsockfd;
	int addr_len;
	char hostname[MAXHOSTNAMELEN];
	char portname[20];
	char *rloginv;
	PROCESS_INFORMATION rempi;			/* Last child in the pipeline */
	PROCESS_INFORMATION inpi, outpi;
	HANDLE err_proc;		/* Process from whick to retrieve the error code */
	DWORD nwait = 0;			/* Children to wait for */
	DWORD exitstatus, bytes, waitret;
	STARTUPINFO rstart;
	WSAOVERLAPPED overlap;
	char lpOutputBuf[1024];
	LINGER linger;

	parse_arguments(argv);
	if (!remotev)
		usage("must specify remote command");
	if (!loginv)
		usage("must specify remote login method");
	if (!inputv && !outputv)
		usage("must specify a local input or output process");

	/* Create socket to remote end */
	if (gethostname(hostname, sizeof(hostname)) < 0)
		fatal("gethostname failed: %s", wstrerror(WSAGetLastError()));

	if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
		fatal("socket allocation failed: %s", wstrerror(WSAGetLastError()));

	memset((char *)&loc_addr, 0, sizeof(loc_addr));
	loc_addr.sin_family = AF_INET;
	loc_addr.sin_addr.s_addr = htonl(INADDR_ANY);	/* Us */
	loc_addr.sin_port = 0;				/* Kernel assigned port */

	if (bind(sockfd, (struct sockaddr *)&loc_addr, sizeof(loc_addr)) < 0)
		fatal("bind to local address failed: %s", wstrerror(WSAGetLastError()));

	if (listen(sockfd, 1) < 0)
		fatal("listen failed: %s", wstrerror(WSAGetLastError()));

	addr_len = sizeof(loc_addr);
	if (getsockname(sockfd, (struct sockaddr *)&loc_addr, &addr_len) < 0)
		fatal("getsockname failed: %s", wstrerror(WSAGetLastError()));
	sprintf(portname, "%d", ntohs(loc_addr.sin_port));

	/* Merge loginv and remotev into rloginv */
	rloginv = (char *)xmalloc(
		strlen(loginv) + 1 +
		strlen("socketpipe -s") + 1 +
		strlen(hostname) + 1 +
		strlen(portname) + 1 +
		strlen(remotev));
	strcpy(rloginv, loginv);
	strcat(rloginv, " socketpipe -s ");
	strcat(rloginv, hostname); strcat(rloginv, " ");
	strcat(rloginv, portname); strcat(rloginv, " ");
	strcat(rloginv, remotev);

	/* Remotely execute the command specified */
	/* XXX handle batch flag? */
	memset(&rstart, 0, sizeof(rstart));
	rstart.cb = sizeof(rstart);
	rstart.wShowWindow = SW_HIDE;

	if (!CreateProcess(NULL, rloginv, NULL, NULL, 0, NORMAL_PRIORITY_CLASS, NULL, NULL, &rstart, &rempi))
		fatal("execution of { %s } failed: %s", rloginv, wstrerror(GetLastError()));
	err_proc = rempi.hProcess;
	/* Accept a connection */
	/*
	 * It seems we have to create newsockfd through WSASocket/AcceptEx for ReadFile/WriteFile
	 * redirection through it to work.  Socket descriptors returned from socket/accept did not
	 * work.  This may just be an implementation accident of the way the socket provider is
	 * found.  In such a case we may need to work harder to find the right provider; see
	 * perl.perl5.porters > socket() call uses non-IFS providers causing subsequent print/read to hang or misbehave
	 * http://groups.google.com/group/perl.perl5.porters/msg/b23f04f62cdbd945?hl=en&utoken=dfS2yzIAAACb-xRze_lVAgVCnJd7_BU4TY4yx6oSFVIPwClGmhM46mb6oOYE8F84y-vJWNEg3rv1Ad0ufVf6pcmlCk7086Li
	 */
	if ((newsockfd = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, 0 /* WSA_FLAG_OVERLAPPED */)) == INVALID_SOCKET)
		fatal("WSASocket failed: %s", wstrerror(WSAGetLastError()));
	linger.l_onoff = 1;
	linger.l_linger = 60;
	if (setsockopt(newsockfd, SOL_SOCKET, SO_LINGER, (const char *)&linger, sizeof(linger)) != 0)
		fatal("setsockopt(SO_LINGER) failed: %s", wstrerror(WSAGetLastError()));
	memset(&overlap, 0, sizeof(overlap));
	if ((overlap.hEvent = CreateEvent(NULL, FALSE, FALSE, NULL)) == NULL)
		fatal("CreateEvent failed: %s", wstrerror(GetLastError()));
	if (!AcceptEx(sockfd, newsockfd, lpOutputBuf, 0, sizeof(SOCKADDR_IN) + 16, sizeof(SOCKADDR_IN) + 16, &bytes, &overlap) &&
	    WSAGetLastError() != ERROR_IO_PENDING)
		fatal("AcceptEx failed: %s", wstrerror(WSAGetLastError()));
	if ((waitret = WaitForMultipleObjects(1, &(overlap.hEvent), FALSE, INFINITE)) == WAIT_FAILED)
		fatal("WaitForMultipleObjects on AcceptEx failed: %s", wstrerror(GetLastError()));
	assert(waitret == WAIT_OBJECT_0);
	nwait++;
	/* Run the I/O commands redirected to newsockfd */
	inpi.hProcess = outpi.hProcess = INVALID_HANDLE_VALUE;

	/* Allow the socket to be inherited */
	if (!SetHandleInformation((HANDLE)newsockfd, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT))
		fatal("SetHandleInformation failed: %s", wstrerror(GetLastError()));
	/* Run the input generation process */
	if (inputv) {
		STARTUPINFO istart;

		nwait++;
		memset(&istart, 0, sizeof(istart));
		istart.cb = sizeof(istart);
		istart.wShowWindow = SW_HIDE;
		istart.dwFlags = STARTF_USESTDHANDLES;
		istart.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
		istart.hStdOutput = (HANDLE)newsockfd;
		istart.hStdInput = GetStdHandle(STD_ERROR_HANDLE);
		if (!CreateProcess(NULL, inputv, NULL, NULL, 1, NORMAL_PRIORITY_CLASS, NULL, NULL, &istart, &inpi))
			fatal("execution of { %s } failed: %s", inputv, wstrerror(GetLastError()));
	}

	/* Run the output processing process */
	if (outputv) {
		STARTUPINFO istart;

		nwait++;
		memset(&istart, 0, sizeof(istart));
		istart.cb = sizeof(istart);
		istart.wShowWindow = SW_HIDE;
		istart.dwFlags = STARTF_USESTDHANDLES;
		istart.hStdInput = (HANDLE)newsockfd;
		istart.hStdOutput = GetStdHandle(STD_OUTPUT_HANDLE);
		istart.hStdInput = GetStdHandle(STD_ERROR_HANDLE);
		if (!CreateProcess(NULL, outputv, NULL, NULL, 1, NORMAL_PRIORITY_CLASS, NULL, NULL, &istart, &outpi))
			fatal("execution of { %s } failed: %s", outputv, wstrerror(GetLastError()));
		err_proc = outpi.hProcess;
	}

	closesocket(newsockfd);
	/* Wait for all our children to terminate */
	while (nwait) {
		HANDLE handles[3];
		DWORD nCount;

		/* Create array of objects we are waiting on */
		nCount = 0;
		if (rempi.hProcess != INVALID_HANDLE_VALUE) handles[nCount++] = rempi.hProcess;
		if (inpi.hProcess != INVALID_HANDLE_VALUE) handles[nCount++] = inpi.hProcess;
		if (outpi.hProcess != INVALID_HANDLE_VALUE) handles[nCount++] = outpi.hProcess;
		assert(nCount == nwait);
		if ((waitret = WaitForMultipleObjects(nCount, handles, FALSE, INFINITE)) == WAIT_FAILED)
			fatal("WaitForMultipleObjects failed: %s", wstrerror(GetLastError()));
		nwait--;

		if (handles[waitret] == err_proc) {
			if (!GetExitCodeProcess(err_proc, &exitstatus))
				fatal("GetExitCodeProcess failed: %s", wstrerror(GetLastError()));
		}

		if (handles[waitret] == rempi.hProcess)
			rempi.hProcess = INVALID_HANDLE_VALUE;
		else if (handles[waitret] == inpi.hProcess)
			inpi.hProcess = INVALID_HANDLE_VALUE;
		else if (handles[waitret] == outpi.hProcess)
			outpi.hProcess = INVALID_HANDLE_VALUE;
	}
	exit(exitstatus);
}

/*
 * Server invocation interface.
 * Run as the remote server executing the specified command and
 * connecting back to the client.
 * This compiles, but has not been tested under Windows.
 */
static void
server(char *argv[])
{
	short port;
	struct sockaddr_in rem_addr;
	int sock;
	char *endptr;
	struct hostent *h;

	port = (short)strtol(argv[3], &endptr, 10);
	if (*argv[3] == 0 || *endptr != 0)
		fatal("bad port specification: %s", argv[3]);

	if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0)
		fatal("socket allocation failed: ", wstrerror(WSAGetLastError()));

	memset((char *)&rem_addr, 0, sizeof(rem_addr));
	rem_addr.sin_port = htons(port);

	if (isalpha(*argv[2])) {		/* host address is a name */
		if ((h = gethostbyname(argv[2])) == NULL)
			fatal("gethostbyname(%s) failed: %s", argv[2], wstrerror(WSAGetLastError()));
	} else  {
		unsigned int addr = inet_addr(argv[2]);
		h = gethostbyaddr((char *)&addr, 4, AF_INET);
	}

	memcpy(&rem_addr.sin_addr, h->h_addr_list[0], sizeof(rem_addr.sin_addr));
	rem_addr.sin_family = AF_INET;

	if (connect(sock, (struct sockaddr *)&rem_addr, sizeof(rem_addr)) < 0)
		fatal("connect(%s) failed: %s", argv[2], wstrerror(WSAGetLastError()));

	/* Redirect I/O to the socket */
	if (dup2(sock, STDIN_FILENO) < 0)
		fatal("input redirection failed: %s", strerror(errno));

	if (dup2(sock, STDOUT_FILENO) < 0)
		fatal("output redirection failed: %s", strerror(errno));

	if (execvp(argv[4], argv + 4) < 0)
		fatal("exec(%s) failed: %s", argv[4], strerror(errno));
	/* NOTREACHED */
}

int
main(int argc, char *argv[])
{
	WSADATA wsaData;
	WORD v;

	progname = argv[0];
	v = MAKEWORD(2, 2);
	if (WSAStartup(v, &wsaData) != 0)
		fatal("WSAStartup failed: %s", wstrerror(GetLastError()));
	if (!argv[1])
		usage("no arguments specified");
	if (strcmp(argv[1], "-s") == 0)
		server(argv);
	else
		client(argv);
	/* NOTREACHED */
}