nbd: Use low-level QIOChannel API in nbd_read_eof()

Instead of using the convenience wrapper qio_channel_read_all_eof(), use
the lower level QIOChannel API. This means duplicating some code, but
we'll need this because this coroutine yield is special: We want it to
be interruptible so that nbd_client_attach_aio_context() can correctly
reenter the coroutine.

This moves the bdrv_dec/inc_in_flight() pair into nbd_read_eof(), so
that connection_co will always sit in this exact qio_channel_yield()
call when bdrv_drain() returns.

Signed-off-by: Kevin Wolf <kwolf@redhat.com>
Reviewed-by: Eric Blake <eblake@redhat.com>

3 files changed, 40 insertions, 18 deletions

diff --git a/block/nbd-client.c b/block/nbd-client.c
index 5ce4aa9520..60f38f0320 100644
--- a/block/nbd-client.c
+++ b/block/nbd-client.c
@@ -84,15 +84,9 @@ static coroutine_fn void nbd_connection_entry(void *opaque)
          *
          * Therefore we keep an additional in_flight reference all the time and
          * only drop it temporarily here.
-         *
-         * FIXME This is not safe because the QIOChannel could wake up the
-         * coroutine for a second time; it is not prepared for coroutine
-         * resumption from external code.
          */
-        bdrv_dec_in_flight(s->bs);
         assert(s->reply.handle == 0);
-        ret = nbd_receive_reply(s->ioc, &s->reply, &local_err);
-        bdrv_inc_in_flight(s->bs);
+        ret = nbd_receive_reply(s->bs, s->ioc, &s->reply, &local_err);
 
         if (local_err) {
             trace_nbd_read_reply_entry_fail(ret, error_get_pretty(local_err));
diff --git a/include/block/nbd.h b/include/block/nbd.h
index cad975e00c..c6ef1ef42e 100644
--- a/include/block/nbd.h
+++ b/include/block/nbd.h
@@ -300,8 +300,8 @@ int nbd_receive_export_list(QIOChannel *ioc, QCryptoTLSCreds *tlscreds,
 int nbd_init(int fd, QIOChannelSocket *sioc, NBDExportInfo *info,
              Error **errp);
 int nbd_send_request(QIOChannel *ioc, NBDRequest *request);
-int coroutine_fn nbd_receive_reply(QIOChannel *ioc, NBDReply *reply,
-                                   Error **errp);
+int coroutine_fn nbd_receive_reply(BlockDriverState *bs, QIOChannel *ioc,
+                                   NBDReply *reply, Error **errp);
 int nbd_client(int fd);
 int nbd_disconnect(int fd);
 int nbd_errno_to_system_errno(int err);
diff --git a/nbd/client.c b/nbd/client.c
index 28d174c0f3..de7da48246 100644
--- a/nbd/client.c
+++ b/nbd/client.c
@@ -1394,30 +1394,58 @@ static int nbd_receive_structured_reply_chunk(QIOChannel *ioc,
  *         negative errno on failure (errp is set)
  */
 static inline int coroutine_fn
-nbd_read_eof(QIOChannel *ioc, void *buffer, size_t size, Error **errp)
+nbd_read_eof(BlockDriverState *bs, QIOChannel *ioc, void *buffer, size_t size,
+             Error **errp)
 {
-    int ret;
+    bool partial = false;
 
     assert(size);
-    ret = qio_channel_read_all_eof(ioc, buffer, size, errp);
-    if (ret < 0) {
-        ret = -EIO;
+    while (size > 0) {
+        struct iovec iov = { .iov_base = buffer, .iov_len = size };
+        ssize_t len;
+
+        len = qio_channel_readv(ioc, &iov, 1, errp);
+        if (len == QIO_CHANNEL_ERR_BLOCK) {
+            bdrv_dec_in_flight(bs);
+            qio_channel_yield(ioc, G_IO_IN);
+            bdrv_inc_in_flight(bs);
+            continue;
+        } else if (len < 0) {
+            return -EIO;
+        } else if (len == 0) {
+            if (partial) {
+                error_setg(errp,
+                           "Unexpected end-of-file before all bytes were read");
+                return -EIO;
+            } else {
+                return 0;
+            }
+        }
+
+        partial = true;
+        size -= len;
+        buffer = (uint8_t*) buffer + len;
     }
-    return ret;
+    return 1;
 }
 
 /* nbd_receive_reply
+ *
+ * Decreases bs->in_flight while waiting for a new reply. This yield is where
+ * we wait indefinitely and the coroutine must be able to be safely reentered
+ * for nbd_client_attach_aio_context().
+ *
  * Returns 1 on success
  *         0 on eof, when no data was read (errp is not set)
  *         negative errno on failure (errp is set)
  */
-int coroutine_fn nbd_receive_reply(QIOChannel *ioc, NBDReply *reply,
-                                   Error **errp)
+int coroutine_fn nbd_receive_reply(BlockDriverState *bs, QIOChannel *ioc,
+                                   NBDReply *reply, Error **errp)
 {
     int ret;
     const char *type;
 
-    ret = nbd_read_eof(ioc, &reply->magic, sizeof(reply->magic), errp);
+    ret = nbd_read_eof(bs, ioc, &reply->magic, sizeof(reply->magic), errp);
     if (ret <= 0) {
         return ret;
     }