/* This file is part of GNU Pies testsuite.
   Copyright (C) 2019-2020 Sergey Poznyakoff

   GNU Pies is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3, or (at your option)
   any later version.

   GNU Pies 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 General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with GNU Pies.  If not, see <http://www.gnu.org/licenses/>. */

#include <config.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <poll.h>
#include <errno.h>
#include <libpies.h>
#include <iobuf.h>

typedef struct netcat_server netcat_server_t;
typedef struct netcat_stream netcat_stream_t;
typedef ssize_t (*netcat_stream_io_t) (netcat_server_t *srv);
typedef int (*netcat_stream_disconnect_t) (netcat_server_t *srv);

enum { IN, OUT };

struct netcat_server
{
  char const *id;
  netcat_server_t *next, *prev;
  struct pollfd *pollfd;
  int state;
  netcat_stream_disconnect_t disconnect;
  netcat_server_t *peer;
  struct iobuf buf[2];  
};

static netcat_server_t *server_head, *server_tail;

netcat_server_t *
netcat_server_create (char const *id, struct pollfd *pfd, int state,
		      netcat_stream_disconnect_t dis, netcat_server_t *peer)
{
  netcat_server_t *srv = grecs_zalloc (sizeof (*srv));
  srv->id = id;
  srv->pollfd = pfd;
  srv->state = state;
  srv->pollfd->events |= srv->state;
  srv->disconnect = dis;
  srv->peer = peer;
  if (peer)
    srv->peer->peer = srv;
  iobuf_reset (&srv->buf[IN]);
  iobuf_reset (&srv->buf[OUT]);

  srv->prev = server_tail;
  srv->next = NULL;
  if (server_tail)
    server_tail->next = srv;
  else
    server_head = srv;
  server_tail = srv;

  return srv;
}

void
netcat_server_remove (netcat_server_t *srv)
{
  netcat_server_t *p;
  if ((p = srv->next) != NULL)
    p->prev = srv->prev;
  else
    server_tail = srv->prev;
  if ((p = srv->prev) != NULL)
    p->next = srv->next;
  else
    server_head = srv->next;
  if (srv->peer)
    srv->peer->peer = NULL;
  free (srv);
}

/*
net_reader
1. If there is free space in the IN buffer, fill it with the data read
from the fd and return
2. Otherwise, clear the POLLIN bit.

stdout_writer (peer -> net_reader)
1. If OUT buffer is empty
  1.1 If the peer's IN buffer is not empty, transfer data from it
  1.2 Else raise the peers POLLIN bit and return.
2. Write as much as possible from OUT to fd

stdin_reader (same as net_reader)

net_writer (peer -> stdin_reader)  same as stdout_writer
*/

static inline int
peer_is_state (netcat_server_t *srv, int state)
{
  return srv->peer && (srv->peer->state & state);
}

static void
netcat_stream_disconnect (netcat_server_t *srv, int mask)
{
  srv->disconnect (srv);
  srv->state &= ~mask;
  srv->pollfd->events &= ~mask;
  if (srv->pollfd->events == 0)
    srv->pollfd->fd = -1;
}

ssize_t
netcat_stream_read (netcat_server_t *srv)
{
  ssize_t n;

  if (iobuf_avail_size (&srv->buf[IN]))
    {
      n = iobuf_fill (&srv->buf[IN], srv->pollfd->fd);
      if (n == -1)
	{
	  if (errno == EINTR)
	    return 0;
	  grecs_error (NULL, errno, "%s: read", srv->id);
	  netcat_stream_disconnect (srv, POLLIN);
	  return -1;
	}
      if (n == 0 || !peer_is_state (srv, POLLOUT))
	{
	  /* No more input is expected || needed */
	  netcat_stream_disconnect (srv, POLLIN);

	  if (srv->peer)
	    netcat_stream_disconnect (srv->peer, POLLOUT);
	  return 0;
	}
      else
	srv->peer->pollfd->events |= POLLOUT;
    }
  else
    {
      srv->pollfd->events &= ~POLLIN;
      n = 0; /* catch timeout? */
    }
  return n;
}

ssize_t
netcat_stream_write (netcat_server_t *srv)
{
  ssize_t n;
  
  if (iobuf_data_size (&srv->buf[OUT]) == 0)
    {
      if (!peer_is_state (srv, POLLIN))
	{
	  // shutdown write end
	  netcat_stream_disconnect (srv, POLLOUT);
	  return -1;
	}
      if (iobuf_copy (&srv->buf[OUT], &srv->peer->buf[IN]) == 0)
	{
	  srv->peer->pollfd->events |= POLLIN;
	  srv->pollfd->events &= ~POLLOUT;
	  return 0;
	}
    }
  n = iobuf_flush (&srv->buf[OUT], srv->pollfd->fd);
  if (n == -1)
    {
      if (errno == EINTR)
	return 0;
      grecs_error (NULL, errno, "%s: write", srv->id);
      netcat_stream_disconnect (srv, POLLOUT);
      return -1;
    }
  if (n == 0)
    {
      netcat_stream_disconnect (srv, POLLOUT);
      return -1;
    }
  return 0;
}

int
disconnect_in (netcat_server_t *srv)
{
  return shutdown (srv->pollfd->fd, SHUT_RD);
}

int
disconnect_out (netcat_server_t *srv)
{
  return shutdown (srv->pollfd->fd, SHUT_WR);
}

int
disconnect_stdin (netcat_server_t *srv)
{
  return close (srv->pollfd->fd);
}

int
disconnect_stdout (netcat_server_t *srv)
{
  close (srv->pollfd->fd);
  exit (0);
}

void
fd_write (int fd, char const *str, size_t len)
{
  while (len)
    {
      ssize_t n = write (fd, str, len);
      if (n == -1)
	{
	  perror ("socket write");
	  exit (1);
	}
      if (n == 0)
	{
	  fprintf (stderr, "zero write\n");
	  exit (1);
	}
      len -= n;
      str += n;
    }
}

void
fd_writeln (int fd, char const *str)
{
  fd_write (fd, str, strlen (str));
  fd_write (fd, "\r\n", 2);
}

int
fd_getc (int fd)
{
  char c;
  ssize_t n = read (fd, &c, 1);
  if (n == -1)
    {
      perror ("socket read");
      exit (1);
    }
  if (n == 0)
    c = EOF;
  return c;
}

static void
tcpmux_init (int fd, char const *service)
{
  int c;
  
  fd_writeln (fd, service);
  
  if (strcmp (service, "help") == 0)
    {
      while ((c = fd_getc (fd)) != EOF)
	{
	  fputc (c, stdout);
	}
      close (fd);
      exit (0);
    }

  c = fd_getc (fd);
  if (c == 0)
    {
      fprintf (stderr, "socket read: unexpected eof\n");
      exit (1);
    }
  if (c == '+')
    {
      while ((c = fd_getc (fd)) != '\n')
	{
	  if (c == EOF)
	    {
	      fprintf (stderr, "socket read: unexpected eof\n");
	      exit (1);
	    }
	}
    }
  else
    {
      fprintf (stderr, "service rejected: ");
      do
	{
	  if (c != '\r')
	    fputc (c, stderr);
	}
      while ((c = fd_getc (fd)) != 0 && c != '\n');
      fputc ('\n', stderr);
      exit (1);
    }
}

static int
netcat (char const *urlstr, char const *tcpmux_service)
{
  int fd;
  struct pies_url *url;
  struct pollfd pfd[3];
  int nfd = sizeof (pfd) / sizeof (pfd[0]);
  netcat_server_t *srvin, *srvout, *srv;
  
  if (pies_url_create (&url, urlstr))
    {
      perror (urlstr);
      return 64;
    }

  fd = url_connect (url, NULL);
  if (fd == -1)
    return 1;

  if (tcpmux_service)
    tcpmux_init (fd, tcpmux_service);
  
  pfd[0].fd = 0;
  pfd[0].events = 0;
  srvin = netcat_server_create ("stdin", &pfd[0], POLLIN, disconnect_stdin, NULL);

  pfd[1].fd = 1;
  pfd[1].events = 0;
  srvout = netcat_server_create ("stdout", &pfd[1], POLLOUT, disconnect_stdout, NULL);

  pfd[2].fd = fd;
  pfd[2].events = 0;

  netcat_server_create ("netread", &pfd[2], POLLIN, disconnect_in, srvout);
  netcat_server_create ("netwrite", &pfd[2], POLLOUT, disconnect_out, srvin);

  while (server_head)
    {
      ssize_t n = poll (pfd, nfd, -1);
      if (n == -1)
	{
	  if (errno != EINTR)
	    grecs_error (NULL, errno, "poll");
	  continue;
	}
      for (srv = server_head; srv; )
	{
	  netcat_server_t *next = srv->next;

	  if ((srv->pollfd->revents & srv->state) & POLLIN)
	    netcat_stream_read (srv);
	  if ((srv->pollfd->revents & srv->state) & POLLOUT)
	    netcat_stream_write (srv);
	  if ((srv->state & POLLOUT) && (srv->pollfd->revents & POLLHUP))
	    {
	      //grecs_error (NULL, 0, "HUP on %s", srv->id);
	      netcat_stream_disconnect (srv, srv->state);
	    }
	  if (srv->state == 0 || srv->pollfd->fd == -1)
	    {
	      netcat_server_t *peer = srv->peer;
	      if (peer && peer->pollfd->events == 0)
		{
		  netcat_stream_disconnect (peer, peer->state);
		  netcat_server_remove (peer);
		}
	      netcat_server_remove (srv);
	    }
	  srv = next;
	}
    }
  return 0;
}

static void
redirect (int sfd, char const *name)
{
  int fd;

  fd = open (name, sfd ? (O_CREAT | O_TRUNC | O_WRONLY) : O_RDONLY, 0644);
  if (fd == -1)
    {
      perror (name);
      exit (1);
    }
  if (dup2 (fd, sfd) == -1)
    {
      perror ("dup2");
      exit (1);
    }
}

static void
usage (FILE *fp)
{
  fprintf (fp, "usage: nt [-i IFILE] [-o OFILE] [-t SERVICE] URL\n");
  fprintf (fp, "Reads data from stdin (or IFILE) and sends them to URL.\n");
  fprintf (fp, "Reads replies from URL and sends them to stdout (or OFILE).\n");
  fprintf (fp, "\nOPTIONS\n\n");
  fprintf (fp, "   -t SERVICE            use TCPMUX service\n");
  fprintf (fp, "   -i IFILE              read input from IFILE\n");
  fprintf (fp, "   -o OFILE              write output to OFILE\n");
}

int
main (int argc, char **argv)
{
  int c;
  char const *tcpmux_service = NULL;
  
  while ((c = getopt (argc, argv, "i:o:t:")) != EOF)
    {
      switch (c)
	{
	case 'i':
	  redirect (0, optarg);
	  break;

	case 'o':
	  redirect (1, optarg);
	  break;

	case 't':
	  tcpmux_service = optarg;
	  break;
	  
	default:
	  exit (64);
	}
    }

  argc -= optind;
  argv += optind;

  if (argc != 1)
    {
      usage (stderr);
      exit (64);
    }

  return netcat (argv[0], tcpmux_service);
}