Lame but usable DTLS defragmentation

2 files changed, 15 insertions, 7 deletions

diff --git a/src/tls/tls_handshake_io.cpp b/src/tls/tls_handshake_io.cpp
index abc2f49f7..08b9f58bb 100644
--- a/src/tls/tls_handshake_io.cpp
+++ b/src/tls/tls_handshake_io.cpp
@@ -220,12 +220,20 @@ void Datagram_Handshake_IO::Handshake_Reassembly::add_fragment(
       }
    else
       {
-      throw Internal_Error("Defragmentation not implemented");
-
-      m_fragments[fragment_offset] =
-         std::deque<byte>(fragment, fragment+fragment_length);
-
-      auto range = m_fragments.equal_range(fragment_offset);
+      /*
+      * FIXME. This is a pretty lame way to do defragmentation, huge
+      * overhead with a tree node per byte.
+      */
+      for(size_t i = 0; i != fragment_length; ++i)
+         m_fragments[fragment_offset+i] = fragment[i];
+
+      if(m_fragments.size() == m_msg_length)
+         {
+         m_message.resize(m_msg_length);
+         for(size_t i = 0; i != m_msg_length; ++i)
+            m_message[i] = m_fragments[i];
+         m_fragments.clear();
+         }
       }
    }
 
diff --git a/src/tls/tls_handshake_io.h b/src/tls/tls_handshake_io.h
index 7628126d6..f5917f004 100644
--- a/src/tls/tls_handshake_io.h
+++ b/src/tls/tls_handshake_io.h
@@ -135,7 +135,7 @@ class Datagram_Handshake_IO : public Handshake_IO
             size_t m_msg_length = 0;
             u16bit m_epoch = 0;
 
-            std::map<size_t, std::deque<byte>> m_fragments;
+            std::map<size_t, byte> m_fragments;
             std::vector<byte> m_message;
          };