/*
 * Copyright (C) 1994-2018 Altair Engineering, Inc.
 * For more information, contact Altair at www.altair.com.
 *
 * This file is part of the PBS Professional ("PBS Pro") software.
 *
 * Open Source License Information:
 *
 * PBS Pro is free software. You can redistribute it and/or modify it under the
 * terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * PBS Pro is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.
 * See the GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Commercial License Information:
 *
 * For a copy of the commercial license terms and conditions,
 * go to: (http://www.pbspro.com/UserArea/agreement.html)
 * or contact the Altair Legal Department.
 *
 * Altair’s dual-license business model allows companies, individuals, and
 * organizations to create proprietary derivative works of PBS Pro and
 * distribute them - whether embedded or bundled with other software -
 * under a commercial license agreement.
 *
 * Use of Altair’s trademarks, including but not limited to "PBS™",
 * "PBS Professional®", and "PBS Pro™" and Altair’s logos is subject to Altair's
 * trademark licensing policies.
 *
 */

/**
 * @file pbs_demux.c
 * @brief
 * pbs_demux - handle I/O from multiple node job
 *
 *	Standard Out and Standard Error of each task is bound to
 *	stream sockets connected to pbs_demux which inputs from the
 *	various streams and writes to the JOB's out and error.
 */

#include <pbs_config.h>   /* the master config generated by configure */
#include <pbs_version.h>

#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>

#if defined(FD_SET_IN_SYS_SELECT_H)
#  include <sys/select.h>
#endif

#include "cmds.h"

enum rwhere {invalid, new_out, new_err, old_out, old_err};
struct routem {
	enum rwhere	r_where;
	short		r_nl;
	short		r_first;
};
fd_set readset;
char   *cookie = 0;

/**
 * @brief
 *	read data from socket
 *
 * @param[in] sock - socket
 * @param[in] prm  - routem structure pointer
 *
 * @return - Void
 *
 */
void
readit(int sock, struct routem *prm)
{
	int   amt;
	char  buf[256];
	FILE *fil;
	int   i;
	char *pc;

	if (prm->r_where == old_out)
		fil = stdout;
	else
		fil = stderr;

	i = 0;
	if ((amt = read(sock, buf, 256)) > 0) {
		if (prm->r_first == 1) {

			/* first data on connection must be the cookie to validate it */

			i = strlen(cookie);
			if (strncmp(buf, cookie, i) != 0) {
				(void)close(sock);
				prm->r_where = invalid;
				FD_CLR(sock, &readset);
			}
			prm->r_first = 0;
		}
		for (pc = buf+i; pc < buf+amt; ++pc) {
#ifdef DEBUG
			if (prm->r_nl) {
				fprintf(fil, "socket %d: ", sock);
				prm->r_nl = 0;
			}
#endif /* DEBUG */
			putc(*pc, fil);
			if (*pc == '\n') {
				prm->r_nl = 1;
				fflush(fil);
			}
		}
	} else {
		close(sock);
		prm->r_where = invalid;
		FD_CLR(sock, &readset);
	}
	return;
}

int
main(int argc, char *argv[])
{
	struct timeval timeout;
	int	i;
	int	maxfd;
	int	main_sock_out = 3;
	int	main_sock_err = 4;
	int	n;
	int	newsock;
	pid_t	parent;
	fd_set selset;
	struct routem *routem;


	/*test for real deal or just version and exit*/

	execution_mode(argc, argv);

	parent = getppid();
	cookie = getenv("PBS_JOBCOOKIE");
	if (cookie == 0) {
		fprintf(stderr, "%s: no PBS_JOBCOOKIE found in the env\n",
			argv[0]);
		exit(3);
	}
#ifdef DEBUG
	printf("Cookie found in environment: %s\n", cookie);
#endif

	maxfd = sysconf(_SC_OPEN_MAX);
	routem = (struct routem *)malloc(maxfd*sizeof(struct routem));
	if (routem == NULL) {
		fprintf(stderr, "%s: out of memory\n", argv[0]);
		exit(2);
	}
	for (i=0; i<maxfd; ++i) {
		(routem+i)->r_where = invalid;
		(routem+i)->r_nl    = 1;
		(routem+i)->r_first = 0;
	}
	(routem+main_sock_out)->r_where = new_out;
	(routem+main_sock_err)->r_where = new_err;


	FD_ZERO(&readset);
	FD_SET(main_sock_out, &readset);
	FD_SET(main_sock_err, &readset);

	if (listen(main_sock_out, 5) < 0) {
		perror("listen on out");
		exit(5);
	}

	if (listen(main_sock_err, 5) < 0) {
		perror("listen on err");
		exit(5);
	}

	while (1) {

		selset = readset;
		timeout.tv_usec = 0;
		timeout.tv_sec  = 10;

		n = select(FD_SETSIZE, &selset, NULL, NULL, &timeout);
		if (n == -1) {
			if (errno == EINTR) {
				n = 0;
			} else {
				fprintf(stderr, "%s: select failed\n", argv[0]);
				exit(1);
			}
		} else if (n == 0) {
			if (kill(parent, 0)  == -1) {
#ifdef DEBUG
				fprintf(stderr, "%s: Parent has gone, and so do I\n",
					argv[0]);
#endif	/* DEBUG */
				break;
			}
		}

		for (i=0; n && i<maxfd; ++i) {
			if (FD_ISSET(i, &selset)) {     /* this socket has data */
				n--;
				switch ((routem+i)->r_where) {
					case new_out:
					case new_err:
						newsock = accept(i, 0, 0);
						(routem+newsock)->r_where =  (routem+i)->r_where== new_out ? old_out : old_err;
						FD_SET(newsock, &readset);
						(routem+newsock)->r_first = 1;
						break;
					case old_out:
					case old_err:
						readit(i, routem+i);
						break;
					default:
						fprintf(stderr, "%s: internal error\n", argv[0]);
						exit(2);
				}
			}
		}
	}
	return 0;
}
