[Libguestfs] [nbdkit PATCH 3/4] api: Add nbdkit_disconnect(int)

Eric Blake
Wed, 26 Oct 2022 22:18:01 -0000

We have already exposed nbdkit_shutdown() as a way for a plugin to
request that nbdkit itself shuts down when convenient (taking away the
current and all other client connections).  But we lacked a way for a
plugin (or filter) to drop the current client, while still leaving
nbdkit up for other clients.  This has utility for testing client
reconnect logic, as well as for emulating things like qemu's
propensity to hard-disconnect on NBD_CMD_WRITE with more than 32M of
data (even though nbdkit would normally accept up to 64M of data).
Add a new API nbdkit_disconnect(int force) for that purpose, utilizing
the logic already built in to connections.c.  If force is true, we
close the socket immediately (the client won't get responses to the
current or any other parallel commands); if force is false, we mark
the socket as not accepting further commands (current commands can
still give their response to the client, but new requests from the
current client are ignored with ESHUTDOWN until the client gives
NBD_CMD_DISC).

This patch provides the C interface (using int, to allow for future
extensions beyond our current two values), and straightforward
bindings to languages that already have nbdkit_shutdown bound (here,
we use bool where it exists, since we don't have to worry as much
about potential future API breaks if we need a third state in the
future).  Upcoming patches will then work on extending the sh/eval
plugin to take advantage of the API, as well as enhancing the
blocksize-policy filter to more closely emulate qemu.

The logic for connection status changes slightly - there is now a
fourth state tracking when we want a soft shutdown, but are still
waiting for the client to send NBD_CMD_DISC - existing commands still
get to send their reply unmolested, but future commands get an
immediate ESHUTDOWN error.  Rather than re-checking the state between
dropping the read lock and grabbing the write lock, we now go by the
state present after recv() first gives us a result at the start of a
transaction.

The testsuite covers both hard and soft disconnects, and in both
plaintext and tls.
---
 docs/nbdkit-plugin.pod         |  17 ++++-
 include/nbdkit-common.h        |   3 +-
 tests/Makefile.am              |  28 ++++++++
 server/internal.h              |   1 +
 server/connections.c           |   3 +-
 server/nbdkit.syms             |   3 +-
 server/protocol.c              |   4 +-
 server/public.c                |  14 +++-
 server/test-public.c           |  10 ++-
 plugins/ocaml/NBDKit.mli       |   5 +-
 plugins/ocaml/NBDKit.ml        |   1 +
 plugins/ocaml/bindings.c       |  12 +++-
 plugins/python/modfunctions.c  |  14 ++++
 plugins/rust/src/lib.rs        |   6 ++
 tests/test-disconnect-tls.sh   | 126 +++++++++++++++++++++++++++++++++
 tests/test-disconnect.sh       | 100 ++++++++++++++++++++++++++
 tests/test-disconnect-plugin.c |  95 +++++++++++++++++++++++++
 17 files changed, 431 insertions(+), 11 deletions(-)
 create mode 100755 tests/test-disconnect-tls.sh
 create mode 100755 tests/test-disconnect.sh
 create mode 100644 tests/test-disconnect-plugin.c

diff --git a/docs/nbdkit-plugin.pod b/docs/nbdkit-plugin.pod
index 6e74ccad..d338cde8 100644
--- a/docs/nbdkit-plugin.pod
+++ b/docs/nbdkit-plugin.pod
@@ -1339,7 +1339,8 @@ Plugins and filters can call L<exit(3)> in the
configuration phase

 Once nbdkit has started serving connections, plugins and filters
 should not call L<exit(3)>.  However they may instruct nbdkit to shut
-down by calling C<nbdkit_shutdown>:
+down by calling C<nbdkit_shutdown>, or to disconnect a single client
+by calling C<nbdkit_disconnect>:

  void nbdkit_shutdown (void);

@@ -1349,6 +1350,20 @@ the plugin and all filters are unloaded cleanly which may
take some
 time.  Further callbacks from nbdkit into the plugin or filter may
 occur after you have called this.

+ void nbdkit_disconnect (int force);
+
+This function requests that the current connection be disconnected
+(I<note> that it does not affect other connections, and the client may
+try to reconnect).  It is only useful from connected callbacks (that
+is, after C<.open> and before C<.close>).  If C<force> is
true, nbdkit
+will disconnect the client immediately, and the client will not
+receive any response to the current command or any other commands in
+flight in parallel threads; otherwise, nbdkit will not accept any new
+commands from the client (failing all commands other than
+C<NBD_CMD_DISC> with C<ESHUTDOWN>), but will allow existing
commands
+to complete gracefully.  Either way, the next callback for the current
+connection should be C<.close>.
+
 =head1 PARSING COMMAND LINE PARAMETERS

 =head2 Parsing numbers
diff --git a/include/nbdkit-common.h b/include/nbdkit-common.h
index b0dcf715..dfdbab68 100644
--- a/include/nbdkit-common.h
+++ b/include/nbdkit-common.h
@@ -1,5 +1,5 @@
 /* nbdkit
- * Copyright (C) 2013-2020 Red Hat Inc.
+ * Copyright (C) 2013-2022 Red Hat Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -137,6 +137,7 @@ NBDKIT_EXTERN_DECL (int64_t, nbdkit_peer_pid, (void));
 NBDKIT_EXTERN_DECL (int64_t, nbdkit_peer_uid, (void));
 NBDKIT_EXTERN_DECL (int64_t, nbdkit_peer_gid, (void));
 NBDKIT_EXTERN_DECL (void, nbdkit_shutdown, (void));
+NBDKIT_EXTERN_DECL (void, nbdkit_disconnect, (int force));

 NBDKIT_EXTERN_DECL (const char *, nbdkit_strdup_intern,
                     (const char *str));
diff --git a/tests/Makefile.am b/tests/Makefile.am
index 530b22bd..e951381d 100644
--- a/tests/Makefile.am
+++ b/tests/Makefile.am
@@ -253,6 +253,8 @@ TESTS += \
 	test-long-name.sh \
 	test-flush.sh \
 	test-swap.sh \
+	test-disconnect.sh \
+	test-disconnect-tls.sh \
 	test-shutdown.sh \
 	test-nbdkit-backend-debug.sh \
 	test-read-password.sh \
@@ -293,6 +295,8 @@ EXTRA_DIST += \
 	test-read-password.sh \
 	test-read-password-interactive.sh \
 	test-read-password-plugin.c \
+	test-disconnect.sh \
+	test-disconnect-tls.sh \
 	test-shutdown.sh \
 	test-single-from-file.sh \
 	test-single-sh.sh \
@@ -376,6 +380,30 @@ test_flush_plugin_la_LDFLAGS = \
 	$(NULL)
 test_flush_plugin_la_LIBADD = $(IMPORT_LIBRARY_ON_WINDOWS)

+# check_LTLIBRARIES won't build a shared library (see automake manual).
+# So we have to do this and add a dependency.
+noinst_LTLIBRARIES += \
+	test-disconnect-plugin.la \
+	$(NULL)
+test-disconnect.sh: test-disconnect-plugin.la
+test-disconnect-tls.sh: test-disconnect-plugin.la keys.psk
+
+test_disconnect_plugin_la_SOURCES = \
+	test-disconnect-plugin.c \
+	$(top_srcdir)/include/nbdkit-plugin.h \
+	$(NULL)
+test_disconnect_plugin_la_CPPFLAGS = \
+	-I$(top_srcdir)/include \
+	-I$(top_builddir)/include \
+	$(NULL)
+test_disconnect_plugin_la_CFLAGS = $(WARNINGS_CFLAGS)
+# For use of the -rpath option, see:
+# https://lists.gnu.org/archive/html/libtool/2007-07/msg00067.html
+test_disconnect_plugin_la_LDFLAGS = \
+	-module -avoid-version -shared $(NO_UNDEFINED_ON_WINDOWS) -rpath /nowhere \
+	$(NULL)
+test_disconnect_plugin_la_LIBADD = $(IMPORT_LIBRARY_ON_WINDOWS)
+
 # check_LTLIBRARIES won't build a shared library (see automake manual).
 # So we have to do this and add a dependency.
 noinst_LTLIBRARIES += \
diff --git a/server/internal.h b/server/internal.h
index 69b4302c..229d707a 100644
--- a/server/internal.h
+++ b/server/internal.h
@@ -235,6 +235,7 @@ struct context {
 typedef enum {
   STATUS_DEAD,         /* Connection is closed */
   STATUS_CLIENT_DONE,  /* Client has sent NBD_CMD_DISC */
+  STATUS_SHUTDOWN,     /* Server wants soft shutdown */
   STATUS_ACTIVE,       /* Client can make requests */
 } conn_status;

diff --git a/server/connections.c b/server/connections.c
index 1b6183df..4d776f2a 100644
--- a/server/connections.c
+++ b/server/connections.c
@@ -90,7 +90,8 @@ connection_set_status (conn_status value)
       pthread_mutex_lock (&conn->status_lock))
     abort ();
   if (value < conn->status) {
-    if (conn->nworkers && conn->status > STATUS_CLIENT_DONE) {
+    if (conn->nworkers && conn->status > STATUS_CLIENT_DONE
&&
+        value <= STATUS_CLIENT_DONE) {
       char c = 0;

       assert (conn->status_pipe[1] >= 0);
diff --git a/server/nbdkit.syms b/server/nbdkit.syms
index 0e897680..45cf3b45 100644
--- a/server/nbdkit.syms
+++ b/server/nbdkit.syms
@@ -1,5 +1,5 @@
 # nbdkit
-# Copyright (C) 2018-2021 Red Hat Inc.
+# Copyright (C) 2018-2022 Red Hat Inc.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are
@@ -44,6 +44,7 @@
     nbdkit_context_get_backend;
     nbdkit_context_set_next;
     nbdkit_debug;
+    nbdkit_disconnect;
     nbdkit_error;
     nbdkit_export_name;
     nbdkit_exports_count;
diff --git a/server/protocol.c b/server/protocol.c
index d1e01502..cc1e4ed8 100644
--- a/server/protocol.c
+++ b/server/protocol.c
@@ -631,10 +631,10 @@ protocol_recv_request_send_reply (void)
   /* Read the request packet. */
   {
     ACQUIRE_LOCK_FOR_CURRENT_SCOPE (&conn->read_lock);
+    r = conn->recv (&request, sizeof request);
     cs = connection_get_status ();
     if (cs <= STATUS_CLIENT_DONE)
       return;
-    r = conn->recv (&request, sizeof request);
     if (r == -1) {
       nbdkit_error ("read request: %m");
       connection_set_status (STATUS_DEAD);
@@ -718,7 +718,7 @@ protocol_recv_request_send_reply (void)
   }

   /* Perform the request.  Only this part happens inside the request lock. */
-  if (quit || connection_get_status () == STATUS_CLIENT_DONE) {
+  if (quit || cs < STATUS_ACTIVE) {
     error = ESHUTDOWN;
   }
   else {
diff --git a/server/public.c b/server/public.c
index 6a9840bb..c2f67451 100644
--- a/server/public.c
+++ b/server/public.c
@@ -728,7 +728,7 @@ nbdkit_nanosleep (unsigned sec, unsigned nsec)
   bool has_quit = quit;
   assert (has_quit ||
           (conn && conn->nworkers > 0 &&
-           connection_get_status () < STATUS_ACTIVE) ||
+           connection_get_status () < STATUS_SHUTDOWN) ||
           (conn && (fds[2].revents & (POLLRDHUP | POLLHUP | POLLERR
|
                                       POLLNVAL))));
   if (has_quit)
@@ -1097,3 +1097,15 @@ nbdkit_printf_intern (const char *fmt, ...)
   va_end (ap);
   return ret;
 }
+
+NBDKIT_DLL_PUBLIC void
+nbdkit_disconnect (int force)
+{
+  struct connection *conn = threadlocal_get_conn ();
+
+  if (!conn) {
+    debug ("no connection in this thread, ignoring disconnect
request");
+    return;
+  }
+  connection_set_status (force ? STATUS_DEAD : STATUS_SHUTDOWN);
+}
diff --git a/server/test-public.c b/server/test-public.c
index 1d83354f..4e4d8a2e 100644
--- a/server/test-public.c
+++ b/server/test-public.c
@@ -63,6 +63,8 @@ nbdkit_debug (const char *fs, ...)

 bool listen_stdin;
 bool configured;
+bool verbose;
+int tls;

 volatile int quit;
 #ifndef WIN32
@@ -89,14 +91,18 @@ connection_get_status (void)
   abort ();
 }

+void
+connection_set_status (conn_status v)
+{
+  abort ();
+}
+
 const char *
 backend_default_export (struct backend *b, int readonly)
 {
   abort ();
 }

-int tls;
-
 /* Unit tests. */

 static bool
diff --git a/plugins/ocaml/NBDKit.mli b/plugins/ocaml/NBDKit.mli
index cc389ca0..cdfacf69 100644
--- a/plugins/ocaml/NBDKit.mli
+++ b/plugins/ocaml/NBDKit.mli
@@ -1,6 +1,6 @@
 (* hey emacs, this is OCaml code: -*- tuareg -*- *)
 (* nbdkit OCaml interface
- * Copyright (C) 2014-2020 Red Hat Inc.
+ * Copyright (C) 2014-2022 Red Hat Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -166,6 +166,9 @@ val export_name : unit -> string
 (** Binding for [nbdkit_shutdown].  Requests the server shut down. *)
 val shutdown : unit -> unit

+(** Binding for [nbdkit_disconnect].  Requests disconnecting current client. *)
+val disconnect : bool -> unit
+
 (** Print a debug message when nbdkit is in verbose mode. *)
 val debug : ('a, unit, string, unit) format4 -> 'a

diff --git a/plugins/ocaml/NBDKit.ml b/plugins/ocaml/NBDKit.ml
index 3ad72acb..c94f4d57 100644
--- a/plugins/ocaml/NBDKit.ml
+++ b/plugins/ocaml/NBDKit.ml
@@ -166,6 +166,7 @@ external realpath : string -> string =
"ocaml_nbdkit_realpath"
 external nanosleep : int -> int -> unit =
"ocaml_nbdkit_nanosleep"
 external export_name : unit -> string = "ocaml_nbdkit_export_name"
 external shutdown : unit -> unit = "ocaml_nbdkit_shutdown"
[@@noalloc]
+external disconnect : bool -> unit = "ocaml_nbdkit_disconnect"
[@@noalloc]
 external _debug : string -> unit = "ocaml_nbdkit_debug"
[@@noalloc]
 let debug fs = ksprintf _debug fs
 external version : unit -> string = "ocaml_nbdkit_version"
diff --git a/plugins/ocaml/bindings.c b/plugins/ocaml/bindings.c
index ba95fb4a..c6c152b9 100644
--- a/plugins/ocaml/bindings.c
+++ b/plugins/ocaml/bindings.c
@@ -1,5 +1,5 @@
 /* nbdkit
- * Copyright (C) 2014-2020 Red Hat Inc.
+ * Copyright (C) 2014-2022 Red Hat Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -165,6 +165,16 @@ ocaml_nbdkit_shutdown (value unitv)
   CAMLreturn (Val_unit);
 }

+/* NB: noalloc function. */
+NBDKIT_DLL_PUBLIC value
+ocaml_nbdkit_disconnect (value boolv)
+{
+  CAMLparam1 (boolv);
+
+  nbdkit_disconnect (Bool_val (boolv));
+  CAMLreturn (Val_unit);
+}
+
 /* NB: noalloc function. */
 NBDKIT_DLL_PUBLIC value
 ocaml_nbdkit_debug (value strv)
diff --git a/plugins/python/modfunctions.c b/plugins/python/modfunctions.c
index 4cd45c3b..ac693923 100644
--- a/plugins/python/modfunctions.c
+++ b/plugins/python/modfunctions.c
@@ -93,6 +93,18 @@ do_shutdown (PyObject *self, PyObject *args)
   Py_RETURN_NONE;
 }

+/* nbdkit.disconnect */
+static PyObject *
+do_disconnect (PyObject *self, PyObject *args)
+{
+  int force;
+
+  if (!PyArg_ParseTuple (args, "p:disconnect", &force))
+    return NULL;
+  nbdkit_disconnect (force);
+  Py_RETURN_NONE;
+}
+
 /* nbdkit.parse_size */
 static PyObject *
 parse_size (PyObject *self, PyObject *args)
@@ -121,6 +133,8 @@ static PyMethodDef NbdkitMethods[] = {
     "Store an errno value prior to throwing an exception" },
   { "shutdown", do_shutdown, METH_NOARGS,
     "Request asynchronous shutdown" },
+  { "disconnect", do_disconnect, METH_VARARGS,
+    "Request disconnection from current client" },
   { NULL }
 };

diff --git a/plugins/rust/src/lib.rs b/plugins/rust/src/lib.rs
index 128334ef..a5b88e85 100644
--- a/plugins/rust/src/lib.rs
+++ b/plugins/rust/src/lib.rs
@@ -1046,6 +1046,7 @@ extern "C" {
     fn nbdkit_peer_name( addr: *mut libc::sockaddr,
                          addrlen: *mut libc::socklen_t) -> c_int;
     fn nbdkit_shutdown();
+    fn nbdkit_disconnect(force: bool);
     fn nbdkit_stdio_safe() -> c_int;
 }

@@ -1106,6 +1107,11 @@ pub fn shutdown() {
     unsafe { nbdkit_shutdown() };
 }

+/// Request nbdkit to disconnect the current client.
+pub fn disconnect(force: bool) {
+    unsafe { nbdkit_disconnect(force) };
+}
+
 #[doc(hidden)]
 #[repr(C)]
 pub struct Plugin {
diff --git a/tests/test-disconnect-tls.sh b/tests/test-disconnect-tls.sh
new file mode 100755
index 00000000..00049b07
--- /dev/null
+++ b/tests/test-disconnect-tls.sh
@@ -0,0 +1,126 @@
+#!/usr/bin/env bash
+# nbdkit
+# Copyright (C) 2019-2022 Red Hat Inc.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+#
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+#
+# * Neither the name of Red Hat nor the names of its contributors may be
+# used to endorse or promote products derived from this software without
+# specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY RED HAT AND CONTRIBUTORS ''AS
IS'' AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL RED HAT OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
+# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+# SUCH DAMAGE.
+
+source ./functions.sh
+set -x
+
+requires nbdsh -c 'exit(not h.supports_tls())'
+
+# Does the nbdkit binary support TLS?
+if ! nbdkit --dump-config | grep -sq tls=yes; then
+    echo "$0: nbdkit built without TLS support"
+    exit 77
+fi
+
+# Did we create the PSK keys file?
+# Probably 'certtool' is missing.
+if [ ! -s keys.psk ]; then
+    echo "$0: PSK keys file was not created by the test harness"
+    exit 77
+fi
+
+plugin=.libs/test-disconnect-plugin.$SOEXT
+requires test -f $plugin
+
+sock=$(mktemp -u /tmp/nbdkit-test-sock.XXXXXX)
+files="disconnect-tls.pid $sock"
+cleanup_fn rm -f $files
+
+# Start nbdkit with the disconnect plugin, which has delayed reads and
+# does disconnect on write based on export name.
+start_nbdkit -P disconnect-tls.pid --tls require --tls-psk=keys.psk \
+             -U $sock $plugin
+
+pid=`cat disconnect-tls.pid`
+
+# We can't use 'nbdsh -u "$uri" because of
nbd_set_uri_allow_local_file.
+# Empty export name does soft disconnect on write; the write and the
+# pending read should still succeed, but second read attempt should fail.
+nbdsh -c '
+import errno
+
+h.set_tls(nbd.TLS_REQUIRE)
+h.set_tls_psk_file("keys.psk")
+h.set_tls_username("qemu")
+h.connect_unix("'"$sock"'")
+
+buf = nbd.Buffer(1)
+c1 = h.aio_pread(buf, 1)
+c2 = h.aio_pwrite(buf, 2)
+h.poll(-1)
+assert h.aio_peek_command_completed() == c2
+h.aio_command_completed(c2)
+c3 = h.aio_pread(buf, 3)
+h.poll(-1)
+assert h.aio_peek_command_completed() == c3
+try:
+  h.aio_command_completed(c3)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ESHUTDOWN
+h.poll(-1)
+assert h.aio_peek_command_completed() == c1
+h.aio_command_completed(c1)
+h.shutdown()
+'
+
+# Non-empty export name does hard disconnect on write. The write and the
+# pending read should fail with lost connection.
+nbdsh -c '
+import errno
+
+h.set_tls(nbd.TLS_REQUIRE)
+h.set_tls_psk_file("keys.psk")
+h.set_tls_username("qemu")
+h.set_export_name("a")
+h.connect_unix("'"$sock"'")
+
+buf = nbd.Buffer(1)
+c1 = h.aio_pread(buf, 1)
+c2 = h.aio_pwrite(buf, 2)
+while h.aio_in_flight() > 1:
+  h.poll(-1)
+assert h.aio_is_ready() is False
+try:
+  h.aio_command_completed(c1)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ENOTCONN
+try:
+  h.aio_command_completed(c2)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ENOTCONN
+'
+
+# nbdkit should still be running
+kill -s 0 $pid
diff --git a/tests/test-disconnect.sh b/tests/test-disconnect.sh
new file mode 100755
index 00000000..1551dc03
--- /dev/null
+++ b/tests/test-disconnect.sh
@@ -0,0 +1,100 @@
+#!/usr/bin/env bash
+# nbdkit
+# Copyright (C) 2019-2022 Red Hat Inc.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+#
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+#
+# * Neither the name of Red Hat nor the names of its contributors may be
+# used to endorse or promote products derived from this software without
+# specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY RED HAT AND CONTRIBUTORS ''AS
IS'' AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL RED HAT OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
+# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+# SUCH DAMAGE.
+
+source ./functions.sh
+set -x
+
+requires_nbdsh_uri
+
+plugin=.libs/test-disconnect-plugin.$SOEXT
+requires test -f $plugin
+
+sock=$(mktemp -u /tmp/nbdkit-test-sock.XXXXXX)
+files="disconnect.pid $sock"
+cleanup_fn rm -f $files
+
+# Start nbdkit with the disconnect plugin, which has delayed reads and
+# does disconnect on write based on export name.
+start_nbdkit -P disconnect.pid -U $sock $plugin
+
+pid=`cat disconnect.pid`
+
+# Empty export name does soft disconnect on write; the write and the
+# pending read should still succeed, but second read attempt should fail.
+nbdsh -u "nbd+unix:///?socket=$sock" -c '
+import errno
+
+buf = nbd.Buffer(1)
+c1 = h.aio_pread(buf, 1)
+c2 = h.aio_pwrite(buf, 2)
+h.poll(-1)
+assert h.aio_peek_command_completed() == c2
+h.aio_command_completed(c2)
+c3 = h.aio_pread(buf, 3)
+h.poll(-1)
+assert h.aio_peek_command_completed() == c3
+try:
+  h.aio_command_completed(c3)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ESHUTDOWN
+h.poll(-1)
+assert h.aio_peek_command_completed() == c1
+h.aio_command_completed(c1)
+h.shutdown()
+'
+
+# Non-empty export name does hard disconnect on write. The write and the
+# pending read should fail with lost connection.
+nbdsh -u "nbd+unix:///a?socket=$sock" -c '
+import errno
+
+buf = nbd.Buffer(1)
+c1 = h.aio_pread(buf, 1)
+c2 = h.aio_pwrite(buf, 2)
+while h.aio_in_flight() > 1:
+  h.poll(-1)
+assert h.aio_is_ready() is False
+try:
+  h.aio_command_completed(c1)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ENOTCONN
+try:
+  h.aio_command_completed(c2)
+  assert False
+except nbd.Error as ex:
+  assert ex.errnum == errno.ENOTCONN
+'
+
+# nbdkit should still be running
+kill -s 0 $pid
diff --git a/tests/test-disconnect-plugin.c b/tests/test-disconnect-plugin.c
new file mode 100644
index 00000000..181b262f
--- /dev/null
+++ b/tests/test-disconnect-plugin.c
@@ -0,0 +1,95 @@
+/* nbdkit
+ * Copyright (C) 2013-2022 Red Hat Inc.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ * * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution.
+ *
+ * * Neither the name of Red Hat nor the names of its contributors may be
+ * used to endorse or promote products derived from this software without
+ * specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY RED HAT AND CONTRIBUTORS ''AS
IS'' AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+ * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+ * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL RED HAT OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
+ * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+ * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include <config.h>
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <stdbool.h>
+
+#include <nbdkit-plugin.h>
+
+static void
+disconnect_unload (void)
+{
+  nbdkit_debug ("clean disconnect");
+}
+
+static void *
+disconnect_open (int readonly)
+{
+  return NBDKIT_HANDLE_NOT_NEEDED;
+}
+
+static int64_t
+disconnect_get_size (void *handle)
+{
+  return 1024*1024;
+}
+
+#define THREAD_MODEL NBDKIT_THREAD_MODEL_PARALLEL
+
+/* Reads are delayed to show effect of disconnect on in-flight commands */
+static int
+disconnect_pread (void *handle, void *buf, uint32_t count, uint64_t offset)
+{
+  memset (buf, 0, count);
+  if (nbdkit_nanosleep (2, 0) == -1)
+    nbdkit_debug ("read delay ended early, returning success
anyway");
+  return 0;
+}
+
+/* Writing causes a disconnect; export name determines severity. */
+static int
+disconnect_pwrite (void *handle, const void *buf, uint32_t count,
+                   uint64_t offset)
+{
+  const char *name = nbdkit_export_name ();
+  bool hard = name && *name;
+  nbdkit_debug ("%s disconnect triggered!", hard ? "hard" :
"soft");
+  nbdkit_disconnect (hard);
+  /* Despite the disconnect, we still claim the write succeeded */
+  return 0;
+}
+
+static struct nbdkit_plugin plugin = {
+  .name              = "disconnect",
+  .version           = PACKAGE_VERSION,
+  .unload            = disconnect_unload,
+  .open              = disconnect_open,
+  .get_size          = disconnect_get_size,
+  .pread             = disconnect_pread,
+  .pwrite            = disconnect_pwrite,
+};
+
+NBDKIT_REGISTER_PLUGIN(plugin)
--
2.37.3