From 6e026fcb663044c356fa87b3f6440fa69ad565b7 Mon Sep 17 00:00:00 2001
From: Matias Fontanini <matias.fontanini@gmail.com>
Date: Mon, 14 Dec 2015 20:22:50 -0800
Subject: [PATCH] Add missing member functions to ICMP extension classes

---
 include/tins/icmp_extension.h | 35 ++++++++++++++++++++++++++++++++++-
 src/icmp_extension.cpp        | 24 ++++++++++++++++++++++++
 tests/src/icmp_extension.cpp  | 30 ++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/include/tins/icmp_extension.h b/include/tins/icmp_extension.h
index 9731d79..8035bab 100644
--- a/include/tins/icmp_extension.h
+++ b/include/tins/icmp_extension.h
@@ -25,6 +25,11 @@ public:
      */
     typedef std::vector<uint8_t> serialization_type;
 
+    /**
+     * \brief Default constructor
+     */
+    ICMPExtension();
+
     /**
      * \brief Constructs an ICMP extension from a buffer
      *
@@ -33,6 +38,27 @@ public:
      */
     ICMPExtension(const uint8_t* buffer, uint32_t total_sz);
 
+    /**
+     * \brief Setter for the extension class field
+     *
+     * \param value The new extension class field value
+     */
+    void extension_class(uint8_t value);
+
+    /**
+     * \brief Setter for the extension sub-type field
+     *
+     * \param value The new extension sub-type field value
+     */
+    void extension_type(uint8_t value);
+
+    /**
+     * \brief Setter for the payload field
+     *
+     * \param value The new payload field value
+     */
+    void payload(const payload_type& value);
+
     /**
      * \brief Getter for the extension class field
      *
@@ -116,7 +142,14 @@ public:
     ICMPExtensionsStructure(const uint8_t* buffer, uint32_t total_sz);
 
     /**
-     * \brief Setter for the checksum field
+     * \brief Setter for the version field
+     *
+     * \param value The new version field value
+     */
+    void version(small_uint<4> value);
+
+    /**
+     * \brief Setter for the reserved field
      *
      * \param value The new reserved field value
      */
diff --git a/src/icmp_extension.cpp b/src/icmp_extension.cpp
index d81d1aa..8c94e92 100644
--- a/src/icmp_extension.cpp
+++ b/src/icmp_extension.cpp
@@ -12,6 +12,11 @@ const uint32_t ICMPExtension::BASE_HEADER_SIZE = sizeof(uint16_t) + sizeof(uint8
 
 // ICMPExtension class
 
+ICMPExtension::ICMPExtension() 
+: extension_class_(0), extension_type_(0) {
+
+} 
+
 ICMPExtension::ICMPExtension(const uint8_t* buffer, uint32_t total_sz) {
     // Check for the base header (u16 length + u8 clss + u8 type)
     if (total_sz < BASE_HEADER_SIZE) {
@@ -31,6 +36,18 @@ ICMPExtension::ICMPExtension(const uint8_t* buffer, uint32_t total_sz) {
     payload_.assign(buffer, buffer + length);
 }
 
+void ICMPExtension::extension_class(uint8_t value) {
+    extension_class_ = value;
+}
+
+void ICMPExtension::extension_type(uint8_t value) {
+    extension_type_ = value;
+}
+
+void ICMPExtension::payload(const payload_type& value) {
+    payload_ = value;
+}
+
 uint32_t ICMPExtension::size() const {
     return BASE_HEADER_SIZE + payload_.size();
 }
@@ -87,6 +104,13 @@ void ICMPExtensionsStructure::reserved(small_uint<12> value) {
     version_and_reserved_ = Endian::host_to_be(current_value);
 }
 
+void ICMPExtensionsStructure::version(small_uint<4> value) {
+    uint16_t current_value = Endian::be_to_host(version_and_reserved_);
+    current_value &= 0xfff;
+    current_value |= value << 12;
+    version_and_reserved_ = Endian::host_to_be(current_value);   
+}
+
 bool ICMPExtensionsStructure::validate_extensions(const uint8_t* buffer, uint32_t total_sz) {
     if (total_sz < BASE_HEADER_SIZE) {
         return false;
diff --git a/tests/src/icmp_extension.cpp b/tests/src/icmp_extension.cpp
index b7b4e80..2b399e7 100644
--- a/tests/src/icmp_extension.cpp
+++ b/tests/src/icmp_extension.cpp
@@ -27,6 +27,28 @@ TEST_F(ICMPExtensionTest, ConstructorFromBuffer) {
     );
 }
 
+TEST_F(ICMPExtensionTest, ExtensionClass) {
+    ICMPExtension extension;
+    extension.extension_class(126);
+    EXPECT_EQ(126, extension.extension_class());
+}
+
+TEST_F(ICMPExtensionTest, ExtensionType) {
+    ICMPExtension extension;
+    extension.extension_type(126);
+    EXPECT_EQ(126, extension.extension_type());
+}
+
+TEST_F(ICMPExtensionTest, Payload) {
+    ICMPExtension::payload_type payload;
+    payload.push_back(0x92);
+    payload.push_back(0x1a);
+    payload.push_back(0xde);
+    ICMPExtension extension;
+    extension.payload(payload);
+    EXPECT_EQ(payload, extension.payload());
+}
+
 TEST_F(ICMPExtensionTest, ExtensionStructureValidation) {
     const uint8_t input[] = { 32, 0, 197, 95, 0, 8, 1, 1, 24, 150, 1, 1 };
     EXPECT_TRUE(ICMPExtensionsStructure::validate_extensions(input, sizeof(input)));
@@ -63,3 +85,11 @@ TEST_F(ICMPExtensionTest, Reserved) {
     EXPECT_EQ(0xdea, structure.reserved());
     EXPECT_EQ(2, structure.version());
 }
+
+TEST_F(ICMPExtensionTest, Version) {
+    ICMPExtensionsStructure structure;
+    structure.reserved(0xdea);
+    structure.version(0xf);
+    EXPECT_EQ(0xdea, structure.reserved());
+    EXPECT_EQ(0xf, structure.version());
+}