From 5426c57ad4cd9de90256af79008f3f0ec048d1eb Mon Sep 17 00:00:00 2001
From: Mark Wooding <mdw@distorted.org.uk>
Date: Wed, 27 Sep 2017 23:49:01 +0100
Subject: [PATCH] pkstream/pkstream.c: Introduce helper functions to fiddle
 with addresses.

The remaining places where pieces of addresses are fiddled with
directly, outside of these new functions, are:

  * in `doaccept', where we continue inspect the peer address to see if
    it's a wildcard, because we'll handle this in a very different way
    later; and

  * in `parseaddr', which needs to fill in addresses and port numbers.
---
 pkstream/pkstream.c | 44 ++++++++++++++++++++++++++++++++++++--------
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/pkstream/pkstream.c b/pkstream/pkstream.c
index 4df37e7e..fa924402 100644
--- a/pkstream/pkstream.c
+++ b/pkstream/pkstream.c
@@ -99,6 +99,35 @@ static int nonblockify(int fd)
 static int cloexec(int fd)
   { return (fdflags(fd, 0, 0, FD_CLOEXEC, FD_CLOEXEC)); }
 
+static socklen_t addrsz(const addr *a)
+{
+  switch (a->sa.sa_family) {
+    case AF_INET: return sizeof(a->sin);
+    default: abort();
+  }
+}
+
+static const char *addrstr(const addr *a)
+{
+  static char buf[128];
+  socklen_t n = sizeof(buf);
+
+  if (getnameinfo(&a->sa, addrsz(a), buf, n, 0, 0, NI_NUMERICHOST))
+    return ("<addrstr failed>");
+  return (buf);
+}
+
+static int addreq(const addr *a, const addr *b)
+{
+  if (a->sa.sa_family != b->sa.sa_family) return (0);
+  switch (a->sa.sa_family) {
+    case AF_INET:
+      return (a->sin.sin_addr.s_addr == b->sin.sin_addr.s_addr);
+    default:
+      abort();
+  }
+}
+
 static void initaddr(addr *a)
 {
   a->sin.sin_family = AF_INET;
@@ -238,9 +267,8 @@ static void doaccept(int fd_s, unsigned mode, void *p)
     moan("couldn't accept incoming connection: %s", strerror(errno));
     return;
   }
-  if (cw.peer.sin.sin_addr.s_addr != INADDR_ANY &&
-      cw.peer.sin.sin_addr.s_addr != a.sin.sin_addr.s_addr) {
-    moan("rejecting connection from %s", inet_ntoa(a.sin.sin_addr));
+  if (cw.peer.sin.sin_addr.s_addr != INADDR_ANY && !addreq(&a, &cw.peer)) {
+    moan("rejecting connection from %s", addrstr(&a));
     close(fd); return;
   }
   if (nonblockify(fd) || cloexec(fd)) {
@@ -259,7 +287,7 @@ static void dolisten(void)
 
   if ((fd = socket(PF_INET, SOCK_STREAM, 0)) < 0 ||
       setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) ||
-      bind(fd, &cw.me.sa, sizeof(cw.me.sin)) ||
+      bind(fd, &cw.me.sa, addrsz(&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);
@@ -401,8 +429,8 @@ int main(int argc, char *argv[])
     parseaddr(connhost, 0, paf_parse, &tmpaddr);
     if ((fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0 ||
 	(bindhost &&
-	 bind(fd, &bindaddr.sa, sizeof(bindaddr.sin))) ||
-	connect(fd, &tmpaddr.sa, sizeof(tmpaddr.sin)))
+	 bind(fd, &bindaddr.sa, addrsz(&bindaddr))) ||
+	connect(fd, &tmpaddr.sa, addrsz(&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));
@@ -414,11 +442,11 @@ int main(int argc, char *argv[])
       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, &tmpaddr.sa, sizeof(tmpaddr.sin)))
+      bind(fd_udp, &tmpaddr.sa, addrsz(&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, &tmpaddr.sa, sizeof(tmpaddr.sin)))
+  if (connect(fd_udp, &tmpaddr.sa, addrsz(&tmpaddr)))
     die(1, "couldn't set up UDP socket: %s", strerror(errno));
 
   if (bindsvc) dolisten();
-- 
2.11.0