--- Linux-NBD/NBD.xs	2010/09/20 11:14:25	1.5
+++ Linux-NBD/NBD.xs	2010/09/21 11:31:05	1.6
@@ -22,8 +22,16 @@
 error, you should not exist
 #endif
 
+struct rstate {
+  struct nbd_request req;
+  u32 req_read; /* how many octets of req are valid */
+  u32 data_read; /* how many octets of the data sv are valid */
+};
+
 MODULE = Linux::NBD		PACKAGE = Linux::NBD::Client
 
+PROTOTYPES: DISABLE
+
 void
 _set_sock (int dev, int fd)
 	CODE:
@@ -88,40 +96,122 @@
 _one_request (SV *obj, int fd)
 	CODE:
 {
-        struct nbd_request req;
+	struct rstate *s;
+        u64 from;
+        u32 len;
+        char *method;
+	MAGIC *mg = mg_find (SvRV (obj), PERL_MAGIC_ext);
+
+        if (!mg)
+          {
+            mg = sv_magicext (SvRV (obj), 0, PERL_MAGIC_ext, 0, 0, 0);
+            mg->mg_len = sizeof (struct rstate);
+            New (0, mg->mg_ptr, mg->mg_len, char);
+
+            ((struct rstate *)mg->mg_ptr)->req_read = 0; /* initialise the state machine */
+          }
+
+        s = (struct rstate *)mg->mg_ptr;
 
-        if (read (fd, &req, sizeof (req)) == sizeof (req))
+        if (s->req_read < sizeof (struct nbd_request))
           {
-            if (req.magic == htonl (NBD_REQUEST_MAGIC))
-              {
-                req.type = htonl (req.type);
-
-                if (req.type < 2)
-                  {
-                    u64 from = ntohll (req.from);
-
-                    PUSHMARK (SP);
-                    EXTEND (SP, 3);
-                    PUSHs (obj);
-                    PUSHs (sv_2mortal (newSVpvn (req.handle, sizeof (req.handle))));
-                    PUSHs (sv_2mortal (sizeof (UV) < 8 && from > (0xffffffffUL)
-                                       ? newSVnv (from)
-                                       : newSVuv (from)));
-                    PUSHs (sv_2mortal (newSVuv (ntohl (req.len))));
-                    PUTBACK;
-                    call_method (req.type ? "req_write" : "req_read", G_DISCARD);
-                    SPAGAIN;
-
-                    XSRETURN_NO;
-                  }
-              }
+            int res = read (fd, s->req_read + (char *)&s->req, sizeof (struct nbd_request) - s->req_read);
+
+            if (res > 0)
+              s->req_read += res;
+            else if (res == 0)
+              XSRETURN_UNDEF; /* should req->eof */
+            else if (errno != EAGAIN && errno != EWOULDBLOCK)
+              XSRETURN_UNDEF; /* should req->error */
+
+            s->data_read = 0;
           }
 
+        if (s->req_read < sizeof (struct nbd_request))
+          XSRETURN_NO;
+
+        /* now we have a full request, so check for data */
+        if (s->req.magic != htonl (NBD_REQUEST_MAGIC))
+          croak ("Linux::NBD::Server received illegal request magic %08lx - protocol error.\n", ntohl (s->req.magic));
+
+        from = ntohll (s->req.from);
+        len  = ntohl  (s->req.len);
+
+        switch (ntohl (s->req.type))
+          {
+            case NBD_CMD_WRITE:
+              if (!mg->mg_obj)
+                {
+                  mg->mg_flags |= MGf_REFCOUNTED;
+                  mg->mg_obj = NEWSV (0, len);
+                  SvPOK_only (mg->mg_obj);
+                  SvCUR_set (mg->mg_obj, len);
+                }
+
+              if (s->data_read < len)
+                {
+                  int res = read (fd, s->data_read + SvPVX (mg->mg_obj), len - s->data_read);
+
+                  if (res > 0)
+                    s->data_read += res;
+                  else if (res == 0)
+                    XSRETURN_UNDEF; /* should req->eof */
+                  else if (errno != EAGAIN && errno != EWOULDBLOCK)
+                    XSRETURN_UNDEF; /* should req->error */
+                }
+
+              if (s->data_read < len)
+                XSRETURN_NO;
+
+              /* fallthrough */
+            case NBD_CMD_READ:
+              s->req_read = 0;
+
+              PUSHMARK (SP);
+              EXTEND (SP, 4);
+              PUSHs (obj);
+              PUSHs (sv_2mortal (newSVpvn (s->req.handle, sizeof (s->req.handle))));
+              PUSHs (sv_2mortal (sizeof (UV) < 8 && from > 0xffffffffUL
+                                 ? newSVnv (from)
+                                 : newSVuv (from)));
+
+              if (mg->mg_obj)
+                {
+                  PUSHs (sv_2mortal (mg->mg_obj));
+                  mg->mg_obj = 0;
+                  method = "req_write";
+                }
+              else
+                {
+                  PUSHs (sv_2mortal (newSVuv (len)));
+                  method = "req_read";
+                }
+
+              break;
+
+            case NBD_CMD_DISC:
+              s->req_read = 0;
+
+              method = "req_disc";
+
+              PUSHMARK (SP);
+              XPUSHs (obj);
+
+              break;
+
+            default:
+              croak ("Linux::NBD::Server received unsupported request type %d.\n", ntohl (s->req.type));
+          }
+
+        PUTBACK;
+        call_method (method, G_DISCARD);
+        SPAGAIN;
+
         XSRETURN_YES;
 }
 
 SV *
-_format_reply (SV *handle, unsigned int error = 0, SV *data = 0)
+format_reply (SV *unused, SV *handle, unsigned int error = 0, SV *data = 0)
 	CODE:
 {
         struct nbd_reply rep;