redirect.c

Path MPICH-V for MPICH the MPI Implementation

#include "redirect.h"
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <signal.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <errno.h>
#include <pty.h> 
#include <string.h>
#include <sys/wait.h>

#define ms_stdin  0
#define ms_stdout 1
#define ms_stderr 2

//#define DEBUG 1

#ifdef DEBUG
#include <stdio.h>
static FILE *debug_f;
#define debug(toto...) do { \
        fprintf(debug_f, "(%d)%s:%d ", getpid(), __FILE__, __LINE__); \
        fprintf(debug_f, toto); \
        fprintf(debug_f, "\n"); \
        fflush(debug_f); \
  } while(0)
#else
#define debug(toto...) do {} while(0)
#endif

static pid_t PID;

struct connector {
  int src;
  int dst;
};

static void abort_run(void)
{
#ifdef DEBUG
  debug("aborting run");
  fclose(debug_f);
#endif

  if(PID != 0)
    kill(PID, SIGTERM);

  _exit(0);
}

static void on_child(int s)
{
  pid_t p;
  debug("chld signal received");
  p = wait(&s);
  if(WIFEXITED(s))
    debug("%d has exited with status %d", p, WEXITSTATUS(s));
  if(WIFSIGNALED(s))
    debug("%d was signaled with signal %d", p, WTERMSIG(s));
  if(WIFSTOPPED(s))
    debug("%d was stopped by signal %d", p, WSTOPSIG(s));
  if(!WIFEXITED(s) && !WIFSIGNALED(s) && !WIFSTOPPED(s))
    debug("%d was neither exited, signaled or stopped...", p);
  abort_run();
}

static void connect_master(int master, struct sockaddr_in *remote)
{
  int fd;
  struct connector connectors[2];
  struct pollfd polls[2];
  int p;
  int nr;
#define NR 4096
  char buf[NR];

  debug("connecting the master");

  fd = socket(AF_INET, SOCK_STREAM, 0);
  if(fd < 0)
    for(;;) pause(); /**< wait for the child to die. No stdio! */

  if(connect(fd, (struct sockaddr*)remote, sizeof(struct sockaddr_in)) < 0)
    {
      debug("unable to connect to %s:%d (stdio)", inet_ntoa(remote->sin_addr), ntohs(remote->sin_port));
      for(;;) pause(); /**< wait for the child to die. No stdio! */
    }

  debug("connected to %s:%d (stdio), dupplicating stdio", inet_ntoa(remote->sin_addr), ntohs(remote->sin_port));

  connectors[0].src = fd;
  connectors[0].dst = master;

  connectors[1].src = master;
  connectors[1].dst = fd;

  for(p = 0; p < 2; p++)
    {
      polls[p].fd = connectors[p].src;
      polls[p].events = POLLIN;
    }

  for(;;)
    {
      p = poll(polls, 2, -1);

      if( (p < 0) && (errno != EINTR) )
	abort_run();
      if( (p == 0) )
	abort_run();

      for(p = 0; p < 2; p++)
	{
	  if( polls[p].revents & POLLIN )
	    {
	      nr = read(connectors[p].src, buf, NR);
	      debug("read %d from %d", nr, p);
	      if(nr < 0)
		{
		  if( errno == EINTR )
		    continue;
		  else
		    abort_run();
		}
	      if(nr == 0)
		abort_run();

	      write(connectors[p].dst, buf, nr);
	    }
	}
    }
#undef NR
}

void redirect_stdio(char *runtime_ip, short stdio_port, short stder_port)
{
  int master;
  struct sockaddr_in remote;
  struct hostent *h;

#ifdef DEBUG
  debug_f = fopen("/tmp/redirect.err", "a");
  debug("==== redirect_stdio  ====");
#endif

  h = gethostbyname(runtime_ip);
  if(h == NULL)
    {
      herror("unable to resolve runtime_ip");
      return;
    }
  remote.sin_family = AF_INET;
  memcpy(&remote.sin_addr, h->h_addr, h->h_length);
 
  signal(SIGCHLD, on_child);
  PID = forkpty(&master, NULL, NULL, NULL);

  debug("forkpty => %d/%d", PID, master);
  if(PID == 0)
    {
      int fd;

      debug("slave:%d/%d/%d", isatty(0), isatty(1), isatty(2));

      /* in the slave */
      signal(SIGCHLD, SIG_DFL);
      close(ms_stderr);
      
      fd = socket(AF_INET, SOCK_STREAM, 0);
      if(fd < 0)
	return;

      debug("stderr fd = %d", fd);

      remote.sin_port = htons(stder_port);
      
      if(connect(fd, (struct sockaddr*)&remote, sizeof(struct sockaddr_in)) < 0)
	{
	  debug("unable to connect stderr to %s:%d", inet_ntoa(remote.sin_addr), ntohs(remote.sin_port));
	  close(fd);
	  return;
	}

      dup2(fd, ms_stderr);
      
      debug("stdio connected to tty, stderr connected to %s:%d", inet_ntoa(remote.sin_addr), ntohs(remote.sin_port));
      /* stdout and stdin are redirected to master by forkpty */
#ifdef DEBUG
      fclose(debug_f);
#endif
      return;
    }

  close(ms_stdin);
  close(ms_stdout);
  close(ms_stderr);

  remote.sin_port = htons(stdio_port);  
  connect_master(master, &remote);
}