1 files changed, 27 insertions, 58 deletions

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,