Break out tpp load and save methods

4 files changed, 207 insertions, 77 deletions

diff --git a/src/Makefile.am b/src/Makefile.am
index 41eeecf..607a056 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -38,7 +38,8 @@ libtp_a_SOURCES = args.cc args.hh xdg.cc xdg.hh terminal.cc terminal.hh \
 			      character.cc character.hh config.cc config.hh \
                   strings.cc strings.hh io.cc io.hh looper.cc looper.hh \
 				  buffer.cc buffer.hh chunked.cc chunked.hh \
-                  package.cc package.hh data.hh common.hh utf.hh utf.cc
+                  package.cc package.hh data.hh common.hh utf.hh utf.cc \
+				  packages.cc packages.hh
 if !HAVE_SSL
 libtp_a_SOURCES += mitm_stub.cc
 endif
diff --git a/src/monitor-gui.cc b/src/monitor-gui.cc
index 6805814..b02c4de 100644
--- a/src/monitor-gui.cc
+++ b/src/monitor-gui.cc
@@ -39,6 +39,7 @@
 #include "looper.hh"
 #include "monitor.hh"
 #include "observers.hh"
+#include "packages.hh"
 #include "protocols.hh"
 #include "proxy.hh"
 #include "resolver.hh"
@@ -682,6 +683,24 @@ private:
     }
   };
 
+  class PackagesReaderDelegate : public PackagesReader::Delegate {
+  public:
+    PackagesReaderDelegate(PackageList* packages)
+      : packages_(packages) {
+    }
+
+    void package(Package const& package) override {
+      packages_->package(package);
+    }
+
+    void data(uint32_t id, char const* data, size_t size, bool last) override {
+      packages_->package_data(id, data, size, last);
+    }
+
+  private:
+    PackageList* const packages_;
+  };
+
 public:
   MonitorGui()
     : packages_(new PackageList()),
@@ -1593,65 +1612,18 @@ private:
       show_error("Unable to open " + file + " for reading");
       return false;
     }
-    uint8_t header[8];
-    in.read(reinterpret_cast<char*>(header), 8);
-    if (!in.good() || memcmp(header, "TPP", 3) || header[3] != 1) {
-      show_error("Not a valid package file: " +file);
+    auto delegate = std::unique_ptr<PackagesReaderDelegate>(
+        new PackagesReaderDelegate(packages));
+    switch (PackagesReader::read(in, delegate.get())) {
+    case PackagesReader::INVALID:
+      show_error("Not a valid package file: " + file);
       return false;
-    }
-    uint32_t count = read_u32(header + 4);
-    bool good = true;
-    while (count--) {
-      good = false;
-      in.read(reinterpret_cast<char*>(header), 8);
-      if (!in.good()) break;
-      size_t size = read_u64(header);
-      if (size == 0) break;
-      uint8_t buf[8192];
-      size_t avail = std::min(size, sizeof(buf));
-      in.read(reinterpret_cast<char*>(buf), avail);
-      if (!in.good()) break;
-      ::Package pkg;
-      auto pkg_size = read_package(&pkg, buf, avail);
-      if (pkg_size == 0) {
-        if (avail == size) break;
-        size_t need = std::min(static_cast<size_t>(1024) * 1024, size);
-        std::unique_ptr<uint8_t[]> mem(new uint8_t[need]);
-        memcpy(mem.get(), buf, avail);
-        in.read(reinterpret_cast<char*>(mem.get()) + avail, need - avail);
-        if (!in.good()) break;
-        pkg_size = read_package(&pkg, mem.get(), need);
-        if (pkg_size == 0) break;
-        packages->package(pkg);
-        packages->package_data(pkg.id, reinterpret_cast<char*>(mem.get())
-                               + pkg_size, need - pkg_size, need == size);
-        size -= need;
-      } else {
-        packages->package(pkg);
-        packages->package_data(pkg.id, reinterpret_cast<char*>(buf) + pkg_size,
-                               avail - pkg_size, avail == size);
-        size -= avail;
-      }
-      while (size) {
-        avail = std::min(sizeof(buf), size);
-        in.read(reinterpret_cast<char*>(buf), avail);
-        if (!in.good()) break;
-        packages->package_data(pkg.id, reinterpret_cast<char*>(buf), avail,
-                               avail == size);
-        size -= avail;
-      }
-      if (size) break;
-      good = true;
-    }
-    if (good) {
-      in.read(reinterpret_cast<char*>(header), 8);
-      if (!in.good() || read_u64(header) != 0) good = false;
-    }
-    if (!good) {
+    case PackagesReader::IO_ERROR:
       show_error("Error reading from " + file);
       return false;
+    case PackagesReader::GOOD:
+      return true;
     }
-    return true;
   }
 
   bool save(PackageList const* packages, std::string const& file) {
@@ -1661,30 +1633,13 @@ private:
       show_error("Unable to open " + file + " for writing");
       return false;
     }
-    uint8_t header[8];
-    memcpy(header, "TPP", 3);
-    header[3] = 0x1;  // Version
-    write_u32(header + 4, packages->rows());  // Count
-    out.write(reinterpret_cast<char*>(header), 8);
+    auto writer = std::unique_ptr<PackagesWriter>(
+        PackagesWriter::create(packages->rows(), &out));
     for (size_t i = 0; i < packages->rows(); ++i) {
       auto const& pkg = packages->package(i);
-      uint8_t buf[8192];
-      std::unique_ptr<uint8_t[]> backup;
-      uint8_t* ptr = buf;
-      size_t need = write_package(pkg.pkg, buf, sizeof(buf));
-      if (need > sizeof(buf)) {
-        backup.reset(new uint8_t[need]);
-        ptr = backup.get();
-        write_package(pkg.pkg, ptr, need);
-      }
-      write_u64(header, need + pkg.data.size());
-      out.write(reinterpret_cast<char*>(header), 8);
-      out.write(reinterpret_cast<char*>(ptr), need);
-      backup.reset();
-      out.write(pkg.data.data(), pkg.data.size());
+      writer->write(pkg.pkg, pkg.data);
     }
-    write_u64(header, 0);  // EOF
-    out.write(reinterpret_cast<char*>(header), 8);
+    writer.reset();
     if (!out.good()) {
       show_error("Error writing " + file);
       return false;
diff --git a/src/packages.cc b/src/packages.cc
new file mode 100644
index 0000000..0309df1
--- /dev/null
+++ b/src/packages.cc
@@ -0,0 +1,125 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#include "common.hh"
+
+#include <istream>
+#include <memory>
+#include <ostream>
+#include <string.h>
+
+#include "data.hh"
+#include "packages.hh"
+
+namespace {
+
+class PackagesWriterImpl : public PackagesWriter {
+public:
+  PackagesWriterImpl(size_t count, std::ostream* out)
+    : count_(count), out_(out) {
+    uint8_t header[8];
+    memcpy(header, "TPP", 3);
+    header[3] = 0x1;  // Version
+    write_u32(header + 4, count_);  // Count
+    out->write(reinterpret_cast<char*>(header), 8);
+
+    if (count == 0) {
+      write_u64(header, 0);  // EOF
+      out_->write(reinterpret_cast<char*>(header), 8);
+    }
+  }
+
+  ~PackagesWriterImpl() {
+    assert(count_ == 0);
+  }
+
+  void write(Package const& package, std::string const& data) override {
+    if (count_ == 0) {
+      assert(false);
+      return;
+    }
+    uint8_t buf[8192];
+    std::unique_ptr<uint8_t[]> backup;
+    uint8_t* ptr = buf;
+    size_t need = write_package(package, buf, sizeof(buf));
+    if (need > sizeof(buf)) {
+      backup.reset(new uint8_t[need]);
+      ptr = backup.get();
+      write_package(package, ptr, need);
+    }
+    uint8_t size[8];
+    write_u64(size, need + data.size());
+    out_->write(reinterpret_cast<char*>(size), 8);
+    out_->write(reinterpret_cast<char*>(ptr), need);
+    backup.reset();
+    out_->write(data.data(), data.size());
+
+    if (--count_ == 0) {
+      write_u64(size, 0);  // EOF
+      out_->write(reinterpret_cast<char*>(size), 8);
+    }
+  }
+
+private:
+  size_t count_;
+  std::ostream* const out_;
+};
+
+}  // namespace
+
+// static
+PackagesWriter* PackagesWriter::create(size_t count, std::ostream* out) {
+  return new PackagesWriterImpl(count, out);
+}
+
+// static
+PackagesReader::Status PackagesReader::read(std::istream& in,
+                                            Delegate* delegate) {
+  uint8_t header[8];
+  in.read(reinterpret_cast<char*>(header), 8);
+  if (!in.good() || memcmp(header, "TPP", 3) || header[3] != 1) {
+    return INVALID;
+  }
+  auto count = read_u32(header + 4);
+  while (count--) {
+    in.read(reinterpret_cast<char*>(header), 8);
+    if (!in.good()) return IO_ERROR;
+    size_t size = read_u64(header);
+    if (size == 0) return INVALID;
+    uint8_t buf[8192];
+    size_t avail = std::min(size, sizeof(buf));
+    in.read(reinterpret_cast<char*>(buf), avail);
+    if (!in.good()) return IO_ERROR;
+    Package pkg;
+    auto pkg_size = read_package(&pkg, buf, avail);
+    if (pkg_size == 0) {
+      if (avail == size) return INVALID;
+      size_t need = std::min(static_cast<size_t>(1024) * 1024, size);
+      std::unique_ptr<uint8_t[]> mem(new uint8_t[need]);
+      memcpy(mem.get(), buf, avail);
+      in.read(reinterpret_cast<char*>(mem.get()) + avail, need - avail);
+      if (!in.good()) return IO_ERROR;
+      pkg_size = read_package(&pkg, mem.get(), need);
+      if (pkg_size == 0) return INVALID;
+      delegate->package(pkg);
+      delegate->data(pkg.id, reinterpret_cast<char*>(mem.get())
+                     + pkg_size, need - pkg_size, need == size);
+      size -= need;
+    } else {
+        delegate->package(pkg);
+        delegate->data(pkg.id, reinterpret_cast<char*>(buf) + pkg_size,
+                       avail - pkg_size, avail == size);
+        size -= avail;
+    }
+    while (size) {
+      avail = std::min(sizeof(buf), size);
+      in.read(reinterpret_cast<char*>(buf), avail);
+      if (!in.good()) return IO_ERROR;
+      delegate->data(pkg.id, reinterpret_cast<char*>(buf), avail,
+                     avail == size);
+      size -= avail;
+    }
+  }
+  in.read(reinterpret_cast<char*>(header), 8);
+  if (!in.good() || read_u64(header) != 0) return INVALID;
+  return GOOD;
+}
diff --git a/src/packages.hh b/src/packages.hh
new file mode 100644
index 0000000..e646f92
--- /dev/null
+++ b/src/packages.hh
@@ -0,0 +1,49 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#ifndef PACKAGES_HH
+#define PACKAGES_HH
+
+#include "package.hh"
+
+class PackagesWriter {
+public:
+  virtual ~PackagesWriter() {}
+
+  static PackagesWriter* create(size_t count, std::ostream* out);
+
+  virtual void write(Package const& package, std::string const& data) = 0;
+
+protected:
+  PackagesWriter() {}
+  PackagesWriter(PackagesWriter const&) = delete;
+  PackagesWriter& operator=(PackagesWriter const&) = delete;
+};
+
+class PackagesReader {
+public:
+  enum Status {
+    GOOD,
+    INVALID,
+    IO_ERROR,
+  };
+
+  class Delegate {
+  public:
+    virtual ~Delegate() {}
+
+    virtual void package(Package const& package) = 0;
+    virtual void data(uint32_t id, char const* data, size_t size,
+                      bool last) = 0;
+
+  protected:
+    Delegate() {}
+  };
+
+  static Status read(std::istream& in, Delegate* delegate);
+
+private:
+  ~PackagesReader() {}
+  PackagesReader() {}
+};
+
+#endif  // PACKAGES_HH