/*
 * 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_win.c
 *
 * @brief
 * process for de-multiplexing output of a multinode job
 */
#include <pbs_config.h>   /* the master config generated by configure */
#include <pbs_version.h>
#include "pbs_ifl.h"
#include <errno.h>
#include <string.h>
#include <limits.h>
#include <windows.h>
#include "win.h"
#include <stdio.h>
#include <stdlib.h>
#include <process.h>
#include "win_remote_shell.h"
#include "pbs_error.h"


/* struct to store attributes required for invoking remote shell cmd */
typedef struct remote_shell_cmd {
	char hostname[PBS_MAXHOSTNAME + 1];
	char pipename_append[PIPENAME_MAX_LENGTH];
	int connect_stdin;
} rshell_cmd;

/**
 * @brief
 *	Thread to handle demux from a client.
 *
 * @param[in] pargs - struct rshell_cmd
 *
 * @return - void
 *
 */
void
remote_shell_command_thread(void *pargs)
{
	rshell_cmd      *prs_cmd = NULL;
	int             rc = 0;

	if (pargs == NULL)
		return;

	prs_cmd = (rshell_cmd*)pargs;
	rc = remote_shell_command(prs_cmd->hostname, prs_cmd->pipename_append, prs_cmd->connect_stdin);
	if (rc == -1) {
		fprintf(stderr, "\npbs_demux: failed to fetch remote command output\n");
	}
	ExitThread(0);
}

/**
 * @brief
 *	Handles multiple incoming requests from mom_open_demux client processes, handles the
 *      output.
 *
 * @param[in]  pipename - pipe name for mom_open_demux clients to connect
 * @param[in]  jobid - job id for which demux is needed
 * @param[in]  num_nodes - job's number of nodes
 *
 * @return  int
 * @retval  0 - success
 * @retval  !0 - error
 *
 */
int
handle_np_conn(char *pipename, char *jobid, int num_nodes)
{
	HANDLE                  hPipe = NULL;
	DWORD                   dwRead=0;
	SECURITY_ATTRIBUTES     SecAttrib = {0};
	SECURITY_DESCRIPTOR     SecDesc;
	int                     i = 0;
	rshell_cmd              rs_cmd;
	HANDLE                  *pdemux_thread_handles = NULL;
	char                    hostname[PBS_MAXHOSTNAME + 1] = {'\0'};

	(void)memset(rs_cmd.hostname, '\0', sizeof(rs_cmd.hostname));
	(void)memset(rs_cmd.pipename_append, '\0', sizeof(rs_cmd.pipename_append));

	hPipe = CreateNamedPipe(
		pipename,
		PIPE_ACCESS_DUPLEX,
		PIPE_TYPE_MESSAGE | PIPE_WAIT,
		PIPE_UNLIMITED_INSTANCES,
		0,
		0,
		(DWORD)-1,
		&SecAttrib);
	if (hPipe == INVALID_HANDLE_VALUE || hPipe == NULL) {
		fprintf(stderr, "pipe creation failed!\n");
		return (-1);
	}

	pdemux_thread_handles = (HANDLE *)malloc(num_nodes * sizeof(HANDLE));
	if (pdemux_thread_handles == NULL) {
		fprintf(stderr, "malloc failed!\n");
		close_valid_handle(&(hPipe));
		return (-1);
	}

	for (i = 0; i < num_nodes && hPipe != NULL; i++) {
		(void)InitializeSecurityDescriptor(&SecDesc, SECURITY_DESCRIPTOR_REVISION);
		(void)SetSecurityDescriptorDacl(&SecDesc, TRUE, NULL, TRUE);
		SecAttrib.nLength = sizeof(SECURITY_ATTRIBUTES);
		SecAttrib.lpSecurityDescriptor = &SecDesc;;
		SecAttrib.bInheritHandle = TRUE;

		(void)do_ConnectNamedPipe(hPipe, NULL);
		if (!ReadFile(hPipe, hostname, PBS_MAXHOSTNAME, &dwRead, NULL) ||
			dwRead == 0) {
			DWORD dwErr = GetLastError();
			/* break when client closes the pipe */
			if (dwErr == ERROR_NO_DATA) {
				(void)DisconnectNamedPipe(hPipe);
				pdemux_thread_handles[i] = INVALID_HANDLE_VALUE;
				continue;
			}
		}
		if (dwRead) {
			char pipename_append[PIPENAME_MAX_LENGTH] = {'\0'};
			hostname[ dwRead / sizeof(char) ] = '\0';
			(void)strncpy_s(pipename_append, _countof(pipename_append), jobid, _TRUNCATE);
			(void)strncat_s(pipename_append, _countof(pipename_append), "mom_demux", _TRUNCATE);
			(void)strncat_s(pipename_append, _countof(pipename_append), hostname, _TRUNCATE);
			/*
			 * run remote command
			 * do not redirect stdin
			 */
			(void)strncpy_s(rs_cmd.hostname, _countof(rs_cmd.hostname), hostname, _TRUNCATE);
			(void)strncpy_s(rs_cmd.pipename_append, _countof(rs_cmd.pipename_append), pipename_append, _TRUNCATE);
			rs_cmd.connect_stdin = 0;
			pdemux_thread_handles[i] = (HANDLE)_beginthread(remote_shell_command_thread, 0, &rs_cmd);
		}
		else {
			pdemux_thread_handles[i] = INVALID_HANDLE_VALUE;
		}
		(void)DisconnectNamedPipe(hPipe);
	}
	/* wait forever for demux threads to exit */
	(void)WaitForMultipleObjects(num_nodes, pdemux_thread_handles, TRUE, INFINITE);
	for (i = 0; i < num_nodes; i++)
		close_valid_handle(&(pdemux_thread_handles[i]));
	free(pdemux_thread_handles);
	close_valid_handle(&(hPipe));
	return 0;
}
/**
 * @brief
 *	The main function in C - entry point
 *
 * @param[in]  argc - argument count
 * @param[in]  argv - pointer to argument array
 * @param[in]  envp - pointer to environment values
 *
 * @return  int
 * @retval  0 - success
 * @retval  !0 - error
 */
int
main(int argc, char **argv, char **envp)
{
	char    jobid[PBS_MAXSVRJOBID + 1];
	char    pipename[PIPENAME_MAX_LENGTH] = "\\\\.\\pipe\\";
	char    pipename_append[] = "_pbs_demux_";
	int     rc = 0;	
	int     num_nodes = 0;

	if (argc < 3)
		exit(1);
	/*
	 * argv[1] is jobid
	 * argv[2] tells whether it is an interactive shell
	 * argv[3] is number of nodes for the job
	 */

	(void)strncpy_s(jobid, _countof(jobid), argv[1], _TRUNCATE);	
	num_nodes = atoi(argv[2]);

	(void)strncat_s(pipename, _countof(pipename), jobid, _TRUNCATE);
	(void)strncat_s(pipename, _countof(pipename), pipename_append, _TRUNCATE);
	while (1) {
		if (handle_np_conn(pipename, jobid, num_nodes) == -1) {
			exit(PBSE_SYSTEM);
		}
	}
	exit(0);
}
