2 files changed, 44 insertions, 29 deletions

diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index 76f2e721a..640267a13 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -23,8 +23,9 @@ void BufferedSink::operator () (const unsigned char * data, size_t len)
     
     while (len) {
         /* Optimisation: bypass the buffer if the data exceeds the
-           buffer size and there is no unflushed data. */
-        if (bufPos == 0 && len >= bufSize) {
+           buffer size. */
+        if (bufPos + len >= bufSize) {
+            flush();
             write(data, len);
             break;
         }
@@ -59,29 +60,37 @@ void FdSink::write(const unsigned char * data, size_t len)
 }
 
 
+void Source::operator () (unsigned char * data, size_t len)
+{
+    while (len) {
+        size_t n = read(data, len);
+        data += n; len -= n;
+    }
+}
+
+
 BufferedSource::~BufferedSource()
 {
     if (buffer) delete[] buffer;
 }
 
 
-void BufferedSource::operator () (unsigned char * data, size_t len)
+size_t BufferedSource::read(unsigned char * data, size_t len)
 {
     if (!buffer) buffer = new unsigned char[bufSize];
 
-    while (len) {
-        if (!bufPosIn) bufPosIn = read(buffer, bufSize);
+    if (!bufPosIn) bufPosIn = readUnbuffered(buffer, bufSize);
             
-        /* Copy out the data in the buffer. */
-        size_t n = len > bufPosIn - bufPosOut ? bufPosIn - bufPosOut : len;
-        memcpy(data, buffer + bufPosOut, n);
-        data += n; bufPosOut += n; len -= n;
-        if (bufPosIn == bufPosOut) bufPosIn = bufPosOut = 0;
-    }
+    /* Copy out the data in the buffer. */
+    size_t n = len > bufPosIn - bufPosOut ? bufPosIn - bufPosOut : len;
+    memcpy(data, buffer + bufPosOut, n);
+    bufPosOut += n;
+    if (bufPosIn == bufPosOut) bufPosIn = bufPosOut = 0;
+    return n;
 }
 
 
-size_t FdSource::read(unsigned char * data, size_t len)
+size_t FdSource::readUnbuffered(unsigned char * data, size_t len)
 {
     ssize_t n;
     do {
@@ -94,6 +103,15 @@ size_t FdSource::read(unsigned char * data, size_t len)
 }
 
 
+size_t StringSource::read(unsigned char * data, size_t len)
+{
+    if (pos == s.size()) throw EndOfFile("end of string reached");
+    size_t n = s.copy((char *) data, len, pos);
+    pos += n;
+    return n;
+}
+
+
 void writePadding(size_t len, Sink & sink)
 {
     if (len % 8) {
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index a155f6681..25398b09d 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -24,7 +24,7 @@ struct BufferedSink : Sink
     BufferedSink(size_t bufSize = 32 * 1024)
         : bufSize(bufSize), bufPos(0), buffer(0) { }
     ~BufferedSink();
-    
+
     void operator () (const unsigned char * data, size_t len);
     
     void flush();
@@ -39,9 +39,14 @@ struct Source
     virtual ~Source() { }
     
     /* Store exactly ‘len’ bytes in the buffer pointed to by ‘data’.
-       It blocks if that much data is not yet available, or throws an
-       error if it is not going to be available. */
-    virtual void operator () (unsigned char * data, size_t len) = 0;
+       It blocks until all the requested data is available, or throws
+       an error if it is not going to be available.   */
+    void operator () (unsigned char * data, size_t len);
+
+    /* Store up to ‘len’ in the buffer pointed to by ‘data’, and
+       return the number of bytes stored.  If blocks until at least
+       one byte is available. */
+    virtual size_t read(unsigned char * data, size_t len) = 0;
 };
 
 
@@ -55,12 +60,10 @@ struct BufferedSource : Source
         : bufSize(bufSize), bufPosIn(0), bufPosOut(0), buffer(0) { }
     ~BufferedSource();
     
-    void operator () (unsigned char * data, size_t len);
+    size_t read(unsigned char * data, size_t len);
     
-    /* Store up to ‘len’ in the buffer pointed to by ‘data’, and
-       return the number of bytes stored.  If should block until at
-       least one byte is available. */
-    virtual size_t read(unsigned char * data, size_t len) = 0;
+    /* Underlying read call, to be overriden. */
+    virtual size_t readUnbuffered(unsigned char * data, size_t len) = 0;
 };
 
 
@@ -83,7 +86,7 @@ struct FdSource : BufferedSource
     int fd;
     FdSource() : fd(-1) { }
     FdSource(int fd) : fd(fd) { }
-    size_t read(unsigned char * data, size_t len);
+    size_t readUnbuffered(unsigned char * data, size_t len);
 };
 
 
@@ -104,13 +107,7 @@ struct StringSource : Source
     const string & s;
     size_t pos;
     StringSource(const string & _s) : s(_s), pos(0) { }
-    virtual void operator () (unsigned char * data, size_t len)
-    {
-        s.copy((char *) data, len, pos);
-        pos += len;
-        if (pos > s.size())
-            throw Error("end of string reached");
-    }
+    size_t read(unsigned char * data, size_t len);    
 };