Fix next header handling on IPv6

This was broken by design. Each header held the next header's
type, which made iterating through them very tricky (you'd have
to look at the previous header's next header value to know the
current one while doing so).

Now each header contains its own value and the next header
value is computed during serialization and then reverted back
after that.

include/tins/ipv6.h

@@ -346,6 +346,7 @@ private:
 
     ipv6_header header_;
     headers_type ext_headers_;
+    uint8_t next_header_;
 };
 }
 

src/ipv6.cpp

@@ -43,7 +43,7 @@
 #include "memory_helpers.h"
 #include "detail/pdu_helpers.h"
 
-using std::copy;
+using std::vector;
 
 using Tins::Memory::InputMemoryStream;
 using Tins::Memory::OutputMemoryStream;
@@ -70,7 +70,7 @@ PDU::metadata IPv6::extract_metadata(const uint8_t *buffer, uint32_t total_sz) {
 }
 
 IPv6::IPv6(address_type ip_dst, address_type ip_src, PDU* /*child*/)
-: header_() {
+: header_(), next_header_() {
     version(6);
     dst_addr(ip_dst);
     src_addr(ip_src);
@@ -95,8 +95,9 @@ IPv6::IPv6(const uint8_t* buffer, uint32_t total_sz) {
             if (!stream.can_read(payload_size)) {
                 throw malformed_packet();
             }
-            // minus one, from the size field
-            add_ext_header(ext_header(ext_type, payload_size, stream.pointer()));
+            // Add a header using the current header type (e.g. what we saw as the next
+            // header type in the previous)
+            add_ext_header(ext_header(current_header, payload_size, stream.pointer()));
             if (actual_payload_length == 0u && current_header == HOP_BY_HOP) {
                 // could be a jumbogram, look for Jumbo Payload Option
                 InputMemoryStream options(stream.pointer(), payload_size);
@@ -153,6 +154,7 @@ IPv6::IPv6(const uint8_t* buffer, uint32_t total_sz) {
             break;
         }
     }
+    next_header_ = current_header;
 }
 
 bool IPv6::is_extension_header(uint8_t header_id) {
@@ -191,7 +193,7 @@ void IPv6::payload_length(uint16_t new_payload_length) {
 }
 
 void IPv6::next_header(uint8_t new_next_header) {
-    header_.next_header = new_next_header;
+    next_header_ = header_.next_header = new_next_header;
 }
 
 void IPv6::hop_limit(uint8_t new_hop_limit) {
@@ -243,6 +245,20 @@ bool IPv6::matches_response(const uint8_t* ptr, uint32_t total_sz) const {
 
 void IPv6::write_serialization(uint8_t* buffer, uint32_t total_sz) {
     OutputMemoryStream stream(buffer, total_sz);
+    vector<uint8_t> header_types;
+    // Iterate the headers and store their current values. At the same time, update header X
+    // so it has the option type of header X + 1
+    for (size_t i = 0; i < ext_headers_.size(); ++i) {
+        const uint8_t option = ext_headers_[i].option();
+        header_types.push_back(option);
+        if (i > 0) {
+            ext_headers_[i - 1].option(option);
+        }
+    }
+    // If we have at least one, then update our IPv6 header's next header type
+    if (!header_types.empty()) {
+        header_.next_header = header_types[0];
+    }
     if (inner_pdu()) {
         uint8_t new_flag = Internals::pdu_flag_to_ip_type(inner_pdu()->pdu_type());
         if (new_flag == 0xff && Internals::pdu_type_registered<IPv6>(inner_pdu()->pdu_type())) {
@@ -250,9 +266,14 @@ void IPv6::write_serialization(uint8_t* buffer, uint32_t total_sz) {
                 Internals::pdu_type_to_id<IPv6>(inner_pdu()->pdu_type())
             );
         }
+        // If we managed to find the next flag, then set it. Otherwise, fall back to the 
+        // original (or user set) next header
         if (new_flag != 0xff) {
             set_last_next_header(new_flag);
         }
+        else {
+            set_last_next_header(next_header_);
+        }
     }
     else {
         set_last_next_header(0);
@@ -262,6 +283,10 @@ void IPv6::write_serialization(uint8_t* buffer, uint32_t total_sz) {
     for (headers_type::const_iterator it = ext_headers_.begin(); it != ext_headers_.end(); ++it) {
         write_header(*it, stream);
     }
+    // Restore our original header types
+    for (size_t i = 0; i < ext_headers_.size(); ++i) {
+        ext_headers_[i].option(header_types[i]);
+    }
 }
 
 #ifndef BSD
@@ -284,16 +309,14 @@ void IPv6::add_ext_header(const ext_header& header) {
 }
 
 const IPv6::ext_header* IPv6::search_header(ExtensionHeader id) const {
-    uint8_t current_header = header_.next_header;
     headers_type::const_iterator it = ext_headers_.begin();
-    while (it != ext_headers_.end() && current_header != id) {
-        current_header = it->option();
+    while (it != ext_headers_.end()) {
+        if (it->option() == id) {
+            return &*it;
+        }
         ++it;
     }
-    if (it == ext_headers_.end()) {
-        return 0;
-    }
-    return &*it;
+    return 0;
 }
 
 void IPv6::set_last_next_header(uint8_t value) {

tests/src/ipv6_test.cpp

@@ -24,7 +24,7 @@ class IPv6Test : public testing::Test {
 public:
     static const uint8_t expected_packet1[], expected_packet2[], 
                          hop_by_hop_options[], broken1[],
-                         fcs_suffix[];
+                         fcs_suffix[], routing_header[];
     
     void test_equals(IPv6& ip1, IPv6& ip2);
 };
@@ -70,6 +70,20 @@ const uint8_t IPv6Test::fcs_suffix[] = {
     0x57, 0xb7
 };
 
+const uint8_t IPv6Test::routing_header[] = {
+    134, 147, 35, 211, 55, 142, 34, 26, 149, 214, 122, 35, 134, 221, 
+    96, 15, 187, 116, 0, 136, 43, 63, 252, 0, 0, 66, 0, 0, 0, 1, 0, 0
+    , 0, 0, 0, 0, 0, 2, 252, 0, 0, 2, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0
+    , 1, 41, 6, 4, 2, 2, 0, 0, 0, 252, 0, 0, 2, 0, 0, 0, 6, 0, 0, 0, 
+    0, 0, 0, 0, 1, 252, 0, 0, 2, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 1, 
+    252, 0, 0, 2, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 1, 96, 15, 187, 
+    116, 0, 40, 6, 64, 252, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
+    1, 252, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 1, 31, 144, 
+    169, 160, 186, 49, 30, 141, 2, 27, 99, 141, 160, 18, 112, 248, 
+    138, 245, 0, 0, 2, 4, 7, 148, 4, 2, 8, 10, 128, 29, 165, 34, 128,
+    29, 165, 34, 1, 3, 3, 7
+};
+
 void IPv6Test::test_equals(IPv6& ip1, IPv6& ip2) {
     EXPECT_EQ(ip1.version(), ip2.version());
     EXPECT_EQ(ip1.traffic_class(), ip2.traffic_class());
@@ -178,6 +192,10 @@ TEST_F(IPv6Test, Serialize) {
 
 TEST_F(IPv6Test, Broken1) {
     EthernetII pkt(broken1, sizeof(broken1));
+    for (auto c : pkt.serialize()) {
+        printf("%2d, ", (int)c);
+    }
+    printf("\n");
     EXPECT_EQ(
         PDU::serialization_type(broken1, broken1 + sizeof(broken1)),
         pkt.serialize()
@@ -317,3 +335,32 @@ TEST_F(IPv6Test, MDLv1Request) {
     );
 }
 
+TEST_F(IPv6Test, OptionIteration) {
+    EthernetII pkt(routing_header, sizeof(routing_header));
+    IPv6& ipv6 = pkt.rfind_pdu<IPv6>();
+    const IPv6::headers_type& headers = ipv6.headers();
+
+    ASSERT_EQ(1, headers.size());
+    const IPv6::ext_header& header = headers[0];
+    EXPECT_EQ(IPv6::ROUTING, header.option());
+}
+
+TEST_F(IPv6Test, OptionAddition) {
+    EthernetII pkt(routing_header, sizeof(routing_header));
+    IPv6& ipv6 = pkt.rfind_pdu<IPv6>();
+    // Add a dummy header
+    ipv6.add_ext_header(IPv6::ext_header(IPv6::AUTHENTICATION));
+
+    const IPv6::headers_type& headers = ipv6.headers();
+
+    ASSERT_EQ(2, headers.size());
+    EXPECT_EQ(IPv6::ROUTING, headers[0].option());
+    EXPECT_EQ(IPv6::AUTHENTICATION, headers[1].option());
+
+    ipv6.serialize();
+    EXPECT_EQ(IPv6::ROUTING, headers[0].option());
+    EXPECT_EQ(IPv6::AUTHENTICATION, headers[1].option());
+
+    EXPECT_TRUE(ipv6.search_header(IPv6::ROUTING) != 0);
+    EXPECT_TRUE(ipv6.search_header(IPv6::AUTHENTICATION) != 0);
+}