Break out Package read/write

12 files changed, 270 insertions, 123 deletions

diff --git a/src/Makefile.am b/src/Makefile.am
index b25378d..ae46e53 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -32,7 +32,7 @@ libproxy_a_CXXFLAGS = $(AM_CXXFLAGS) -DVERSION='"@VERSION@"' @THREAD_CFLAGS@
 
 libtp_a_SOURCES = args.cc xdg.cc terminal.cc http.cc url.cc paths.cc \
 			      character.cc config.cc strings.cc io.cc looper.cc \
-				  buffer.cc chunked.cc
+				  buffer.cc chunked.cc package.cc
 if !HAVE_SSL
 libtp_a_SOURCES += mitm_stub.cc
 endif
diff --git a/src/data.hh b/src/data.hh
new file mode 100644
index 0000000..e949b47
--- /dev/null
+++ b/src/data.hh
@@ -0,0 +1,35 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#ifndef DATA_HH
+#define DATA_HH
+
+#include <cstdint>
+
+inline uint16_t read_u16(uint8_t const* data) {
+  return data[0] << 8 | data[1];
+}
+
+inline uint32_t read_u32(uint8_t const* data) {
+  return static_cast<uint32_t>(read_u16(data)) << 16 | read_u16(data + 2);
+}
+
+inline uint64_t read_u64(uint8_t const* data) {
+  return static_cast<uint64_t>(read_u32(data)) << 32 | read_u32(data + 4);
+}
+
+inline void write_u16(uint8_t* dst, uint16_t value) {
+  dst[0] = value >> 8;
+  dst[1] = value & 0xff;
+}
+
+inline void write_u32(uint8_t* dst, uint32_t value) {
+  write_u16(dst, value >> 16);
+  write_u16(dst + 2, value & 0xffff);
+}
+
+inline void write_u64(uint8_t* dst, uint64_t value) {
+  write_u32(dst, value >> 32);
+  write_u32(dst + 4, value & 0xffffffff);
+}
+
+#endif  // DATA_HH
diff --git a/src/monitor-cmd.cc b/src/monitor-cmd.cc
index 449656e..ba02703 100644
--- a/src/monitor-cmd.cc
+++ b/src/monitor-cmd.cc
@@ -53,8 +53,7 @@ public:
     std::cerr << "# Error: " << error << std::endl;
   }
 
-  void package(
-      Monitor* UNUSED(monitor), Monitor::Package const& package) override {
+  void package(Monitor* UNUSED(monitor), Package const& package) override {
     packages_.insert(std::make_pair(package.id, package));
   }
 
@@ -93,8 +92,7 @@ public:
   }
 
 private:
-  void print_package(
-      Monitor::Package& pkg, bool last, char const* data, size_t size) {
+  void print_package(Package& pkg, bool last, char const* data, size_t size) {
     if (size == 0 && !last) return;
     {
       ios_save save(out_);
@@ -170,7 +168,7 @@ private:
   bool interleave_;
   Looper* looper_;
   bool attached_;
-  std::unordered_map<uint32_t, Monitor::Package> packages_;
+  std::unordered_map<uint32_t, Package> packages_;
   // Used when interleaving
   std::unordered_map<uint32_t, uint64_t> offset_;
   // Used when not interleaving
diff --git a/src/monitor-gui.cc b/src/monitor-gui.cc
index 29c856f..3494855 100644
--- a/src/monitor-gui.cc
+++ b/src/monitor-gui.cc
@@ -145,7 +145,7 @@ public:
     listeners_.erase(listener);
   }
 
-  void package(Monitor::Package const& package) {
+  void package(::Package const& package) {
     auto const index = packages_.size();
     open_.emplace(package.id, index);
     packages_.emplace_back();
@@ -924,7 +924,7 @@ public:
     assert(monitor == monitor_.get());
   }
 
-  void package(Monitor* monitor, Monitor::Package const& package) override {
+  void package(Monitor* monitor, ::Package const& package) override {
     assert(monitor == monitor_.get());
     packages_->package(package);
   }
diff --git a/src/monitor.cc b/src/monitor.cc
index ba95596..e07b036 100644
--- a/src/monitor.cc
+++ b/src/monitor.cc
@@ -6,6 +6,7 @@
 
 #include "buffer.hh"
 #include "chunked.hh"
+#include "data.hh"
 #include "http.hh"
 #include "io.hh"
 #include "looper.hh"
@@ -150,18 +151,6 @@ private:
     return true;
   }
 
-  static uint64_t read_u64(uint8_t const* data) {
-    return static_cast<uint64_t>(read_u32(data)) << 32 | read_u32(data + 4);
-  }
-
-  static uint32_t read_u32(uint8_t const* data) {
-    return static_cast<uint32_t>(read_u16(data)) << 16 | read_u16(data + 2);
-  }
-
-  static uint16_t read_u16(uint8_t const* data) {
-    return data[0] << 8 | data[1];
-  }
-
   void package(void const* data, size_t size) {
     auto d = reinterpret_cast<char const*>(data);
     auto const end = d + size;
@@ -177,41 +166,20 @@ private:
         uint16_t size = read_u16(package_ + offset + 3);
         if (offset + size > package_fill_) break;
         size_t o = 5;
-        if (size >= 29 && memcmp(package_ + offset, "PKG", 3) == 0) {
+        if (size >= 3 && memcmp(package_ + offset, "PKG", 3) == 0) {
           Package pkg;
-          pkg.id = read_u32(package_ + offset + o);
-          o += 4;
-          pkg.timestamp.tv_sec = read_u64(package_ + offset + o);
-          o += 8;
-          pkg.timestamp.tv_nsec = read_u32(package_ + offset + o);
-          o += 4;
-          pkg.flags = read_u16(package_ + offset + o);
-          o += 2;
-          pkg.source_port = read_u16(package_ + offset + o);
-          o += 2;
-          pkg.target_port = read_u16(package_ + offset + o);
-          o += 2;
-          auto len = read_u16(package_ + offset + o);
-          o += 2;
-          if (o + len + 2 <= size) {
-            pkg.source_host.assign(
-                reinterpret_cast<char*>(package_) + offset + o, len);
-            o += len;
-            len = read_u16(package_ + offset + o);
-            o += 2;
-            if (o + len <= size) {
-              pkg.target_host.assign(
-                  reinterpret_cast<char*>(package_) + offset + o, len);
-              o += len;
-              bool last = !(pkg.flags & 0x01);
-              pkg.flags >>= 1;
-              delegate_->package(this, pkg);
-              if (o < size || last) {
-                delegate_->package_data(
-                    this, pkg.id,
-                    reinterpret_cast<char*>(package_) + offset + o, size - o,
-                    last);
-              }
+          auto ret = read_package(&pkg,
+                                  package_ + offset + o, size - offset - o);
+          if (ret > 0) {
+            o += ret;
+            bool last = !(pkg.flags & 0x01);
+            pkg.flags >>= 1;
+            delegate_->package(this, pkg);
+            if (o < size || last) {
+              delegate_->package_data(
+                  this, pkg.id,
+                  reinterpret_cast<char*>(package_) + offset + o, size - o,
+                  last);
             }
           }
         } else if (size >= 10 && memcmp(package_ + offset, "DAT", 3) == 0) {
diff --git a/src/monitor.hh b/src/monitor.hh
index 6c2a557..7c6cc40 100644
--- a/src/monitor.hh
+++ b/src/monitor.hh
@@ -7,6 +7,8 @@
 #include <memory>
 #include <string>
 
+#include "package.hh"
+
 class Looper;
 class Resolver;
 
@@ -23,16 +25,6 @@ public:
     ATTACHED,
   };
 
-  struct Package {
-    uint32_t id;
-    struct timespec timestamp;
-    uint16_t flags;
-    std::string source_host;
-    uint16_t source_port;
-    std::string target_host;
-    uint16_t target_port;
-  };
-
   class Delegate {
   public:
     virtual ~Delegate() {}
diff --git a/src/package.cc b/src/package.cc
new file mode 100644
index 0000000..8a7ef24
--- /dev/null
+++ b/src/package.cc
@@ -0,0 +1,55 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#include "common.hh"
+
+#include <cstring>
+
+#include "data.hh"
+#include "package.hh"
+
+size_t read_package(Package* pkg, uint8_t const* data, size_t max) {
+  if (max < 26) return 0;
+  size_t offset = 0;
+  pkg->id = read_u32(data + offset);
+  offset += 4;
+  pkg->timestamp.tv_sec = read_u64(data + offset);
+  offset += 8;
+  pkg->timestamp.tv_nsec = read_u32(data + offset);
+  offset += 4;
+  pkg->flags = read_u16(data + offset);
+  offset += 2;
+  pkg->source_port = read_u16(data + offset);
+  offset += 2;
+  pkg->target_port = read_u16(data + offset);
+  offset += 2;
+  auto len = read_u16(data + offset);
+  offset += 2;
+  if (offset + len + 2 > max) return 0;
+  pkg->source_host.assign(reinterpret_cast<char const*>(data) + offset, len);
+  offset += len;
+  len = read_u16(data + offset);
+  offset += 2;
+  if (offset + len > max) return 0;
+  pkg->target_host.assign(reinterpret_cast<char const*>(data) + offset, len);
+  offset += len;
+  return offset;
+}
+
+size_t write_package(Package const& pkg, uint8_t* data, size_t max) {
+  auto len = 26 + pkg.source_host.size() + pkg.target_host.size();
+  if (!data || max < len) {
+    return len;
+  }
+  write_u32(data, pkg.id);
+  write_u64(data + 4, pkg.timestamp.tv_sec);
+  write_u32(data + 12, pkg.timestamp.tv_nsec);
+  write_u16(data + 16, pkg.flags);
+  write_u16(data + 18, pkg.source_port);
+  write_u16(data + 20, pkg.target_port);
+  write_u16(data + 22, pkg.source_host.size());
+  memcpy(data + 24, pkg.source_host.data(), pkg.source_host.size());
+  write_u16(data + 24 + pkg.source_host.size(), pkg.target_host.size());
+  memcpy(data + 26 + pkg.source_host.size(),
+         pkg.target_host.data(), pkg.target_host.size());
+  return len;
+}
diff --git a/src/package.hh b/src/package.hh
new file mode 100644
index 0000000..e83271f
--- /dev/null
+++ b/src/package.hh
@@ -0,0 +1,22 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#ifndef PACKAGE_HH
+#define PACKAGE_HH
+
+#include <cstdint>
+#include <string>
+
+struct Package {
+  uint32_t id;
+  struct timespec timestamp;
+  uint16_t flags;
+  std::string source_host;
+  uint16_t source_port;
+  std::string target_host;
+  uint16_t target_port;
+};
+
+size_t read_package(Package* package, uint8_t const* data, size_t max);
+size_t write_package(Package const& package, uint8_t* data, size_t max);
+
+#endif  // PACKAGE_HH
diff --git a/src/proxy.cc b/src/proxy.cc
index 05a53bb..878b40c 100644
--- a/src/proxy.cc
+++ b/src/proxy.cc
@@ -23,12 +23,14 @@
 #include "buffer.hh"
 #include "chunked.hh"
 #include "config.hh"
+#include "data.hh"
 #include "http.hh"
 #include "io.hh"
 #include "logger.hh"
 #include "looper.hh"
 #include "mitm.hh"
 #include "resolver.hh"
+#include "package.hh"
 #include "paths.hh"
 #include "proxy.hh"
 #include "url.hh"
@@ -379,13 +381,6 @@ private:
                              std::string const& target_host,
                              uint16_t target_port,
                              bool last);
-  void send_attached_package2(uint8_t* buffer, size_t size,
-                              uint32_t id, uint16_t flags,
-                              std::string const& source_host,
-                              uint16_t source_port,
-                              std::string const& target_host,
-                              uint16_t target_port,
-                              bool last);
   void send_attached_data(uint32_t id, void const* ptr, size_t size, bool last);
   void send_attached(void const* header, size_t header_size,
                      void const* data, size_t data_size);
@@ -1850,21 +1845,6 @@ uint32_t ProxyImpl::get_next_package_id() {
   return next_package_id_++;
 }
 
-void write_u16(uint8_t* dst, uint16_t value) {
-  dst[0] = value >> 8;
-  dst[1] = value & 0xff;
-}
-
-void write_u32(uint8_t* dst, uint32_t value) {
-  write_u16(dst, value >> 16);
-  write_u16(dst + 2, value & 0xffff);
-}
-
-void write_u64(uint8_t* dst, uint64_t value) {
-  write_u32(dst, value >> 32);
-  write_u32(dst + 4, value & 0xffffffff);
-}
-
 void ProxyImpl::send_attached_package(uint32_t id, uint16_t flags,
                                       std::string const& source_host,
                                       uint16_t source_port,
@@ -1877,45 +1857,34 @@ void ProxyImpl::send_attached_package(uint32_t id, uint16_t flags,
   }
   if (attached_.empty()) return;
   uint8_t data[256];
-  size_t need = 2 + 3 + 4 + 8 + 4 + 2 + 2 + 2 + 2 + 2 + source_host.size() +
-    target_host.size();
-  if (need <= sizeof(data)) {
-    send_attached_package2(data, need, id, flags, source_host, source_port,
-                           target_host, target_port, last);
-  } else {
-    // TODO: Might need better handling of really long source_host/target_host
-    auto p = std::unique_ptr<uint8_t[]>(new uint8_t[need]);
-    send_attached_package2(p.get(), need, id, flags, source_host, source_port,
-                           target_host, target_port, last);
-  }
-}
 
-void ProxyImpl::send_attached_package2(uint8_t* buffer, size_t size,
-                                       uint32_t id, uint16_t flags,
-                                       std::string const& source_host,
-                                       uint16_t source_port,
-                                       std::string const& target_host,
-                                       uint16_t target_port,
-                                       bool last) {
-  buffer[0] = 'P';
-  buffer[1] = 'K';
-  buffer[2] = 'G';
-  write_u16(buffer + 3, size);
-  write_u32(buffer + 5, id);
+  Package pkg;
+  pkg.id = id;
   auto dur = looper_->now().time_since_epoch();
   auto sec = std::chrono::duration_cast<std::chrono::seconds>(dur);
-  auto nsec = std::chrono::duration_cast<std::chrono::nanoseconds>(dur - sec);
-  write_u64(buffer + 9, sec.count());
-  write_u32(buffer + 17, nsec.count());
-  write_u16(buffer + 21, (last ? 0 : 1) | (flags << 1));
-  write_u16(buffer + 23, source_port);
-  write_u16(buffer + 25, target_port);
-  write_u16(buffer + 27, source_host.size());
-  memcpy(buffer + 29, source_host.data(), source_host.size());
-  write_u16(buffer + 29 + source_host.size(), target_host.size());
-  memcpy(buffer + 31 + source_host.size(),
-         target_host.data(), target_host.size());
-  send_attached(buffer, size, nullptr, 0);
+  pkg.timestamp.tv_sec = sec.count();
+  pkg.timestamp.tv_nsec =
+    std::chrono::duration_cast<std::chrono::nanoseconds>(dur - sec).count();
+  pkg.flags = (last ? 0 : 1) | (flags << 1);
+  pkg.source_port = source_port;
+  pkg.source_host = source_host;
+  pkg.target_port = target_port;
+  pkg.target_host = target_host;
+
+  auto size = 5 + write_package(pkg, data + 5, sizeof(data) - 5);
+  auto ptr = data;
+  std::unique_ptr<uint8_t[]> extra;
+  if (size > sizeof(data)) {
+    // TODO: Might need better handling of really long source_host/target_host
+    extra.reset(new uint8_t[size]);
+    ptr = extra.get();
+    write_package(pkg, ptr + 5, size - 5);
+  }
+  ptr[0] = 'P';
+  ptr[1] = 'K';
+  ptr[2] = 'G';
+  write_u16(ptr + 3, size);
+  send_attached(ptr, size, nullptr, 0);
 }
 
 void ProxyImpl::send_attached_data(uint32_t id, void const* ptr, size_t size,
diff --git a/test/.gitignore b/test/.gitignore
index aa33417..8bf01f1 100644
--- a/test/.gitignore
+++ b/test/.gitignore
@@ -9,3 +9,4 @@
 /test-xdg
 /test-observers
 /test-htmlattrtext
+/test-package
diff --git a/test/Makefile.am b/test/Makefile.am
index b5d2cee..5264665 100644
--- a/test/Makefile.am
+++ b/test/Makefile.am
@@ -3,7 +3,7 @@ MAINTAINERCLEANFILES = Makefile.in
 AM_CXXFLAGS = @DEFINES@
 
 TESTS = test-url test-http test-args test-xdg test-paths test-strings \
-        test-observers test-htmlattrtext
+        test-observers test-htmlattrtext test-package
 
 check_PROGRAMS = $(TESTS)
 
@@ -29,3 +29,6 @@ test_observers_SOURCES = test-observers.cc
 
 test_htmlattrtext_SOURCES = test-htmlattrtext.cc
 test_htmlattrtext_LDADD = $(top_builddir)/src/libattrstr.a
+
+test_package_SOURCES = test-package.cc
+test_package_LDADD = $(top_builddir)/src/libtp.a
diff --git a/test/test-package.cc b/test/test-package.cc
new file mode 100644
index 0000000..ce15c88
--- /dev/null
+++ b/test/test-package.cc
@@ -0,0 +1,104 @@
+// -*- mode: c++; c-basic-offset: 2; -*-
+
+#include "common.hh"
+#include "test.hh"
+
+#include "package.hh"
+
+namespace {
+
+void setup(Package* pkg) {
+  pkg->id = 42;
+  pkg->timestamp.tv_sec = 123;
+  pkg->timestamp.tv_nsec = 999999999;
+  pkg->source_port = 0;
+  pkg->source_host = "source";
+  pkg->target_port = 65535;
+  pkg->target_host = "target";
+}
+
+bool pkg_eq(Package const& p1, Package const& p2) {
+  return p1.id == p2.id
+    && p1.timestamp.tv_sec == p2.timestamp.tv_sec
+    && p1.timestamp.tv_nsec == p2.timestamp.tv_nsec
+    && p1.source_port == p2.source_port
+    && p1.source_host.compare(p2.source_host) == 0
+    && p1.target_port == p2.target_port
+    && p1.target_host.compare(p2.target_host) == 0;
+}
+
+bool test_sanity() {
+  Package pkg1;
+  Package pkg2;
+
+  setup(&pkg1);
+
+  uint8_t tmp[4096];
+  auto size = write_package(pkg1, tmp, sizeof(tmp));
+  ASSERT_TRUE(size < sizeof(tmp));
+
+  ASSERT_EQ(size, read_package(&pkg2, tmp, sizeof(tmp)));
+
+  ASSERT_TRUE(pkg_eq(pkg1, pkg2));
+
+  return true;
+}
+
+bool test_overflow() {
+  Package pkg1;
+  setup(&pkg1);
+
+  uint8_t tmp[4096];
+  auto size = write_package(pkg1, tmp, sizeof(tmp));
+  for (size_t i = 0; i < size; ++i) {
+    memset(tmp + size, 0, sizeof(tmp) - size);
+    auto s2 = write_package(pkg1, tmp + size, i);
+    ASSERT_EQ(size, s2);
+    ASSERT_EQ(0, tmp[size + i]);
+  }
+  auto s2 = write_package(pkg1, tmp + size, size);
+  ASSERT_EQ(size, s2);
+  ASSERT_EQ(0, memcmp(tmp, tmp + size, size));
+  return true;
+}
+
+bool test_null() {
+  Package pkg1;
+  setup(&pkg1);
+
+  uint8_t tmp[4096];
+  auto size = write_package(pkg1, tmp, sizeof(tmp));
+  auto s2 = write_package(pkg1, NULL, 0);
+  ASSERT_EQ(size, s2);
+  s2 = write_package(pkg1, NULL, sizeof(tmp));
+  ASSERT_EQ(size, s2);
+
+  return true;
+}
+
+bool test_empty() {
+  Package pkg1;
+  Package pkg2;
+  setup(&pkg1);
+  pkg1.source_host.clear();
+  pkg1.target_host.clear();
+
+  uint8_t tmp[4096];
+  auto size = write_package(pkg1, tmp, sizeof(tmp));
+  ASSERT_EQ(size, read_package(&pkg2, tmp, sizeof(tmp)));
+
+  ASSERT_TRUE(pkg_eq(pkg1, pkg2));
+
+  return true;
+}
+
+}  // namespace
+
+int main(void) {
+  BEFORE;
+  RUN(test_sanity());
+  RUN(test_overflow());
+  RUN(test_null());
+  RUN(test_empty());
+  AFTER;
+}