[Libguestfs] [PATCH libnbd] states: In recv_into_rbuf and send_from_wbuf loop until EAGAIN.

[Richard W.M. Jones]

Previously we performed a single call to recv(2) or send(2) (or the
GnuTLS equivalents), and even if more data/space was immediately
available to receive/send we would return to poll.  Instead of this,
loop until the socket returns EAGAIN.
---
 generator/states.c | 91 ++++++++++++++++++++++------------------------
 1 file changed, 43 insertions(+), 48 deletions(-)

diff --git a/generator/states.c b/generator/states.c
index 145e8c1..cde934a 100644
--- a/generator/states.c
+++ b/generator/states.c
@@ -46,43 +46,40 @@ recv_into_rbuf (struct nbd_handle *h)
   void *rbuf;
   size_t rlen;
 
-  if (h->rlen == 0)
-    return 0;                   /* move to next state */
+  while (h->rlen > 0) {
+    /* As a special case h->rbuf is allowed to be NULL, meaning
+     * throw away the data.
+     */
+    if (h->rbuf) {
+      rbuf = h->rbuf;
+      rlen = h->rlen;
+    }
+    else {
+      rbuf = &buf;
+      rlen = h->rlen > sizeof buf ? sizeof buf : h->rlen;
+    }
 
-  /* As a special case h->rbuf is allowed to be NULL, meaning
-   * throw away the data.
-   */
-  if (h->rbuf) {
-    rbuf = h->rbuf;
-    rlen = h->rlen;
-  }
-  else {
-    rbuf = &buf;
-    rlen = h->rlen > sizeof buf ? sizeof buf : h->rlen;
-  }
-
-  r = h->sock->ops->recv (h, h->sock, rbuf, rlen);
-  if (r == -1) {
-    if (errno == EAGAIN || errno == EWOULDBLOCK)
-      return 1;                 /* more data */
-    /* sock->ops->recv called set_error already. */
-    return -1;
-  }
-  if (r == 0) {
-    set_error (0, "recv: server disconnected unexpectedly");
-    return -1;
-  }
+    r = h->sock->ops->recv (h, h->sock, rbuf, rlen);
+    if (r == -1) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK)
+        return 1;      /* more data */
+      /* sock->ops->recv called set_error already. */
+      return -1;
+    }
+    if (r == 0) {
+      set_error (0, "recv: server disconnected unexpectedly");
+      return -1;
+    }
 #ifdef DUMP_PACKETS
-  if (h->rbuf != NULL)
-    nbd_internal_hexdump (h->rbuf, r, stderr);
+    if (h->rbuf != NULL)
+      nbd_internal_hexdump (h->rbuf, r, stderr);
 #endif
-  if (h->rbuf)
-    h->rbuf += r;
-  h->rlen -= r;
-  if (h->rlen == 0)
-    return 0;                   /* move to next state */
-  else
-    return 1;                   /* more data */
+    if (h->rbuf)
+      h->rbuf += r;
+    h->rlen -= r;
+  }
+
+  return 0;                     /* move to next state */
 }
 
 static int
@@ -90,21 +87,19 @@ send_from_wbuf (struct nbd_handle *h)
 {
   ssize_t r;
 
-  if (h->wlen == 0)
-    return 0;                   /* move to next state */
-  r = h->sock->ops->send (h, h->sock, h->wbuf, h->wlen);
-  if (r == -1) {
-    if (errno == EAGAIN || errno == EWOULDBLOCK)
-      return 1;                 /* more data */
-    /* sock->ops->send called set_error already. */
-    return -1;
+  while (h->wlen > 0) {
+    r = h->sock->ops->send (h, h->sock, h->wbuf, h->wlen);
+    if (r == -1) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK)
+        return 1;      /* more data */
+      /* sock->ops->send called set_error already. */
+      return -1;
+    }
+    h->wbuf += r;
+    h->wlen -= r;
   }
-  h->wbuf += r;
-  h->wlen -= r;
-  if (h->wlen == 0)
-    return 0;                   /* move to next state */
-  else
-    return 1;                   /* more data */
+
+  return 0;                     /* move to next state */
 }
 
 /*----- End of prologue. -----*/
-- 
2.21.0