nbd-client: support newstyle protocol, -b, -d, -p, -N

Recognize the "newstyle" protocol and switch to it automatically.
Add options for setting blocksize (-b) and for disconnecting
a nbd device (-d).

function                                             old     new   delta
nbdclient_main                                       492    1179    +687
static.long_options                                    -      80     +80
print_direc                                          438     457     +19
brctl_main                                          1169    1183     +14
bb_strtoll                                            84      89      +5
packed_usage                                       33082   33066     -16
------------------------------------------------------------------------------
(add/remove: 1/0 grow/shrink: 4/1 up/down: 805/-16)           Total: 789 bytes

Signed-off-by: Elvira Khabirova <lineprinter@altlinux.org>
Signed-off-by: Denys Vlasenko <vda.linux@googlemail.com>

1 files changed, 186 insertions, 44 deletions

diff --git a/networking/nbd-client.c b/networking/nbd-client.c
index bedb01a..103756b 100644
--- a/networking/nbd-client.c
+++ b/networking/nbd-client.c
@@ -16,6 +16,7 @@
 #include "libbb.h"
 #include <netinet/tcp.h>
 #include <linux/fs.h>
+#include <getopt.h>
 
 #define NBD_SET_SOCK          _IO(0xab, 0)
 #define NBD_SET_BLKSIZE       _IO(0xab, 1)
@@ -27,57 +28,144 @@
 #define NBD_SET_SIZE_BLOCKS   _IO(0xab, 7)
 #define NBD_DISCONNECT        _IO(0xab, 8)
 #define NBD_SET_TIMEOUT       _IO(0xab, 9)
+#define NBD_SET_FLAGS         _IO(0xab, 10)
 
 //usage:#define nbdclient_trivial_usage
-//usage:       "HOST PORT BLOCKDEV"
+//usage:       "{ [-b BLKSIZE] [-N NAME] [-t SEC] [-p] HOST [PORT] | -d } BLOCKDEV"
 //usage:#define nbdclient_full_usage "\n\n"
-//usage:       "Connect to HOST and provide a network block device on BLOCKDEV"
+//usage:       "Connect to HOST and provide network block device on BLOCKDEV"
 
-//TODO: more compat with nbd-client version 2.9.13 -
-//Usage: nbd-client [bs=blocksize] [timeout=sec] host port nbd_device [-swap] [-persist] [-nofork]
-//Or   : nbd-client -d nbd_device
-//Or   : nbd-client -c nbd_device
-//Default value for blocksize is 1024 (recommended for ethernet)
+//TODO: more compat with nbd-client version 3.17 -
+//nbd-client host [ port ] nbd-device [ -connections num ] [ -sdp ] [ -swap ]
+//	[ -persist ] [ -nofork ] [ -nonetlink ] [ -systemd-mark ]
+//	[ -block-size block size ] [ -timeout seconds ] [ -name name ]
+//	[ -certfile certfile ] [ -keyfile keyfile ] [ -cacertfile cacertfile ]
+//	[ -tlshostname hostname ]
+//nbd-client -unix path nbd-device [ -connections num ] [ -sdp ] [ -swap ]
+//	[ -persist ] [ -nofork ] [ -nonetlink ] [ -systemd-mark ]
+//	[ -block-size block size ] [ -timeout seconds ] [ -name name ]
+//nbd-client nbd-device
+//nbd-client -d nbd-device
+//nbd-client -c nbd-device
+//nbd-client -l host [ port ]
+//nbd-client [ -netlink ] -l host
+//
+//Default value for blocksize is 4096
 //Allowed values for blocksize are 512,1024,2048,4096
-//Note, that kernel 2.4.2 and older ones do not work correctly with
-//blocksizes other than 1024 without patches
 
 int nbdclient_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
-int nbdclient_main(int argc UNUSED_PARAM, char **argv)
+int nbdclient_main(int argc, char **argv)
 {
-	unsigned long timeout = 0;
 #if BB_MMU
-	int nofork = 0;
+	bool nofork;
 #endif
-	char *host, *port, *device;
+	bool opt_d;
+	bool opt_p;
+	const char *host, *port, *device;
+	const char *name;
+	unsigned blksize, size_blocks;
+	unsigned timeout;
+	int ch;
 	struct nbd_header_t {
 		uint64_t magic1; // "NBDMAGIC"
-		uint64_t magic2; // 0x420281861253 big endian
+		uint64_t magic2; // old style: 0x420281861253 big endian
+		//               // new style: 0x49484156454F5054 (IHAVEOPT)
+	} nbd_header;
+	struct old_nbd_header_t {
 		uint64_t devsize;
 		uint32_t flags;
 		char data[124];
-	} nbd_header;
+	} old_nbd_header;
+	struct new_nbd_header_t {
+		uint64_t devsize;
+		uint16_t transmission_flags;
+		char data[124];
+	} new_nbd_header;
+	struct nbd_opt_t {
+		uint64_t magic;
+		uint32_t opt;
+		uint32_t len;
+	} nbd_opts;
 
-	BUILD_BUG_ON(offsetof(struct nbd_header_t, data) != 8+8+8+4);
+	static const struct option long_options[] = {
+		{ "block-size", required_argument, NULL, 'b' },
+		{ "timeout"   , required_argument, NULL, 't' },
+		{ "name"      , required_argument, NULL, 'n' },
+		{ "persist"   , no_argument      , NULL, 'p' },
+		{ NULL }
+	};
 
-	// Parse command line stuff (just a stub now)
-	if (!argv[1] || !argv[2] || !argv[3] || argv[4])
-		bb_show_usage();
+	BUILD_BUG_ON(offsetof(struct old_nbd_header_t, data) != 8 + 4);
+	BUILD_BUG_ON(offsetof(struct new_nbd_header_t, data) != 8 + 2);
 
 #if !BB_MMU
 	bb_daemonize_or_rexec(DAEMON_CLOSE_EXTRA_FDS, argv);
 #endif
 
-	host = argv[1];
-	port = argv[2];
-	device = argv[3];
+	// Parse args. nbd-client uses stupid "one-dash long options" style :(
+	// Even though short forms (-b,-t,-N,-p) exist for all long opts,
+	// older manpages only contained long forms, which probably resulted
+	// in many scripts using them.
+	blksize = 4096;
+	timeout = 0;
+	name = ""; // use of "" instead of NULL simplifies strlen() later
+	opt_d = opt_p = 0;
+	while ((ch = getopt_long_only(argc, argv, "dN:", long_options, NULL)) != -1) {
+		switch (ch) {
+		case 'p':	// -persist
+			opt_p = 1;
+			break;
+		case 'd':	// -d
+			opt_d = 1;
+			break;
+		case 'b':	// -block-size
+			blksize = xatou(optarg);
+			break;
+		case 't':	// -timeout
+			timeout = xatou(optarg);
+			break;
+		case 'N':	// -N
+		case 'n':	// -name
+			name = optarg;
+			break;
+		default:
+			bb_show_usage();
+		}
+	}
+	argv += optind;
+
+	if (opt_d) { // -d
+		if (argv[0] && !argv[1]) {
+			int nbd = xopen(argv[0], O_RDWR);
+			ioctl(nbd, NBD_DISCONNECT);
+			ioctl(nbd, NBD_CLEAR_SOCK);
+			if (ENABLE_FEATURE_CLEAN_UP)
+				close(nbd);
+			return 0;
+		}
+		bb_show_usage();
+	}
+
+	// Allow only argv[] of: HOST [PORT] BLOCKDEV
+	if (!argv[0] || !argv[1] || (argv[2] && argv[3])) {
+		bb_show_usage();
+	}
 
-	// Repeat until spanked (-persist behavior)
-	for (;;) {
+	host = argv[0];
+	port = argv[2] ? argv[1] : "10809";
+	device = argv[2] ? argv[2] : argv[1];
+
+	// Repeat until spanked if -persist
+#if BB_MMU
+	nofork = 0;
+#endif
+	do {
 		int sock, nbd;
 		int ro;
+		int proto_new; // 0 for old, 1 for new
+		char *data;
 
-		// Make sure the /dev/nbd exists
+		// Make sure BLOCKDEV exists
 		nbd = xopen(device, O_RDWR);
 
 		// Find and connect to server
@@ -85,40 +173,95 @@ int nbdclient_main(int argc UNUSED_PARAM, char **argv)
 		setsockopt_1(sock, IPPROTO_TCP, TCP_NODELAY);
 
 		// Log on to the server
-		xread(sock, &nbd_header, 8+8+8+4 + 124);
-		if (memcmp(&nbd_header.magic1, "NBDMAGIC""\x00\x00\x42\x02\x81\x86\x12\x53", 16) != 0)
+		xread(sock, &nbd_header, 8 + 8);
+		if (memcmp(&nbd_header.magic1, "NBDMAGIC",
+				sizeof(nbd_header.magic1)) != 0
+		) {
+			bb_error_msg_and_die("login failed");
+		}
+		if (memcmp(&nbd_header.magic2,
+				"\x00\x00\x42\x02\x81\x86\x12\x53",
+				sizeof(nbd_header.magic2)) == 0
+		) {
+			proto_new = 0;
+		} else if (memcmp(&nbd_header.magic2, "IHAVEOPT", 8) == 0) {
+			proto_new = 1;
+		} else {
 			bb_error_msg_and_die("login failed");
+		}
+
+		if (!proto_new) {
+			xread(sock, &old_nbd_header,
+					sizeof(old_nbd_header.devsize) +
+					sizeof(old_nbd_header.flags) +
+					sizeof(old_nbd_header.data));
+			size_blocks = SWAP_BE64(old_nbd_header.devsize) / blksize;
+			ioctl(nbd, NBD_SET_BLKSIZE, (unsigned long) blksize);
+			ioctl(nbd, NBD_SET_SIZE_BLOCKS, size_blocks);
+			ioctl(nbd, NBD_CLEAR_SOCK);
+			ro = !!(old_nbd_header.flags & htons(2));
+			data = old_nbd_header.data;
+		} else {
+			unsigned namelen;
+			uint16_t handshake_flags;
+
+			xread(sock, &handshake_flags, sizeof(handshake_flags));
+			xwrite(sock, &const_int_0, sizeof(const_int_0)); // client_flags
+
+			memcpy(&nbd_opts.magic, "IHAVEOPT",
+					sizeof(nbd_opts.magic));
+			nbd_opts.opt = htonl(1); // NBD_OPT_EXPORT_NAME
+			namelen = strlen(name);
+			nbd_opts.len = htonl(namelen);
+			xwrite(sock, &nbd_opts,
+					sizeof(nbd_opts.magic) +
+					sizeof(nbd_opts.opt) +
+					sizeof(nbd_opts.len));
+			xwrite(sock, name, namelen);
 
-		// Set 4k block size.  Everything uses that these days
-		ioctl(nbd, NBD_SET_BLKSIZE, 4096);
-		ioctl(nbd, NBD_SET_SIZE_BLOCKS, SWAP_BE64(nbd_header.devsize) / 4096);
-		ioctl(nbd, NBD_CLEAR_SOCK);
+			xread(sock, &new_nbd_header,
+					sizeof(new_nbd_header.devsize) +
+					sizeof(new_nbd_header.transmission_flags) +
+					sizeof(new_nbd_header.data));
+			size_blocks = SWAP_BE64(new_nbd_header.devsize) / blksize;
+			ioctl(nbd, NBD_SET_BLKSIZE, (unsigned long) blksize);
+			ioctl(nbd, NBD_SET_SIZE_BLOCKS, size_blocks);
+			ioctl(nbd, NBD_CLEAR_SOCK);
+			ioctl(nbd, NBD_SET_FLAGS,
+					ntohs(new_nbd_header.transmission_flags));
+			ro = !!(new_nbd_header.transmission_flags & htons(2));
+			data = new_nbd_header.data;
+		}
 
-		// If the sucker was exported read only, respect that locally
-		ro = (nbd_header.flags & SWAP_BE32(2)) / SWAP_BE32(2);
-		if (ioctl(nbd, BLKROSET, &ro) < 0)
+		if (ioctl(nbd, BLKROSET, &ro) < 0) {
 			bb_perror_msg_and_die("BLKROSET");
+		}
 
-		if (timeout)
-			if (ioctl(nbd, NBD_SET_TIMEOUT, timeout))
+		if (timeout) {
+			if (ioctl(nbd, NBD_SET_TIMEOUT, (unsigned long) timeout)) {
 				bb_perror_msg_and_die("NBD_SET_TIMEOUT");
-		if (ioctl(nbd, NBD_SET_SOCK, sock))
-			bb_perror_msg_and_die("NBD_SET_SOCK");
+			}
+		}
 
-		// if (swap) mlockall(MCL_CURRENT|MCL_FUTURE);
+		if (ioctl(nbd, NBD_SET_SOCK, sock)) {
+			bb_perror_msg_and_die("NBD_SET_SOCK");
+		}
 
+		//if (swap) mlockall(MCL_CURRENT|MCL_FUTURE);
 #if BB_MMU
 		// Open the device to force reread of the partition table.
 		// Need to do it in a separate process, since open(device)
 		// needs some other process to sit in ioctl(nbd, NBD_DO_IT).
 		if (fork() == 0) {
+			/* child */
 			char *s = strrchr(device, '/');
-			sprintf(nbd_header.data, "/sys/block/%.32s/pid", s ? s + 1 : device);
+			sprintf(data, "/sys/block/%.32s/pid", s ? s + 1 : device);
 			// Is it up yet?
 			for (;;) {
-				int fd = open(nbd_header.data, O_RDONLY);
+				int fd = open(data, O_RDONLY);
 				if (fd >= 0) {
-					//close(fd);
+					if (ENABLE_FEATURE_CLEAN_UP)
+						close(fd);
 					break;
 				}
 				sleep(1);
@@ -133,7 +276,6 @@ int nbdclient_main(int argc UNUSED_PARAM, char **argv)
 			nofork = 1;
 		}
 #endif
-
 		// This turns us (the process that calls this ioctl)
 		// into a dedicated NBD request handler.
 		// We block here for a long time.
@@ -148,7 +290,7 @@ int nbdclient_main(int argc UNUSED_PARAM, char **argv)
 
 		close(sock);
 		close(nbd);
-	}
+	} while (opt_p);
 
 	return 0;
 }