Add remove_option member to IP, TCP, Dot11, ICMPv6, DHCP and DHCPv6

src/dhcp.cpp

@@ -33,11 +33,13 @@
 #include "endianness.h"
 #include "dhcp.h"
 #include "ethernetII.h"
+#include "internals.h"
 #include "exceptions.h"
 
 using std::string;
 using std::list;
 using std::runtime_error;
+using std::find_if;
 
 namespace Tins {
 // Magic cookie: uint32_t. 
@@ -89,12 +91,30 @@ void DHCP::internal_add_option(const option &opt) {
     _size += static_cast<uint32_t>(opt.data_size() + (sizeof(uint8_t) << 1));
 }
 
-const DHCP::option *DHCP::search_option(OptionTypes opt) const {
-    for(options_type::const_iterator it = _options.begin(); it != _options.end(); ++it) {
-        if(it->option() == opt)
-            return &(*it);
+bool DHCP::remove_option(OptionTypes type) {
+    options_type::iterator iter = search_option_iterator(type);
+    if (iter == _options.end()) {
+        return false;
     }
-    return 0;
+    _size -= static_cast<uint32_t>(iter->data_size() + (sizeof(uint8_t) << 1));
+    _options.erase(iter);
+    return true;
+}
+
+const DHCP::option *DHCP::search_option(OptionTypes opt) const {
+    // Search for the iterator. If we found something, return it, otherwise return nullptr.
+    options_type::const_iterator iter = search_option_iterator(opt);
+    return (iter != _options.end()) ? &*iter : 0;
+}
+
+DHCP::options_type::const_iterator DHCP::search_option_iterator(OptionTypes opt) const {
+    Internals::option_type_equality_comparator<option> comparator(opt);
+    return find_if(_options.begin(), _options.end(), comparator);
+}
+
+DHCP::options_type::iterator DHCP::search_option_iterator(OptionTypes opt) {
+    Internals::option_type_equality_comparator<option> comparator(opt);
+    return find_if(_options.begin(), _options.end(), comparator);
 }
 
 void DHCP::type(Flags type) {

src/dhcpv6.cpp

@@ -32,6 +32,8 @@
 #include "dhcpv6.h"
 #include "exceptions.h"
 
+using std::find_if;
+
 namespace Tins {
 DHCPv6::DHCPv6() : options_size() {
     std::fill(header_data, header_data + sizeof(header_data), 0);
@@ -84,12 +86,30 @@ void DHCPv6::add_option(const option &opt) {
     options_size += opt.data_size() + sizeof(uint16_t) * 2;
 }
 
-const DHCPv6::option *DHCPv6::search_option(OptionTypes id) const {
-    for(options_type::const_iterator it = options_.begin(); it != options_.end(); ++it) {
-        if(it->option() == static_cast<uint16_t>(id))
-            return &*it;
+bool DHCPv6::remove_option(OptionTypes type) {
+    options_type::iterator iter = search_option_iterator(type);
+    if (iter == options_.end()) {
+        return false;
     }
-    return 0;
+    options_size -= iter->data_size() + sizeof(uint16_t) * 2;
+    options_.erase(iter);
+    return true;
+}
+
+const DHCPv6::option *DHCPv6::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);
+    return (iter != options_.end()) ? &*iter : 0;
+}
+
+DHCPv6::options_type::const_iterator DHCPv6::search_option_iterator(OptionTypes type) const {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(options_.begin(), options_.end(), comparator);
+}
+
+DHCPv6::options_type::iterator DHCPv6::search_option_iterator(OptionTypes type) {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(options_.begin(), options_.end(), comparator);
 }
 
 uint8_t* DHCPv6::write_option(const option &opt, uint8_t* buffer) const {
@@ -151,8 +171,9 @@ void DHCPv6::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *
         buffer = link_addr.copy(buffer);
         buffer = peer_addr.copy(buffer);
     }
-    for(options_type::const_iterator it = options_.begin(); it != options_.end(); ++it)
+    for(options_type::const_iterator it = options_.begin(); it != options_.end(); ++it) {
         buffer = write_option(*it, buffer);
+    }
 }
 
 

src/dot11/dot11_base.cpp

@@ -106,16 +106,35 @@ void Dot11::internal_add_option(const option &opt) {
     _options_size += static_cast<uint32_t>(opt.data_size() + sizeof(uint8_t) * 2);
 }
 
+bool Dot11::remove_option(OptionTypes type) {
+    options_type::iterator iter = search_option_iterator(type);
+    if (iter == _options.end()) {
+        return false;
+    }
+    _options_size -= static_cast<uint32_t>(iter->data_size() + sizeof(uint8_t) * 2);
+    _options.erase(iter);
+    return true;
+}
+
 void Dot11::add_option(const option &opt) {
     internal_add_option(opt);
     _options.push_back(opt);
 }
 
-const Dot11::option *Dot11::search_option(OptionTypes opt) const {
-    for(std::list<option>::const_iterator it = _options.begin(); it != _options.end(); ++it)
-        if(it->option() == (uint8_t)opt)
-            return &(*it);
-    return 0;
+const Dot11::option *Dot11::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);
+    return (iter != _options.end()) ? &*iter : 0;
+}
+
+Dot11::options_type::const_iterator Dot11::search_option_iterator(OptionTypes type) const {
+    Internals::option_type_equality_comparator<option> comparator(static_cast<uint8_t>(type));
+    return find_if(_options.begin(), _options.end(), comparator);
+}
+
+Dot11::options_type::iterator Dot11::search_option_iterator(OptionTypes type) {
+    Internals::option_type_equality_comparator<option> comparator(static_cast<uint8_t>(type));
+    return find_if(_options.begin(), _options.end(), comparator);
 }
 
 void Dot11::protocol(small_uint<2> new_proto) {

src/icmpv6.cpp

@@ -263,18 +263,36 @@ void ICMPv6::internal_add_option(const option &option) {
     _options_size += static_cast<uint32_t>(option.data_size() + sizeof(uint8_t) * 2);
 }
 
+bool ICMPv6::remove_option(OptionTypes type) {
+    options_type::iterator iter = search_option_iterator(type);
+    if (iter == _options.end()) {
+        return false;
+    }
+    _options_size -= static_cast<uint32_t>(iter->data_size() + sizeof(uint8_t) * 2);
+    _options.erase(iter);
+    return true;
+}
+
 uint8_t *ICMPv6::write_option(const option &opt, uint8_t *buffer) {
     *buffer++ = opt.option();
     *buffer++ = static_cast<uint8_t>((opt.length_field() + sizeof(uint8_t) * 2) / 8);
     return std::copy(opt.data_ptr(), opt.data_ptr() + opt.data_size(), buffer);
 }
 
-const ICMPv6::option *ICMPv6::search_option(OptionTypes id) const {
-    for(options_type::const_iterator it = _options.begin(); it != _options.end(); ++it) {
-        if(it->option() == id)
-            return &*it;
-    }
-    return 0;
+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);
+    return (iter != _options.end()) ? &*iter : 0;
+}
+
+ICMPv6::options_type::const_iterator ICMPv6::search_option_iterator(OptionTypes type) const {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(_options.begin(), _options.end(), comparator);
+}
+
+ICMPv6::options_type::iterator ICMPv6::search_option_iterator(OptionTypes type) {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(_options.begin(), _options.end(), comparator);
 }
 
 // ********************************************************************

src/ip.cpp

@@ -297,18 +297,40 @@ void IP::add_option(const option &opt) {
     _ip_options.push_back(opt);
 }
 
-void IP::internal_add_option(const option &opt) {
-    _options_size += static_cast<uint16_t>(1 + opt.data_size());
+void IP::update_padded_options_size() {
     uint8_t padding = _options_size % 4;
     _padded_options_size = padding ? (_options_size - padding + 4) : _options_size;
 }
 
-const IP::option *IP::search_option(option_identifier id) const {
-    for(options_type::const_iterator it = _ip_options.begin(); it != _ip_options.end(); ++it) {
-        if(it->option() == id)
-            return &(*it);
+void IP::internal_add_option(const option &opt) {
+    _options_size += static_cast<uint16_t>(1 + opt.data_size());
+    update_padded_options_size();
+}
+
+bool IP::remove_option(option_identifier id) {
+    options_type::iterator iter = search_option_iterator(id);
+    if (iter == _ip_options.end()) {
+        return false;
     }
-    return 0;
+    _options_size -= static_cast<uint16_t>(1 + iter->data_size());
+    _ip_options.erase(iter);
+    update_padded_options_size();
+    return true;
+}
+
+const IP::option *IP::search_option(option_identifier id) const {
+    options_type::const_iterator iter = search_option_iterator(id);
+    return (iter != _ip_options.end()) ? &*iter : 0;
+}
+
+IP::options_type::const_iterator IP::search_option_iterator(option_identifier id) const {
+    Internals::option_type_equality_comparator<option> comparator(id);
+    return find_if(_ip_options.begin(), _ip_options.end(), comparator);
+}
+
+IP::options_type::iterator IP::search_option_iterator(option_identifier id) {
+    Internals::option_type_equality_comparator<option> comparator(id);
+    return find_if(_ip_options.begin(), _ip_options.end(), comparator);
 }
 
 uint8_t* IP::write_option(const option &opt, uint8_t* buffer) {

src/tcp.cpp

@@ -28,6 +28,7 @@
  */
 
 #include <cstring>
+#include <algorithm>
 #include <cassert>
 #include "tcp.h"
 #include "ip.h"
@@ -36,6 +37,9 @@
 #include "rawpdu.h"
 #include "utils.h"
 #include "exceptions.h"
+#include "internals.h"
+
+using std::find_if;
 
 namespace Tins {
 
@@ -321,12 +325,20 @@ void TCP::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *par
     }
 }
 
-const TCP::option *TCP::search_option(OptionTypes opt) const {
-    for(options_type::const_iterator it = _options.begin(); it != _options.end(); ++it) {
-        if(it->option() == opt)
-            return &(*it);
-    }
-    return 0;
+const TCP::option *TCP::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);
+    return (iter != _options.end()) ? &*iter : 0;
+}
+
+TCP::options_type::const_iterator TCP::search_option_iterator(OptionTypes type) const {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(_options.begin(), _options.end(), comparator);
+}
+
+TCP::options_type::iterator TCP::search_option_iterator(OptionTypes type) {
+    Internals::option_type_equality_comparator<option> comparator(type);
+    return find_if(_options.begin(), _options.end(), comparator);
 }
 
 /* options */
@@ -347,18 +359,35 @@ uint8_t *TCP::write_option(const option &opt, uint8_t *buffer) {
     }
 }
 
+void TCP::update_options_size() {
+    uint8_t padding = _options_size & 3;
+    _total_options_size = (padding) ? (_options_size - padding + 4) : _options_size;
+}
+
 void TCP::internal_add_option(const option &opt) {
-    uint8_t padding;
-    
     _options_size += sizeof(uint8_t);
     // SACK_OK contains length but not data....
-    if(opt.data_size() || opt.option() == SACK_OK)
-        _options_size += sizeof(uint8_t);
-        
-    _options_size += static_cast<uint16_t>(opt.data_size());
-    
-    padding = _options_size & 3;
-    _total_options_size = (padding) ? _options_size - padding + 4 : _options_size;
+    if(opt.data_size() || opt.option() == SACK_OK) {
+        _options_size += sizeof(uint8_t);    
+        _options_size += static_cast<uint16_t>(opt.data_size());
+    }
+    update_options_size();
+}
+
+bool TCP::remove_option(OptionTypes type) {
+    options_type::iterator iter = search_option_iterator(type);
+    if (iter == _options.end()) {
+        return false;
+    }
+    _options_size -= sizeof(uint8_t);
+    // SACK_OK contains length but not data....
+    if(iter->data_size() || iter->option() == SACK_OK) {
+        _options_size -= sizeof(uint8_t);
+        _options_size -= static_cast<uint16_t>(iter->data_size());
+    }
+    _options.erase(iter);
+    update_options_size();
+    return true;
 }
 
 bool TCP::matches_response(const uint8_t *ptr, uint32_t total_sz) const {