From 1f4be63d0885e29d55fbde7f0e62a6e72dd0f273 Mon Sep 17 00:00:00 2001
From: Matias Fontanini <matias.fontanini@gmail.com>
Date: Thu, 17 Mar 2016 21:49:06 -0700
Subject: [PATCH] Properly handle MLDv1 on ICMP

---
 include/tins/icmpv6.h | 12 +++++++++++
 src/icmpv6.cpp        | 46 ++++++++++++++++++++++++++++---------------
 tests/src/ipv6.cpp    | 15 ++++++++++++++
 3 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/include/tins/icmpv6.h b/include/tins/icmpv6.h
index 7a34523..9cb3dad 100644
--- a/include/tins/icmpv6.h
+++ b/include/tins/icmpv6.h
@@ -1115,6 +1115,17 @@ public:
     ICMPv6* clone() const {
         return new ICMPv6(*this);
     }
+
+    /** 
+     * \brief Indicates whether to use MLDv2
+     *
+     * If this is set to true, then MLDv2 will be used rather than MLDv1 when
+     * serializing Multicast Listener Discovery messages. By default,
+     * MLDv2 will be used.
+     *
+     * \param value The value to set
+     */
+    void use_mldv2(bool value);
     
     // ****************************************************************
     //                          Option setters
@@ -1590,6 +1601,7 @@ private:
     multicast_listener_query_message_fields mlqm_;
     sources_list sources_;
     ICMPExtensionsStructure extensions_;
+    bool use_mldv2_;
 };
 
 } // Tins
diff --git a/src/icmpv6.cpp b/src/icmpv6.cpp
index 585bcd6..0dc4d76 100644
--- a/src/icmpv6.cpp
+++ b/src/icmpv6.cpp
@@ -47,13 +47,13 @@ using Tins::Memory::OutputMemoryStream;
 namespace Tins {
 
 ICMPv6::ICMPv6(Types tp)
-: options_size_(), reach_time_(0), retrans_timer_(0), mlqm_() {
+: options_size_(), reach_time_(0), retrans_timer_(0), mlqm_(), use_mldv2_(true) {
     memset(&header_, 0, sizeof(header_));
     type(tp);
 }
 
 ICMPv6::ICMPv6(const uint8_t* buffer, uint32_t total_sz) 
-: options_size_(), reach_time_(0), retrans_timer_(0), mlqm_() {
+: options_size_(), reach_time_(0), retrans_timer_(0), mlqm_(), use_mldv2_(true) {
     InputMemoryStream stream(buffer, total_sz);
     stream.read(header_);
     if (has_target_addr()) {
@@ -77,12 +77,16 @@ ICMPv6::ICMPv6(const uint8_t* buffer, uint32_t total_sz)
     }
     else if (type() == MGM_QUERY) {
         stream.read(multicast_address_);
-        stream.read(mlqm_);
-        int sources_count = stream.read_be<uint16_t>();
-        while (sources_count--) {
-            ipaddress_type address;
-            stream.read(address);
-            sources_.push_back(address);
+        // MLDv1 ends here
+        use_mldv2_ = stream;
+        if (stream) {
+            stream.read(mlqm_);
+            int sources_count = stream.read_be<uint16_t>();
+            while (sources_count--) {
+                ipaddress_type address;
+                stream.read(address);
+                sources_.push_back(address);
+            }
         }
     }
     // Retrieve options
@@ -232,8 +236,11 @@ uint32_t ICMPv6::header_size() const {
         }
     }
     else if (type() == MGM_QUERY) {
-        extra += ipaddress_type::address_size + sizeof(mlqm_) + sizeof(uint16_t) +
-                 ipaddress_type::address_size * sources_.size();
+        extra += ipaddress_type::address_size;
+        if (use_mldv2_) {
+            extra += sizeof(mlqm_) + sizeof(uint16_t) + 
+                     ipaddress_type::address_size * sources_.size();
+        }
     }
     return sizeof(header_) + options_size_ + extra + 
         (has_target_addr() ? ipaddress_type::address_size : 0) +
@@ -315,12 +322,15 @@ void ICMPv6::write_serialization(uint8_t* buffer, uint32_t total_sz, const PDU*
     }
     else if (type() == MGM_QUERY) {
         stream.write(multicast_address_);
-        stream.write(mlqm_);
-        stream.write_be<uint16_t>(sources_.size());
-        typedef sources_list::const_iterator iterator;
-        for (iterator iter = sources_.begin(); iter != sources_.end(); ++iter) {
-            stream.write(*iter);
-        } 
+        // Only write this if we're using MLDv2
+        if (use_mldv2_) {
+            stream.write(mlqm_);
+            stream.write_be<uint16_t>(sources_.size());
+            typedef sources_list::const_iterator iterator;
+            for (iterator iter = sources_.begin(); iter != sources_.end(); ++iter) {
+                stream.write(*iter);
+            } 
+        }
     }
     for (options_type::const_iterator it = options_.begin(); it != options_.end(); ++it) {
         write_option(*it, stream);
@@ -414,6 +424,10 @@ void ICMPv6::write_option(const option& opt, OutputMemoryStream& stream) {
     stream.write(opt.data_ptr(), opt.data_size());
 }
 
+void ICMPv6::use_mldv2(bool value) {
+    use_mldv2_ = value;
+}
+
 const ICMPv6::option* ICMPv6::search_option(OptionTypes type) const {
     // Search for the iterator. If we found something, return it, otherwise return nullptr.
     options_type::const_iterator iter = search_option_iterator(type);
diff --git a/tests/src/ipv6.cpp b/tests/src/ipv6.cpp
index f4782e2..c02a449 100644
--- a/tests/src/ipv6.cpp
+++ b/tests/src/ipv6.cpp
@@ -283,3 +283,18 @@ TEST_F(IPv6Test, SerializePacketHavingICMPExtensionsWithoutLengthAndShortPayload
     ASSERT_EQ(1UL, serialized.rfind_pdu<ICMPv6>().extensions().extensions().size());
     EXPECT_EQ(ext_payload, serialized.rfind_pdu<ICMPv6>().extensions().extensions().begin()->payload());
 }
+
+TEST_F(IPv6Test, MDLv1Request) {
+    const uint8_t mldv1[] = {
+        51, 51, 0, 0, 0, 1, 100, 112, 2, 226, 169, 250, 134, 221, 96, 0, 0, 0, 0, 32, 0, 1, 254,
+        128, 0, 0, 0, 0, 0, 0, 102, 112, 2, 255, 254, 226, 169, 250, 255, 2, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 1, 58, 0, 5, 2, 0, 0, 0, 0, 130, 0, 70, 203, 39, 16, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+    };
+    EthernetII pkt(mldv1, sizeof(mldv1));
+    EXPECT_EQ(
+        PDU::serialization_type(mldv1, mldv1 + sizeof(mldv1)),
+        pkt.serialize()
+    );
+}
+