Make netcat work with nommu and factor out poll() loop into net/net.c.

3 files changed, 114 insertions, 100 deletions

diff --git a/lib/lib.h b/lib/lib.h
index 69692b73..2afe558b 100644
--- a/lib/lib.h
+++ b/lib/lib.h
@@ -281,6 +281,7 @@ void xsetsockopt(int fd, int level, int opt, void *val, socklen_t len);
 int xconnect(char *host, char *port, int family, int socktype, int protocol,
   int flags);
 int xpoll(struct pollfd *fds, int nfds, int timeout);
+int pollinate(int in1, int in2, int out1, int out2, int timeout, int shutdown_timeout);
 
 // password.c
 int get_salt(char *salt, char * algo);
diff --git a/lib/net.c b/lib/net.c
index 2e72b268..1a46a5a9 100644
--- a/lib/net.c
+++ b/lib/net.c
@@ -55,3 +55,40 @@ int xpoll(struct pollfd *fds, int nfds, int timeout)
     } else return i;
   }
 }
+
+// Loop forwarding data from in1 to out1 and in2 to out2, handling
+// half-connection shutdown. timeouts return if no data for X miliseconds.
+// Returns 0: both closed, 1 shutdown_timeout, 2 timeout
+int pollinate(int in1, int in2, int out1, int out2, int timeout, int shutdown_timeout)
+{
+  struct pollfd pollfds[2];
+  int i, pollcount = 2;
+
+  memset(pollfds, 0, 2*sizeof(struct pollfd));
+  pollfds[0].events = pollfds[1].events = POLLIN;
+  pollfds[0].fd = in1;
+  pollfds[1].fd = in2;
+
+  // Poll loop copying data from each fd to the other one.
+  for (;;) {
+    if (!xpoll(pollfds, pollcount, timeout)) return pollcount;
+
+    for (i=0; i<pollcount; i++) {
+      if (pollfds[i].revents & POLLIN) {
+        int len = read(pollfds[i].fd, libbuf, sizeof(libbuf));
+        if (len<1) pollfds[i].revents = POLLHUP;
+        else xwrite(i ? out2 : out1, libbuf, len);
+      }
+      if (pollfds[i].revents & POLLHUP) {
+        // Close half-connection.  This is needed for things like
+        // "echo GET / | netcat landley.net 80"
+        // Note that in1 closing triggers timeout, in2 returns now.
+        if (i) {
+          shutdown(pollfds[0].fd, SHUT_WR);
+          pollcount--;
+          timeout = shutdown_timeout;
+        } else return 0;
+      }
+    }
+  }
+}
diff --git a/toys/net/netcat.c b/toys/net/netcat.c
index 1c75eb26..24420d6b 100644
--- a/toys/net/netcat.c
+++ b/toys/net/netcat.c
@@ -5,7 +5,7 @@
  * TODO: udp, ipv6, genericize for telnet/microcom/tail-f
 
 USE_NETCAT(OLDTOY(nc, netcat, TOYFLAG_USR|TOYFLAG_BIN))
-USE_NETCAT(NEWTOY(netcat, USE_NETCAT_LISTEN("^tlL")"w#p#s:q#f:", TOYFLAG_BIN))
+USE_NETCAT(NEWTOY(netcat, USE_NETCAT_LISTEN("^tlL")"w#W#p#s:q#f:"USE_NETCAT_LISTEN("[!tlL][!Lw]"), TOYFLAG_BIN))
 
 config NETCAT
   bool "netcat"
@@ -15,9 +15,10 @@ config NETCAT
 
     -f	use FILENAME (ala /dev/ttyS0) instead of network
     -p	local port number
-    -q	SECONDS quit this many seconds after EOF on stdin.
-    -s	local ipv4 address
-    -w	SECONDS timeout for connection
+    -q	quit SECONDS after EOF on stdin, even if stdout hasn't closed yet.
+    -s	local source address
+    -w	SECONDS timeout to establish connection
+    -W	SECONDS timeout for idle connection
 
     Use "stty 115200 -F /dev/ttyS0 && stty raw -echo -ctlecho" with
     netcat -f to connect to a serial port.
@@ -26,28 +27,20 @@ config NETCAT_LISTEN
   bool "netcat server options (-let)"
   default y
   depends on NETCAT
-  depends on TOYBOX_FORK
   help
-    usage: netcat [-lL COMMAND...]
+    usage: netcat [-t] [-lL COMMAND...]
 
-    -l	listen for one incoming connection.
-    -L	listen for multiple incoming connections (server mode).
+    -l	listen for one incoming connection
+    -L	listen for multiple incoming connections (server mode)
+    -t	allocate tty (must come before -l or -L)
 
-    The command line after -l or -L is executed to handle each incoming
-    connection. If none, the connection is forwarded to stdin/stdout.
+    The command line after -l or -L is executed (as a child process) to handle
+    each incoming connection. If blank -l waits for a connection and forwards
+    it to stdin/stdout. If no -p specified, -l prints port it bound to and
+    backgrounds itself (returning immediately).
 
     For a quick-and-dirty server, try something like:
     netcat -s 127.0.0.1 -p 1234 -tL /bin/bash -l
-
-config NETCAT_LISTEN_TTY
-  bool
-  default y
-  depends on NETCAT_LISTEN
-  depends on TOYBOX_FORK
-  help
-    usage: netcat [-t]
-
-    -t	allocate tty (must come before -l or -L)
 */
 
 #define FOR_netcat
@@ -58,6 +51,7 @@ GLOBALS(
   long quit_delay;       // -q Exit after EOF from stdin after # seconds.
   char *source_address;  // -s Bind to a specific source address.
   long port;             // -p Bind to a specific source port.
+  long idle;             // -W Wait # seconds for more data
   long wait;             // -w Wait # seconds for a connection.
 )
 
@@ -92,36 +86,34 @@ static void lookup_port(char *str, uint16_t *port)
 
 void netcat_main(void)
 {
-  int sockfd=-1, pollcount=2;
-  struct pollfd pollfds[2];
+  struct sockaddr_in *address = (void *)toybuf;
+  int sockfd=-1, in1 = 0, in2 = 0, out1 = 1, out2 = 1;
+  pid_t child;
+
+  // Addjust idle and quit_delay to miliseconds or -1 for no timeout
+  TT.idle = TT.idle ? TT.idle*1000 : -1;
+  TT.quit_delay = TT.quit_delay ? TT.quit_delay*1000 : -1;
 
-  memset(pollfds, 0, 2*sizeof(struct pollfd));
-  pollfds[0].events = pollfds[1].events = POLLIN;
   set_alarm(TT.wait);
 
   // The argument parsing logic can't make "<2" conditional on other
   // arguments like -f and -l, so we do it by hand here.
   if ((toys.optflags&FLAG_f) ? toys.optc :
       (!(toys.optflags&(FLAG_l|FLAG_L)) && toys.optc!=2))
-        help_exit("Argument count wrong");
+        help_exit("bad argument count");
 
-  if (TT.filename) pollfds[0].fd = xopen(TT.filename, O_RDWR);
+  if (TT.filename) in1 = out2 = xopen(TT.filename, O_RDWR);
   else {
-    int temp;
-    struct sockaddr_in address;
-
     // Setup socket
     sockfd = xsocket(AF_INET, SOCK_STREAM, 0);
     fcntl(sockfd, F_SETFD, FD_CLOEXEC);
-    temp = 1;
-    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &temp, sizeof(temp));
-    memset(&address, 0, sizeof(address));
-    address.sin_family = AF_INET;
+    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &out1, sizeof(out1));
+    address->sin_family = AF_INET;
     if (TT.source_address || TT.port) {
-      address.sin_port = SWAP_BE16(TT.port);
+      address->sin_port = SWAP_BE16(TT.port);
       if (TT.source_address)
-        lookup_name(TT.source_address, (uint32_t *)&address.sin_addr);
-      if (bind(sockfd, (struct sockaddr *)&address, sizeof(address)))
+        lookup_name(TT.source_address, (uint32_t *)&(address->sin_addr));
+      if (bind(sockfd, (struct sockaddr *)address, sizeof(*address)))
         perror_exit("bind");
     }
 
@@ -129,102 +121,86 @@ void netcat_main(void)
 
     if (!CFG_NETCAT_LISTEN || !(toys.optflags&(FLAG_L|FLAG_l))) {
       // Figure out where to dial out to.
-      lookup_name(*toys.optargs, (uint32_t *)&address.sin_addr);
-      lookup_port(toys.optargs[1], &address.sin_port);
-      temp = connect(sockfd, (struct sockaddr *)&address, sizeof(address));
-      if (temp<0) perror_exit("connect");
-      pollfds[0].fd = sockfd;
+      lookup_name(*toys.optargs, (uint32_t *)&(address->sin_addr));
+      lookup_port(toys.optargs[1], &(address->sin_port));
+// TODO xconnect
+      if (connect(sockfd, (struct sockaddr *)address, sizeof(*address))<0)
+        perror_exit("connect");
+      in1 = out2 = sockfd;
 
     // Listen for incoming connections
 
     } else {
-      socklen_t len = sizeof(address);
+      socklen_t len = sizeof(*address);
 
       if (listen(sockfd, 5)) error_exit("listen");
       if (!TT.port) {
-        getsockname(sockfd, (struct sockaddr *)&address, &len);
-        printf("%d\n", SWAP_BE16(address.sin_port));
+        getsockname(sockfd, (struct sockaddr *)address, &len);
+        printf("%d\n", SWAP_BE16(address->sin_port));
         fflush(stdout);
-      }
-      // Do we need to return immediately because -l has arguments?
-
-      if ((toys.optflags & FLAG_l) && toys.optc) {
-        if (CFG_TOYBOX_FORK && xfork()) goto cleanup;
-        close(0);
-        close(1);
-        close(2);
+        // Return immediately if no -p and -Ll has arguments, so wrapper
+        // script can use port number.
+        if (CFG_TOYBOX_FORK && toys.optc && xfork()) goto cleanup;
       }
 
       for (;;) {
-        pid_t child = 0;
+        child = 0;
+        len = sizeof(*address); // gcc's insane optimizer can overwrite this
+        in1 = out2 = accept(sockfd, (struct sockaddr *)address, &len);
 
-        // For -l, call accept from the _new_ process.
+        if (in1<0) perror_exit("accept");
 
-        pollfds[0].fd = accept(sockfd, (struct sockaddr *)&address, &len);
-        if (pollfds[0].fd<0) perror_exit("accept");
+        // We can't exit this loop or the optimizer's "liveness analysis"
+        // combines badly with vfork() to corrupt or local variables
+        // (the child's call stack gets trimmed and the next function call
+        // stops the variables the parent tries to re-use next loop)
+        // So there's a bit of redundancy here
 
-        // Do we need a tty?
+        // We have a connection. Disarm timeout.
+        set_alarm(0);
 
-        if (toys.optflags&FLAG_t)
-          child = forkpty(&(pollfds[1].fd), NULL, NULL, NULL);
+        if (toys.optc) {
+          // Do we need a tty?
 
-        // Do we need to fork and/or redirect for exec?
+// TODO nommu, and -t only affects server mode...? Only do -t with optc
+//        if (CFG_TOYBOX_FORK && (toys.optflags&FLAG_t))
+//          child = forkpty(&fdout, NULL, NULL, NULL);
+//        else
+
+          // Do we need to fork and/or redirect for exec?
 
-        else {
           if (toys.optflags&FLAG_L) {
             toys.stacktop = 0;
             child = vfork();
           }
-          if (!child && toys.optc) {
-            int fd = pollfds[0].fd;
-
-            dup2(fd, 0);
-            dup2(fd, 1);
-            if (toys.optflags&FLAG_L) dup2(fd, 2);
-            if (fd>2) close(fd);
+          if (child<0) error_msg("vfork failed\n");
+          else {
+            if (child) {
+              close(in1);
+              continue;
+            }
+            dup2(in1, 0);
+            dup2(in1, 1);
+            if (toys.optflags&FLAG_L) dup2(in1, 2);
+            if (in1>2) close(in1);
+            xexec(toys.optargs);
           }
         }
 
-        if (child<0) error_msg("Fork failed\n");
-        if (child<1) break;
-        close(pollfds[0].fd);
+        pollinate(in1, in2, out1, out2, TT.idle, TT.quit_delay);
+        close(in1);
       }
     }
   }
 
-  // We have a connection.  Disarm timeout.
-  // (Does not play well with -L, but what _should_ that do?)
+  // We have a connection. Disarm timeout.
   set_alarm(0);
 
-  if (CFG_NETCAT_LISTEN && ((toys.optflags&(FLAG_L|FLAG_l)) && toys.optc))
-    xexec(toys.optargs);
-
-  // Poll loop copying stdin->socket and socket->stdout.
-  for (;;) {
-    int i;
-
-    if (0>poll(pollfds, pollcount, -1)) perror_exit("poll");
-
-    for (i=0; i<pollcount; i++) {
-      if (pollfds[i].revents & POLLIN) {
-        int len = read(pollfds[i].fd, toybuf, sizeof(toybuf));
-        if (len<1) goto dohupnow;
-        xwrite(i ? pollfds[0].fd : 1, toybuf, len);
-      } else if (pollfds[i].revents & POLLHUP) {
-dohupnow:
-        // Close half-connection.  This is needed for things like
-        // "echo GET / | netcat landley.net 80"
-        if (i) {
-          shutdown(pollfds[0].fd, SHUT_WR);
-          pollcount--;
-          set_alarm(TT.quit_delay);
-        } else goto cleanup;
-      }
-    }
-  }
+  pollinate(in1, in2, out1, out2, TT.idle, TT.quit_delay);
+
 cleanup:
   if (CFG_TOYBOX_FREE) {
-    close(pollfds[0].fd);
+    close(in1);
     close(sockfd);
   }
 }