From bcd6d11fbe6f41ad392f0cab322cac2689d03543 Mon Sep 17 00:00:00 2001
From: Rob Landley <rob@landley.net>
Date: Tue, 27 Dec 2022 06:24:39 -0600
Subject: [PATCH] Cleanup pass on nbd-client: try to work with nommu, FLAG()
 macros, cleanup when signaled.

Needs testing: the current nbd-server dropped support for this protocol.
---
 toys/other/nbd_client.c | 83 ++++++++++++++++++++++-------------------
 1 file changed, 45 insertions(+), 38 deletions(-)

diff --git a/toys/other/nbd_client.c b/toys/other/nbd_client.c
index 262bc0ce..79929dd7 100644
--- a/toys/other/nbd_client.c
+++ b/toys/other/nbd_client.c
@@ -12,12 +12,11 @@ USE_NBD_CLIENT(OLDTOY(nbd-client, nbd_client, TOYFLAG_USR|TOYFLAG_BIN))
 
 config NBD_CLIENT
   bool "nbd-client"
-  depends on TOYBOX_FORK
   default y
   help
     usage: nbd-client [-ns] HOST PORT DEVICE
 
-    -n	Do not fork into background
+    -n	Do not daemonize
     -s	nbd swap support (lock server into memory)
 */
 
@@ -37,78 +36,86 @@ config NBD_CLIENT
 #include "toys.h"
 #include <linux/nbd.h>
 
+GLOBALS(
+  int nbd;
+)
+
+static void sig_cleanup(int catch)
+{
+  // Flush on the way out
+  ioctl(TT.nbd, NBD_CLEAR_QUE);
+  ioctl(TT.nbd, NBD_CLEAR_SOCK);
+  _exit(catch ? 128+catch : 0);
+}
+
 void nbd_client_main(void)
 {
-  int sock = -1, nbd, flags;
+  int sock = -1, flags, temp;
   unsigned long timeout = 0;
   char *host=toys.optargs[0], *port=toys.optargs[1], *device=toys.optargs[2];
-  uint64_t devsize;
+  unsigned long long devsize;
 
-  // Repeat until spanked
+  // Daemonize in a nommu-friendly way, but retain stderr
+  if (toys.stacktop && !FLAG(n)) {
+    dup2(2, 222);
+    xvdaemon();
+  }
+  dup2(222, 2);
+  close(222);
 
-  nbd = xopen(device, O_RDWR);
-  for (;;) {
-    int temp;
+  TT.nbd = xopen(device, O_RDWR);
+  xsignal(SIGINT, sig_cleanup);
+  xsignal(SIGTERM, sig_cleanup);
 
+  for (;;) {
     // Find and connect to server
-
     sock = xconnectany(xgetaddrinfo(host, port, AF_UNSPEC, SOCK_STREAM, 0, 0));
     temp = 1;
     setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &temp, sizeof(int));
 
     // Read login data
-
     xreadall(sock, toybuf, 152);
     if (xmemcmp(toybuf, "NBDMAGIC\x00\x00\x42\x02\x81\x86\x12\x53", 16))
       error_exit("bad login %s:%s", host, port);
-    devsize = SWAP_BE64(*(uint64_t *)(toybuf+16));
+    devsize = SWAP_BE64(*(unsigned long long *)(toybuf+16));
     flags = SWAP_BE32(*(int *)(toybuf+24));
 
-    // Set 4k block size.  Everything uses that these days.
-    ioctl(nbd, NBD_SET_BLKSIZE, 4096);
-    ioctl(nbd, NBD_SET_SIZE_BLOCKS, devsize/4096);
-    ioctl(nbd, NBD_CLEAR_SOCK);
+    // Use 4k block size
+    ioctl(TT.nbd, NBD_SET_BLKSIZE, 4096);
+    ioctl(TT.nbd, NBD_SET_SIZE_BLOCKS, devsize/4096);
+    ioctl(TT.nbd, NBD_CLEAR_SOCK);
 
-    // If the sucker was exported read only, respect that locally.
-    temp = (flags & 2) ? 1 : 0;
-    xioctl(nbd, BLKROSET, &temp);
+    // Locally respect read only exports
+    flags = (flags>>1)&1;
+    xioctl(TT.nbd, BLKROSET, &flags);
 
-    if (timeout && ioctl(nbd, NBD_SET_TIMEOUT, timeout)<0) break;
-    if (ioctl(nbd, NBD_SET_SOCK, sock) < 0) break;
+    if (timeout && ioctl(TT.nbd, NBD_SET_TIMEOUT, timeout)<0) break;
+    if (ioctl(TT.nbd, NBD_SET_SOCK, sock) < 0) break;
 
-    if (toys.optflags & FLAG_s) mlockall(MCL_CURRENT|MCL_FUTURE);
+    if (FLAG(s)) mlockall(MCL_CURRENT|MCL_FUTURE);
 
     // Open the device to force reread of the partition table.
-    if ((toys.optflags & FLAG_n) || !xfork()) {
+    if (!CFG_TOYBOX_FORK || !xfork()) {
       char *s = strrchr(device, '/');
       int i;
 
+      // Give device up to 10 seconds to come up
       sprintf(toybuf, "/sys/block/%.32s/pid", s ? s+1 : device);
-      // Is it up yet? (Give it 10 seconds.)
       for (i=0; i<100; i++) {
-        temp = open(toybuf, O_RDONLY);
-        if (temp == -1) msleep(100);
-        else {
-          close(temp);
-          break;
-        }
+        if (access(toybuf, F_OK)) break;
+        msleep(100);
       }
       close(open(device, O_RDONLY));
-      if (!(toys.optflags & FLAG_n)) exit(0);
+      if (CFG_TOYBOX_FORK) _exit(0);
     }
 
-    // Daemonize here.
-
-    if (daemon(0,0)) perror_exit("daemonize");
-
     // Process NBD requests until further notice.
 
-    if (ioctl(nbd, NBD_DO_IT)>=0 || errno==EBADR) break;
+    if (ioctl(TT.nbd, NBD_DO_IT)>=0 || errno==EBADR) break;
     close(sock);
+    ioctl(TT.nbd, NBD_CLEAR_QUE);
   }
 
   // Flush queue and exit.
-  ioctl(nbd, NBD_CLEAR_QUE);
-  ioctl(nbd, NBD_CLEAR_SOCK);
-  if (CFG_TOYBOX_FREE) close(nbd);
+  if (CFG_TOYBOX_FREE) close(TT.nbd);
 }
-- 
2.39.2