Use first IP fragment as base for reassembly

Fixes #225 for good

3 files changed, 20 insertions, 16 deletions

diff --git a/include/tins/ip_reassembler.h b/include/tins/ip_reassembler.h
index bdb994a..3fd0524 100644
--- a/include/tins/ip_reassembler.h
+++ b/include/tins/ip_reassembler.h
@@ -86,7 +86,6 @@ private:
     size_t total_size_;
     IP first_fragment_;
     bool received_end_;
-    uint8_t transport_proto_;
 };
 } // namespace Internals
 
diff --git a/src/ip_reassembler.cpp b/src/ip_reassembler.cpp
index cdbc039..744b4c3 100644
--- a/src/ip_reassembler.cpp
+++ b/src/ip_reassembler.cpp
@@ -38,19 +38,13 @@ namespace Tins {
 namespace Internals {
 
 IPv4Stream::IPv4Stream() 
-: received_size_(), total_size_(), received_end_(false), transport_proto_(0xff) {
+: received_size_(), total_size_(), received_end_(false) {
 
 }
 
 void IPv4Stream::add_fragment(IP* ip) {
-    if (fragments_.empty()) {
-        // Release the inner PDU, store this first fragment and restore the inner PDU
-        PDU* inner_pdu = ip->release_inner_pdu();
-        first_fragment_ = *ip;
-        ip->inner_pdu(inner_pdu);
-    }
+    const uint16_t offset = extract_offset(ip);
     fragments_type::iterator it = fragments_.begin();
-    uint16_t offset = extract_offset(ip);
     while (it != fragments_.end() && offset > it->offset()) {
         ++it;
     }
@@ -66,28 +60,37 @@ void IPv4Stream::add_fragment(IP* ip) {
         received_end_ = true;
     }
     if (offset == 0) {
-        transport_proto_ = ip->protocol();
+        // Release the inner PDU, store this first fragment and restore the inner PDU
+        PDU* inner_pdu = ip->release_inner_pdu();
+        first_fragment_ = *ip;
+        ip->inner_pdu(inner_pdu);
     }
 }
 
 bool IPv4Stream::is_complete() const {
-    return received_end_ && received_size_ == total_size_;
+    // If we haven't received the last chunk of we haven't received all the data,
+    // then we're not complete
+    if (!received_end_ || received_size_ != total_size_) {
+        return false;
+    }
+    // Make sure the first fragment has offset 0
+    return fragments_.begin()->offset() == 0;
 }
 
 PDU* IPv4Stream::allocate_pdu() const {
     PDU::serialization_type buffer;
     buffer.reserve(total_size_);
     // Check if we actually have all the data we need. Otherwise return nullptr;
-    uint16_t expected = 0;
+    size_t expected = 0;
     for (fragments_type::const_iterator it = fragments_.begin(); it != fragments_.end(); ++it) {
         if (expected != it->offset()) {
             return 0;
         }
-        expected = static_cast<uint16_t>(it->offset() + it->payload().size());
+        expected = it->offset() + it->payload().size();
         buffer.insert(buffer.end(), it->payload().begin(), it->payload().end());
     }
     return Internals::pdu_from_flag(
-        static_cast<Constants::IP::e>(transport_proto_),
+        static_cast<Constants::IP::e>(first_fragment_.protocol()),
         buffer.empty() ? 0 :& buffer[0],
         static_cast<uint32_t>(buffer.size())
     );
diff --git a/tests/src/ip_reassembler_test.cpp b/tests/src/ip_reassembler_test.cpp
index 97a6bcb..fb049b4 100644
--- a/tests/src/ip_reassembler_test.cpp
+++ b/tests/src/ip_reassembler_test.cpp
@@ -51,7 +51,9 @@ void IPv4ReassemblerTest::test_packets(const vector<pair<const uint8_t*, size_t>
     IPv4Reassembler reassembler;
     for(size_t i = 0; i < vt.size(); ++i) {
         EthernetII eth(vt[i].first, (uint32_t)vt[i].second);
-        if (i != 0) {
+        // Set the TTL for the first fragment to 32 so we can make sure the right "base"
+        // packet was used to build the reassembled one
+        if (eth.rfind_pdu<IP>().fragment_offset() == 0) {
             eth.rfind_pdu<IP>().ttl(32);
         }
         IPv4Reassembler::PacketStatus status = reassembler.process(eth);
@@ -60,7 +62,7 @@ void IPv4ReassemblerTest::test_packets(const vector<pair<const uint8_t*, size_t>
             ASSERT_EQ(static_cast<size_t>(vt.size() - 1), i);
             ASSERT_TRUE(eth.find_pdu<UDP>() != NULL);
             const IP& ip = eth.rfind_pdu<IP>();
-            EXPECT_EQ(64, ip.ttl());
+            EXPECT_EQ(32, ip.ttl());
 
             RawPDU* raw = eth.find_pdu<RawPDU>();
             ASSERT_TRUE(raw != NULL);