pkstream/pkstream.c: Set a flag if we're listening.

[tripe] / pkstream / pkstream.c

/* -*-c-*-
 *
 * Forwarding UDP packets over a stream
 *
 * (c) 2003 Straylight/Edgeware
 */

/*----- Licensing notice --------------------------------------------------*
 *
 * This file is part of Trivial IP Encryption (TrIPE).
 *
 * TrIPE 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 of the License, or (at your
 * option) any later version.
 *
 * TrIPE 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 TrIPE.  If not, see <https://www.gnu.org/licenses/>.
 */

/*----- Header files ------------------------------------------------------*/

#include "config.h"

#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/uio.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>

#include <mLib/alloc.h>
#include <mLib/bits.h>
#include <mLib/dstr.h>
#include <mLib/fdflags.h>
#include <mLib/mdwopt.h>
#include <mLib/quis.h>
#include <mLib/report.h>
#include <mLib/sel.h>
#include <mLib/selpk.h>

#include "util.h"

/*----- Data structures ---------------------------------------------------*/

typedef struct pk {
  struct pk *next;			/* Next packet in the chain */
  octet *p, *o;				/* Buffer start and current posn */
  size_t n;				/* Size of packet remaining */
} pk;

typedef struct pkstream {
  unsigned f;				/* Flags... */
#define PKF_FULL 1u			/*   Buffer is full: stop reading */
  sel_file r, w;			/* Read and write selectors */
  pk *pks, **pk_tail;			/* Packet queue */
  size_t npk, szpk;			/* Number and size of data */
  selpk p;				/* Packet parser */
} pkstream;

typedef struct connwait {
  unsigned f;				/* Various flags */
#define cwf_port 1u			/*   Port is defined => listen */
  sel_file a;				/* Selector */
  struct sockaddr_in me, peer;	       /* Who I'm meant to be; who peer is */
} connwait;

/*----- Static variables --------------------------------------------------*/

static sel_state sel;
static connwait cw;
static int fd_udp;
static size_t pk_nmax = 128, pk_szmax = 1024*1024;

/*----- Main code ---------------------------------------------------------*/

static int nonblockify(int fd)
  { return (fdflags(fd, O_NONBLOCK, O_NONBLOCK, 0, 0)); }

static int cloexec(int fd)
  { return (fdflags(fd, 0, 0, FD_CLOEXEC, FD_CLOEXEC)); }

static void initaddr(struct sockaddr_in *sin)
{
  sin->sin_family = AF_INET;
  sin->sin_addr.s_addr = INADDR_ANY;
  sin->sin_port = 0;
}

static void dolisten(void);

static void doclose(pkstream *p)
{
  pk *pk, *ppk;
  close(p->w.fd);
  close(p->p.reader.fd);
  selpk_destroy(&p->p);
  if (!(p->f&PKF_FULL)) sel_rmfile(&p->r);
  if (p->npk) sel_rmfile(&p->w);
  for (pk = p->pks; pk; pk = ppk) {
    ppk = pk->next;
    xfree(pk->p);
    xfree(pk);
  }
  xfree(p);
  if (cw.f&cwf_port) dolisten();
  else exit(0);
}

static void rdtcp(octet *b, size_t sz, pkbuf *pk, size_t *k, void *vp)
{
  pkstream *p = vp;
  size_t pksz;

  if (!sz) { doclose(p); return; }
  pksz = LOAD16(b);
  if (pksz + 2 == sz) {
    DISCARD(write(fd_udp, b + 2, pksz));
    selpk_want(&p->p, 2);
  } else {
    selpk_want(&p->p, pksz + 2);
    *k = sz;
  }
}

static void wrtcp(int fd, unsigned mode, void *vp)
{
#define NPK 16
  struct iovec iov[NPK];
  pkstream *p = vp;
  size_t i;
  ssize_t n;
  pk *pk, *ppk;

  for (i = 0, pk = p->pks; i < NPK && pk; i++, pk = pk->next) {
    iov[i].iov_base = pk->o;
    iov[i].iov_len = pk->n;
  }

  if ((n = writev(fd, iov, i)) < 0) {
    if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) return;
    moan("couldn't write to TCP socket: %s", strerror(errno));
    doclose(p);
    return;
  }

  p->szpk -= n;
  for (pk = p->pks; n && pk; pk = ppk) {
    ppk = pk->next;
    if (pk->n <= n) {
      p->npk--;
      n -= pk->n;
      xfree(pk->p);
      xfree(pk);
    } else {
      pk->n -= n;
      pk->o += n;
      break;
    }
  }
  p->pks = pk;
  if (!pk) { p->pk_tail = &p->pks; sel_rmfile(&p->w); }
  if ((p->f&PKF_FULL) && p->npk < pk_nmax && p->szpk < pk_szmax)
    { p->f &= ~PKF_FULL; sel_addfile(&p->r); }
}

static void rdudp(int fd, unsigned mode, void *vp)
{
  octet buf[65536];
  ssize_t n;
  pkstream *p = vp;
  pk *pk;

  if ((n = read(fd, buf, sizeof(buf))) < 0) {
    if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
      return;
    moan("couldn't read from UDP socket: %s", strerror(errno));
    return;
  }
  pk = xmalloc(sizeof(*pk));
  pk->next = 0;
  pk->p = xmalloc(n + 2);
  STORE16(pk->p, n);
  memcpy(pk->p + 2, buf, n);
  pk->o = pk->p;
  pk->n = n + 2;
  *p->pk_tail = pk;
  p->pk_tail = &pk->next;
  if (!p->npk) sel_addfile(&p->w);
  sel_force(&p->w);
  p->npk++;
  p->szpk += n + 2;
  if (p->npk >= pk_nmax || p->szpk >= pk_szmax)
    { sel_rmfile(&p->r); p->f |= PKF_FULL; }
}

static void dofwd(int fd_in, int fd_out)
{
  pkstream *p = xmalloc(sizeof(*p));
  sel_initfile(&sel, &p->r, fd_udp, SEL_READ, rdudp, p);
  sel_initfile(&sel, &p->w, fd_out, SEL_WRITE, wrtcp, p);
  selpk_init(&p->p, &sel, fd_in, rdtcp, p);
  selpk_want(&p->p, 2);
  p->pks = 0;
  p->pk_tail = &p->pks;
  p->npk = p->szpk = 0;
  p->f = 0;
  sel_addfile(&p->r);
}

static void doaccept(int fd_s, unsigned mode, void *p)
{
  int fd;
  struct sockaddr_in sin;
  socklen_t sz = sizeof(sin);

  if ((fd = accept(fd_s, (struct sockaddr *)&sin, &sz)) < 0) {
    if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) return;
    moan("couldn't accept incoming connection: %s", strerror(errno));
    return;
  }
  if (cw.peer.sin_addr.s_addr != INADDR_ANY &&
      cw.peer.sin_addr.s_addr != sin.sin_addr.s_addr) {
    moan("rejecting connection from %s", inet_ntoa(sin.sin_addr));
    close(fd); return;
  }
  if (nonblockify(fd) || cloexec(fd)) {
    moan("couldn't accept incoming connection: %s", strerror(errno));
    close(fd); return;
  }
  dofwd(fd, fd);
  close(fd_s);
  sel_rmfile(&cw.a);
}

static void dolisten(void)
{
  int fd;
  int opt = 1;

  if ((fd = socket(PF_INET, SOCK_STREAM, 0)) < 0 ||
      setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) ||
      bind(fd, (struct sockaddr *)&cw.me, sizeof(cw.me)) ||
      listen(fd, 1) || nonblockify(fd) || cloexec(fd))
    die(1, "couldn't set up listening socket: %s", strerror(errno));
  sel_initfile(&sel, &cw.a, fd, SEL_READ, doaccept, 0);
  sel_addfile(&cw.a);
}

#define paf_parse 1u
static void parseaddr(const char *host, const char *svc, unsigned f,
		      struct sockaddr_in *sin)
{
  char *alloc = 0, *sep;
  struct hostent *h;
  struct servent *s;
  char *qq;
  unsigned long n;

  if (f&paf_parse) {
    alloc = xstrdup(host);
    if ((sep = strchr(alloc, ':')) == 0)
      die(1, "missing port number in address `%s'", host);
    host = alloc; *sep = 0; svc = sep + 1;
  }

  if (host) {
    if ((h = gethostbyname(host)) == 0) die(1, "unknown host `%s'", host);
    memcpy(&sin->sin_addr, h->h_addr, sizeof(sin->sin_addr));
  }

  if (svc) {
    if ((n = strtoul(svc, &qq, 0)) > 0 && !*qq && n <= 0xffff)
      sin->sin_port = htons(n);
    else if ((s = getservbyname(svc, "tcp")) != 0)
      sin->sin_port = s->s_port;
    else
      die(1, "bad service name/number `%s'", svc);
  }

  xfree(alloc);
}

static void usage(FILE *fp)
{
  pquis(fp,
	"Usage: $ [-l PORT] [-b ADDR] [-p ADDR] [-c ADDR:PORT]\n\
	ADDR:PORT ADDR:PORT\n");
}

static void version(FILE *fp)
  { pquis(fp, "$, tripe version " VERSION "\n"); }

static void help(FILE *fp)
{
  version(fp);
  fputc('\n', fp);
  usage(fp);
  fputs("\n\
Options:\n\
\n\
-h, --help		Display this help text.\n\
-v, --version		Display version number.\n\
-u, --usage		Display pointless usage message.\n\
\n\
-l, --listen=PORT	Listen for connections to TCP PORT.\n\
-p, --peer=ADDR		Only accept connections from IP ADDR.\n\
-b, --bind=ADDR		Bind to ADDR before connecting.\n\
-c, --connect=ADDR:PORT	Connect to IP ADDR, TCP PORT.\n\
\n\
Forwards UDP packets over a reliable stream.  By default, uses stdin and\n\
stdout; though it can use TCP sockets instead.\n\
", fp);
}

int main(int argc, char *argv[])
{
  unsigned f = 0;
  const char *bindhost = 0, *bindsvc = 0, *peerhost = 0;
  struct sockaddr_in bindaddr;
  const char *connhost = 0;
  struct sockaddr_in tmpaddr;
  int fd = -1;
  int len = 65536;

#define f_bogus 1u

  cw.f = 0;

  ego(argv[0]);
  sel_init(&sel);
  for (;;) {
    static struct option opt[] = {
      { "help",			0,		0,	'h' },
      { "version",		0,		0,	'v' },
      { "usage",		0,		0,	'u' },
      { "listen",		OPTF_ARGREQ,	0,	'l' },
      { "peer",			OPTF_ARGREQ,	0,	'p' },
      { "bind",			OPTF_ARGREQ,	0,	'b' },
      { "connect",		OPTF_ARGREQ,	0,	'c' },
      { 0,			0,		0,	0 }
    };
    int i;

    i = mdwopt(argc, argv, "hvul:p:b:c:", opt, 0, 0, 0);
    if (i < 0)
      break;
    switch (i) {
      case 'h': help(stdout); exit(0);
      case 'v': version(stdout); exit(0);
      case 'u': usage(stdout); exit(0);
      case 'l': bindsvc = optarg; break;
      case 'p': peerhost = optarg; break;
      case 'b': bindhost = optarg; break;
      case 'c': connhost = optarg; break;
      default: f |= f_bogus; break;
    }
  }
  if (optind + 2 != argc || (f&f_bogus)) { usage(stderr); exit(1); }

  if (bindhost && !bindsvc && !connhost)
    die(1, "bind addr only makes sense when listening or connecting");
  if (peerhost && !bindsvc)
    die(1, "peer addr only makes sense when listening");
  if (bindsvc && connhost)
    die(1, "can't listen and connect");

  if (bindhost || bindsvc) {
    initaddr(&bindaddr);
    if (!bindsvc) parseaddr(bindhost, 0, 0, &bindaddr);
    else {
      initaddr(&cw.me);
      parseaddr(bindhost, bindsvc, 0, &cw.me);
      cw.f |= cwf_port;
    }
  }

  initaddr(&cw.peer);
  if (peerhost) parseaddr(peerhost, 0, 0, &cw.peer);

  if (connhost) {
    initaddr(&tmpaddr);
    parseaddr(connhost, 0, paf_parse, &tmpaddr);
    if ((fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0 ||
	(bindhost &&
	 bind(fd, (struct sockaddr *)&bindaddr, sizeof(bindaddr))) ||
	connect(fd, (struct sockaddr *)&tmpaddr, sizeof(tmpaddr)))
      die(1, "couldn't connect to TCP server: %s", strerror(errno));
    if (nonblockify(fd) || cloexec(fd))
      die(1, "couldn't connect to TCP server: %s", strerror(errno));
  }

  initaddr(&tmpaddr);
  parseaddr(argv[optind], 0, paf_parse, &tmpaddr);
  if ((fd_udp = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) < 0 ||
      nonblockify(fd_udp) || cloexec(fd_udp) ||
      setsockopt(fd_udp, SOL_SOCKET, SO_RCVBUF, &len, sizeof(len)) ||
      setsockopt(fd_udp, SOL_SOCKET, SO_SNDBUF, &len, sizeof(len)) ||
      bind(fd_udp, (struct sockaddr *)&tmpaddr, sizeof(tmpaddr)))
    die(1, "couldn't set up UDP socket: %s", strerror(errno));
  initaddr(&tmpaddr);
  parseaddr(argv[optind + 1], 0, paf_parse, &tmpaddr);
  if (connect(fd_udp, (struct sockaddr *)&tmpaddr, sizeof(tmpaddr)))
    die(1, "couldn't set up UDP socket: %s", strerror(errno));

  if (bindsvc) dolisten();
  else if (connhost) dofwd(fd, fd);
  else dofwd(STDIN_FILENO, STDOUT_FILENO);

  for (;;) {
    if (sel_select(&sel) && errno != EINTR)
      die(1, "select failed: %s", strerror(errno));
  }
  return (0);
}

/*----- That's all, folks -------------------------------------------------*/