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(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(transport_proto_), + static_cast(first_fragment_.protocol()), buffer.empty() ? 0 :& buffer[0], static_cast(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 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().fragment_offset() == 0) { eth.rfind_pdu().ttl(32); } IPv4Reassembler::PacketStatus status = reassembler.process(eth); @@ -60,7 +62,7 @@ void IPv4ReassemblerTest::test_packets(const vector ASSERT_EQ(static_cast(vt.size() - 1), i); ASSERT_TRUE(eth.find_pdu() != NULL); const IP& ip = eth.rfind_pdu(); - EXPECT_EQ(64, ip.ttl()); + EXPECT_EQ(32, ip.ttl()); RawPDU* raw = eth.find_pdu(); ASSERT_TRUE(raw != NULL);