Some work I did over the weekend on nbd_client, not sure where I left off...

1 files changed, 30 insertions, 59 deletions

diff --git a/toys/pending/nbd_client.c b/toys/pending/nbd_client.c
index 0465eea1..86972303 100644
--- a/toys/pending/nbd_client.c
+++ b/toys/pending/nbd_client.c
@@ -6,17 +6,23 @@
  *
  * Not in SUSv4.
 
-USE_NBD_CLIENT(NEWTOY(nbd_client, "<3>3", TOYFLAG_USR|TOYFLAG_BIN))
+USE_NBD_CLIENT(NEWTOY(nbd_client, "<3>3ns", TOYFLAG_USR|TOYFLAG_BIN))
 
 config NBD_CLIENT
   bool "nbd-client"
   default n
   help
+    usage: nbd-client [-ns] HOST PORT DEVICE
+
+    -n	Do not fork into background
+    -s	nbd swap support (lock server into memory)
+*/
+
+/*
     Usage: nbd-client [-sSpn] [-b BLKSZ] [-t SECS] [-N name] HOST PORT DEVICE
 
     -b	block size
     -t	timeout in seconds
-    -s	swap
     -S	sdp
     -p	persist
     -n	nofork
@@ -26,22 +32,7 @@ config NBD_CLIENT
 
 #define FOR_nbd_client
 #include "toys.h"
-
-//#include <errno.h>
-//#include <fcntl.h>
-//#include <limits.h>
-#include <netdb.h>
-//#include <stdint.h>
-//#include <stdio.h>
-//#include <stdlib.h>
-//#include <string.h>
-//#include <unistd.h>
-#include <netinet/tcp.h>
-//#include <sys/ioctl.h>
-//#include <sys/mount.h>
-//#include <sys/types.h>
-//#include <sys/socket.h>
-//#include <sys/stat.h>
+#include "toynet.h"
 
 #define NBD_SET_SOCK          _IO(0xab, 0)
 #define NBD_SET_BLKSIZE       _IO(0xab, 1)
@@ -61,19 +52,10 @@ void nbd_client_main(void)
   struct addrinfo *addr, *p;
   char *host=toys.optargs[0], *port=toys.optargs[1], *device=toys.optargs[2];
   uint64_t devsize;
-  char data[124];
-
-  // Parse command line stuff (just a stub now)
-
-  // Make sure the /dev/nbd exists.
-
-  if (0>(nbd = open(device, O_RDWR))) {
-    fprintf(stderr, "Can't open '%s'\n", device);
-    exit(1);
-  }
 
   // Repeat until spanked
 
+  nbd = xopen(device, O_RDWR);
   for (;;) {
     int temp;
     struct addrinfo hints;
@@ -87,30 +69,22 @@ void nbd_client_main(void)
     for (p = addr; p; p = p->ai_next) {
       sock = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
       if (-1 != connect(sock, p->ai_addr, p->ai_addrlen)) break;
+      close(sock);
     }
     freeaddrinfo(addr);
 
-    if (!p) {
-      fprintf(stderr, "Can't connect '%s' port '%s'\n", host, port);
-      exit(1);
-    }
+    if (!p) perror_exit("%s:%s", host, port);
 
     temp = 1;
     setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &temp, sizeof(int));
 
-    // Log on to the server.  (Todo: one big 8+8+8+4+124=152 read)
+    // Read login data
 
-    if (read(sock, data, 8) != 8 || memcmp(data, "NBDMAGIC", 8)
-      || read(sock, &devsize, 8) != 8
-      || SWAP_BE64(devsize) != 0x420281861253LL
-      || read(sock, &devsize, 8) != 8 || read(sock, &flags, 4) != 4
-      || read(sock, data, 124) != 124)
-    {
-      fprintf(stderr, "Login fail\n");
-      exit(1);
-    }
-    devsize = SWAP_BE64(devsize);
-    flags = SWAP_BE32(flags);
+    xreadall(sock, toybuf, 152);
+    if (memcmp(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));
+    flags = SWAP_BE32(*(int *)(toybuf+24));
 
     // Set 4k block size.  Everything uses that these days.
     ioctl(nbd, NBD_SET_BLKSIZE, 4096);
@@ -119,31 +93,30 @@ void nbd_client_main(void)
 
     // If the sucker was exported read only, respect that locally.
     temp = (flags & 2) ? 1 : 0;
-    if (ioctl(nbd, BLKROSET, &temp)<0) {
-      fprintf(stderr, "Login fail\n");
-      exit(1);
-    }
+    xioctl(nbd, BLKROSET, &temp);
 
     if (timeout && ioctl(nbd, NBD_SET_TIMEOUT, timeout)<0) break;
     if (ioctl(nbd, NBD_SET_SOCK, sock) < 0) break;
 
-    // if (swap) mlockall(MCL_CURRENT|MCL_FUTURE);
+    if (toys.optflags & FLAG_s) mlockall(MCL_CURRENT|MCL_FUTURE);
 
     // Open the device to force reread of the partition table.
-    if (!fork()) {
+    if ((toys.optflags & FLAG_n) || !fork()) {
       char *s = strrchr(device, '/');
-      sprintf(data, "/sys/block/%.32s/pid", s ? s+1 : device);
-      // Is it up yet?
-      for (;;) {
-        temp = open(data, O_RDONLY);
-        if (temp == -1) sleep(1);
+      int i;
+
+      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;
         }
       }
       close(open(device, O_RDONLY));
-      exit(0);
+      if (!(toys.optflags & FLAG_n)) exit(0);
     }
 
     // Daemonize here.
@@ -154,13 +127,11 @@ void nbd_client_main(void)
 
     if (ioctl(nbd, NBD_DO_IT)>=0 || errno==EBADR) break;
     close(sock);
-    close(nbd);
   }
+  close(nbd);
 
   // Flush queue and exit.
 
   ioctl(nbd, NBD_CLEAR_QUEUE);
   ioctl(nbd, NBD_CLEAR_SOCK);
-
-  exit(0);
 }