[PATCH 2/2] Support preinitialized NBD connections

[Josh Triplett <josh@joshtriplett.org>]

Add an option to skip negotiation. Add an option to specify the device
size, required if using preinit.
---
 man/nbd-client.8.in.sgml | 26 +++++++++++++++++++++++
 nbd-client.c             | 45 ++++++++++++++++++++++++++++++++++------
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/man/nbd-client.8.in.sgml b/man/nbd-client.8.in.sgml
index cd4dacd..7138809 100644
--- a/man/nbd-client.8.in.sgml
+++ b/man/nbd-client.8.in.sgml
@@ -66,7 +66,9 @@ manpage.1: manpage.sgml
       <arg>-nonetlink</arg>
       <arg>-systemd-mark</arg>
       <arg>-readonly</arg>
+      <arg>-preinit</arg>
       <arg>-block-size <replaceable>block size</replaceable></arg>
+      <arg>-size <replaceable>bytes</replaceable></arg>
       <arg>-timeout <replaceable>seconds</replaceable></arg>
       <arg>-name <replaceable>name</replaceable></arg>
       <arg>-certfile <replaceable>certfile</replaceable></arg>
@@ -86,7 +88,9 @@ manpage.1: manpage.sgml
       <arg>-nonetlink</arg>
       <arg>-systemd-mark</arg>
       <arg>-readonly</arg>
+      <arg>-preinit</arg>
       <arg>-block-size <replaceable>block size</replaceable></arg>
+      <arg>-size <replaceable>bytes</replaceable></arg>
       <arg>-timeout <replaceable>seconds</replaceable></arg>
       <arg>-name <replaceable>name</replaceable></arg>
     </cmdsynopsis>
@@ -266,6 +270,18 @@ manpage.1: manpage.sgml
 	    server or something similar.</para>
 	</listitem>
       </varlistentry>
+      <varlistentry>
+        <term><option>-preinit</option></term>
+        <term><option>-P</option></term>
+        <listitem>
+          <para>When this option is specified, &dhpackage; will skip
+            the usual negotiation with the server, and hand the socket
+            to the kernel immediately after connecting. Only use this
+            when connecting to specialized NBD servers specifically
+            designed for it. This requires specifying the size of the
+            device via the -B option, and does not support TLS.</para>
+        </listitem>
+      </varlistentry>
       <varlistentry>
         <term><option>-readonly</option></term>
         <term><option>-R</option></term>
@@ -275,6 +291,16 @@ manpage.1: manpage.sgml
             would allow writes.</para>
         </listitem>
       </varlistentry>
+      <varlistentry>
+	<term><option>-size <replaceable>bytes</replaceable></option></term>
+	<term><option>-B <replaceable>bytes</replaceable></option></term>
+	<listitem>
+          <para>Force the device size to the specified number of bytes,
+            rather than using the value from server negotiation. Must
+            be a multiple of the block size. If using preinit (-P) to
+            skip negotiation, this option is required.</para>
+	</listitem>
+      </varlistentry>
       <varlistentry>
         <term><option>-sdp</option></term>
 	<term><option>-S</option></term>
diff --git a/nbd-client.c b/nbd-client.c
index 2023962..aa314df 100644
--- a/nbd-client.c
+++ b/nbd-client.c
@@ -896,9 +896,9 @@ void usage(char* errmsg, ...) {
 		fprintf(stderr, "%s version %s\n", PROG_NAME, PACKAGE_VERSION);
 	}
 #if HAVE_NETLINK
-	fprintf(stderr, "Usage: nbd-client -name|-N name host [port] nbd_device\n\t[-block-size|-b block size] [-timeout|-t timeout] [-swap|-s] [-sdp|-S]\n\t[-persist|-p] [-nofork|-n] [-systemd-mark|-m] [-nonetlink|-L]\n\t[-readonly|-R]\n");
+	fprintf(stderr, "Usage: nbd-client -name|-N name host [port] nbd_device\n\t[-block-size|-b block size] [-timeout|-t timeout] [-swap|-s] [-sdp|-S]\n\t[-persist|-p] [-nofork|-n] [-systemd-mark|-m] [-nonetlink|-L]\n\t[-readonly|-R] [-size|-B bytes] [-preinit|-P]\n");
 #else
-	fprintf(stderr, "Usage: nbd-client -name|-N name host [port] nbd_device\n\t[-block-size|-b block size] [-timeout|-t timeout] [-swap|-s] [-sdp|-S]\n\t[-persist|-p] [-nofork|-n] [-systemd-mark|-m]\n\t[-readonly|-R]\n");
+	fprintf(stderr, "Usage: nbd-client -name|-N name host [port] nbd_device\n\t[-block-size|-b block size] [-timeout|-t timeout] [-swap|-s] [-sdp|-S]\n\t[-persist|-p] [-nofork|-n] [-systemd-mark|-m]\n\t[-readonly|-R] [-size|-B bytes] [-preinit|-P]\n");
 #endif
 	fprintf(stderr, "Or   : nbd-client -u (with same arguments as above)\n");
 	fprintf(stderr, "Or   : nbd-client nbdX\n");
@@ -934,9 +934,9 @@ void disconnect(char* device) {
 }
 
 #if HAVE_NETLINK
-static const char *short_opts = "-A:b:c:C:d:gH:hK:LlnN:pRSst:uVx";
+static const char *short_opts = "-A:B:b:c:C:d:gH:hK:LlnN:PpRSst:uVx";
 #else
-static const char *short_opts = "-A:b:c:C:d:gH:hK:lnN:pRSst:uVx";
+static const char *short_opts = "-A:B:b:c:C:d:gH:hK:lnN:PpRSst:uVx";
 #endif
 
 int main(int argc, char *argv[]) {
@@ -951,9 +951,11 @@ int main(int argc, char *argv[]) {
 	int sdp=0;
 	int G_GNUC_UNUSED nofork=0; // if -dNOFORK
 	pid_t main_pid;
-	u64 size64;
+	u64 size64 = 0;
+	u64 force_size64 = 0;
 	uint16_t flags = 0;
 	bool force_read_only = false;
+	bool preinit = false;
 	int c;
 	int nonspecial=0;
 	int b_unix=0;
@@ -975,6 +977,7 @@ int main(int argc, char *argv[]) {
 	struct option long_options[] = {
 		{ "cacertfile", required_argument, NULL, 'A' },
 		{ "block-size", required_argument, NULL, 'b' },
+		{ "size", required_argument, NULL, 'B' },
 		{ "check", required_argument, NULL, 'c' },
 		{ "connections", required_argument, NULL, 'C'},
 		{ "disconnect", required_argument, NULL, 'd' },
@@ -991,6 +994,7 @@ int main(int argc, char *argv[]) {
 		{ "nofork", no_argument, NULL, 'n' },
 		{ "name", required_argument, NULL, 'N' },
 		{ "persist", no_argument, NULL, 'p' },
+		{ "preinit", no_argument, NULL, 'P' },
 		{ "readonly", no_argument, NULL, 'R' },
 		{ "sdp", no_argument, NULL, 'S' },
 		{ "swap", no_argument, NULL, 's' },
@@ -1060,6 +1064,13 @@ int main(int argc, char *argv[]) {
 				exit(EXIT_FAILURE);
 			}
 			break;
+		case 'B':
+			force_size64=(u64)strtoull(optarg, NULL, 0);
+			if(force_size64 == 0) {
+				fprintf(stderr, "E: Invalid size\n");
+				exit(EXIT_FAILURE);
+			}
+			break;
 		case 'c':
 			return check_conn(optarg, 1);
 		case 'C':
@@ -1097,6 +1108,9 @@ int main(int argc, char *argv[]) {
 		case 'p':
 			cont=1;
 			break;
+		case 'P':
+			preinit = true;
+			break;
 		case 'R':
 			force_read_only = true;
 			break;
@@ -1181,12 +1195,28 @@ int main(int argc, char *argv[]) {
 		tls = true;
 	}
 
+	if (preinit) {
+		if (tls) {
+			fprintf(stderr, "E: preinit connection cannot be used with TLS\n");
+			exit(EXIT_FAILURE);
+		}
+		if (!force_size64) {
+			fprintf(stderr, "E: preinit connection requires specifying size\n");
+			exit(EXIT_FAILURE);
+		}
+	}
+
         if (!tlshostname && hostname)
                 tlshostname = strdup(hostname);
 
 	if (netlink)
 		nofork = 1;
 
+	if((force_size64 % blocksize) != 0) {
+		fprintf(stderr, "E: size (" PRIu64 " bytes) is not a multiple of blocksize (%d)!\n", force_size64, blocksize);
+		exit(EXIT_FAILURE);
+	}
+
 	if(strlen(name)==0 && !(opts & NBDC_DO_LIST)) {
 		printf("Warning: the oldstyle protocol is no longer supported.\nThis method now uses the newstyle protocol with a default export\n");
 	}
@@ -1211,9 +1241,12 @@ int main(int argc, char *argv[]) {
 		if (sock < 0)
 			exit(EXIT_FAILURE);
 
-		negotiate(&sock, &size64, &flags, name, needed_flags, cflags, opts, certfile, keyfile, cacertfile, tlshostname, tls, can_opt_go);
+		if (!preinit)
+			negotiate(&sock, &size64, &flags, name, needed_flags, cflags, opts, certfile, keyfile, cacertfile, tlshostname, tls, can_opt_go);
 		if (force_read_only)
 			flags |= NBD_FLAG_READ_ONLY;
+		if (force_size64)
+			size64 = force_size64;
 		if (netlink) {
 			sockfds[i] = sock;
 			continue;
-- 
2.27.0