1 files changed, 129 insertions, 15 deletions

diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index 9b422271323f..c4563ffd1212 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -2,24 +2,117 @@
 #include "util.hh"
 
 #include <cstring>
+#include <cerrno>
 
 
 namespace nix {
 
 
-void FdSink::operator () (const unsigned char * data, unsigned int len)
+BufferedSink::~BufferedSink()
+{
+    /* We can't call flush() here, because C++ for some insane reason
+       doesn't allow you to call virtual methods from a destructor. */
+    assert(!bufPos);
+    if (buffer) delete[] buffer;
+}
+
+    
+void BufferedSink::operator () (const unsigned char * data, size_t len)
+{
+    if (!buffer) buffer = new unsigned char[bufSize];
+    
+    while (len) {
+        /* Optimisation: bypass the buffer if the data exceeds the
+           buffer size. */
+        if (bufPos + len >= bufSize) {
+            flush();
+            write(data, len);
+            break;
+        }
+        /* Otherwise, copy the bytes to the buffer.  Flush the buffer
+           when it's full. */
+        size_t n = bufPos + len > bufSize ? bufSize - bufPos : len;
+        memcpy(buffer + bufPos, data, n);
+        data += n; bufPos += n; len -= n;
+        if (bufPos == bufSize) flush();
+    }
+}
+
+
+void BufferedSink::flush()
+{
+    if (bufPos == 0) return;
+    size_t n = bufPos;
+    bufPos = 0; // don't trigger the assert() in ~BufferedSink()
+    write(buffer, n);
+}
+
+
+FdSink::~FdSink()
+{
+    try { flush(); } catch (...) { ignoreException(); }
+}
+
+
+void FdSink::write(const unsigned char * data, size_t len)
 {
     writeFull(fd, data, len);
 }
 
 
-void FdSource::operator () (unsigned char * data, unsigned int len)
+void Source::operator () (unsigned char * data, size_t len)
+{
+    while (len) {
+        size_t n = read(data, len);
+        data += n; len -= n;
+    }
+}
+
+
+BufferedSource::~BufferedSource()
 {
-    readFull(fd, data, len);
+    if (buffer) delete[] buffer;
 }
 
 
-void writePadding(unsigned int len, Sink & sink)
+size_t BufferedSource::read(unsigned char * data, size_t len)
+{
+    if (!buffer) buffer = new unsigned char[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);
+    bufPosOut += n;
+    if (bufPosIn == bufPosOut) bufPosIn = bufPosOut = 0;
+    return n;
+}
+
+
+size_t FdSource::readUnbuffered(unsigned char * data, size_t len)
+{
+    ssize_t n;
+    do {
+        checkInterrupt();
+        n = ::read(fd, (char *) data, bufSize);
+    } while (n == -1 && errno == EINTR);
+    if (n == -1) throw SysError("reading from file");
+    if (n == 0) throw EndOfFile("unexpected end-of-file");
+    return n;
+}
+
+
+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) {
         unsigned char zero[8];
@@ -56,28 +149,36 @@ void writeLongLong(unsigned long long n, Sink & sink)
 }
 
 
-void writeString(const string & s, Sink & sink)
+void writeString(const unsigned char * buf, size_t len, Sink & sink)
 {
-    unsigned int len = s.length();
     writeInt(len, sink);
-    sink((const unsigned char *) s.c_str(), len);
+    sink(buf, len);
     writePadding(len, sink);
 }
 
 
-void writeStringSet(const StringSet & ss, Sink & sink)
+void writeString(const string & s, Sink & sink)
+{
+    writeString((const unsigned char *) s.c_str(), s.size(), sink);
+}
+
+
+template<class T> void writeStrings(const T & ss, Sink & sink)
 {
     writeInt(ss.size(), sink);
-    for (StringSet::iterator i = ss.begin(); i != ss.end(); ++i)
+    foreach (typename T::const_iterator, i, ss)
         writeString(*i, sink);
 }
 
+template void writeStrings(const Paths & ss, Sink & sink);
+template void writeStrings(const PathSet & ss, Sink & sink);
 
-void readPadding(unsigned int len, Source & source)
+
+void readPadding(size_t len, Source & source)
 {
     if (len % 8) {
         unsigned char zero[8];
-        unsigned int n = 8 - (len % 8);
+        size_t n = 8 - (len % 8);
         source(zero, n);
         for (unsigned int i = 0; i < n; i++)
             if (zero[i]) throw SerialisationError("non-zero padding");
@@ -115,9 +216,19 @@ unsigned long long readLongLong(Source & source)
 }
 
 
+size_t readString(unsigned char * buf, size_t max, Source & source)
+{
+    size_t len = readInt(source);
+    if (len > max) throw Error("string is too long");
+    source(buf, len);
+    readPadding(len, source);
+    return len;
+}
+
+ 
 string readString(Source & source)
 {
-    unsigned int len = readInt(source);
+    size_t len = readInt(source);
     unsigned char * buf = new unsigned char[len];
     AutoDeleteArray<unsigned char> d(buf);
     source(buf, len);
@@ -126,14 +237,17 @@ string readString(Source & source)
 }
 
  
-StringSet readStringSet(Source & source)
+template<class T> T readStrings(Source & source)
 {
     unsigned int count = readInt(source);
-    StringSet ss;
+    T ss;
     while (count--)
-        ss.insert(readString(source));
+        ss.insert(ss.end(), readString(source));
     return ss;
 }
 
+template Paths readStrings(Source & source);
+template PathSet readStrings(Source & source);
+
 
 }