Add OutputMemoryStream and port most classes to use it

include/tins/dhcpv6.h

@@ -38,7 +38,11 @@
 #include "ipv6_address.h"
 #include "pdu_option.h"
 
-namespace Tins {    
+namespace Tins {
+namespace Memory  {
+class OutputMemoryStream;
+} // Memory
+
 /**
  * \class DHCPv6
  * \brief Represents a DHCPv6 PDU.
@@ -869,7 +873,7 @@ class DHCPv6 : public PDU {
     }
 private:
     void write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *);
-    uint8_t* write_option(const option &option, uint8_t* buffer) const;
+    void write_option(const option &option, Memory::OutputMemoryStream& stream) const;
     options_type::const_iterator search_option_iterator(OptionTypes type) const;
     options_type::iterator search_option_iterator(OptionTypes type);
 

include/tins/eapol.h

@@ -38,7 +38,10 @@
 
 
 namespace Tins {
-
+namespace Memory {
+class OutputMemoryStream;
+} // Memory
+
     /** \cond 
      * Forward declaration. Avoid header inclusion.
      */
@@ -163,7 +166,7 @@ namespace Tins {
          * \param buffer The pointer in which to save the serialization.
          * \param total_sz The total size of the buffer.
          */
-        virtual void write_body(uint8_t *buffer, uint32_t total_sz) = 0;
+        virtual void write_body(Memory::OutputMemoryStream& stream) = 0;
     private:
         /** 
          * \brief Serialices this EAPOL PDU.
@@ -353,7 +356,7 @@ namespace Tins {
             uint8_t key_sign[16];
         } TINS_END_PACK;
 
-        void write_body(uint8_t *buffer, uint32_t total_sz);
+        void write_body(Memory::OutputMemoryStream& stream);
 
 
         key_type _key;
@@ -731,7 +734,7 @@ namespace Tins {
         #endif
         } TINS_END_PACK;
 
-        void write_body(uint8_t *buffer, uint32_t total_sz);
+        void write_body(Memory::OutputMemoryStream& stream);
 
 
         rsnhdr _header;

include/tins/exceptions.h

@@ -49,7 +49,7 @@ class exception_base : public std::runtime_error {
 /**
  * \brief Exception thrown when an option is not found.
  */
-class option_not_found : exception_base {
+class option_not_found : public exception_base {
 public:
     // try to avoid allocations by doing this.
     const char* what() const throw() {
@@ -67,6 +67,16 @@ class malformed_packet : public exception_base {
     }
 };
 
+/**
+ * \brief Exception thrown when serializing a packet fails.
+ */
+class serialization_error : public exception_base {
+public:
+    const char* what() const throw() {
+        return "Serialization error";
+    }
+};
+
 /**
  * \brief Exception thrown when a PDU is not found when using PDU::rfind_pdu.
  */

include/tins/icmpv6.h

@@ -46,6 +46,7 @@
 namespace Tins {
 namespace Memory {
 class InputMemoryStream;
+class OutputMemoryStream;
 } // memory
 
 /**
@@ -1379,7 +1380,7 @@ class ICMPv6 : public PDU {
     void internal_add_option(const option &option);
     void write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent);
     bool has_options() const;
-    uint8_t *write_option(const option &opt, uint8_t *buffer);
+    void write_option(const option &opt, Memory::OutputMemoryStream& stream);
     void parse_options(Memory::InputMemoryStream& stream);
     void add_addr_list(uint8_t type, const addr_list_type &value);
     addr_list_type search_addr_list(OptionTypes type) const;

include/tins/ip.h

@@ -40,6 +40,9 @@
 #include "cxxstd.h"
 
 namespace Tins {
+namespace Memory {
+class OutputMemoryStream;
+} // Memory
 
     /**
      * \class IP
@@ -725,7 +728,7 @@ namespace Tins {
         void internal_add_option(const option &option);
         void init_ip_fields();
         void write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent);
-        uint8_t* write_option(const option &opt, uint8_t* buffer);
+        void write_option(const option &opt, Memory::OutputMemoryStream& stream);
         void add_route_option(option_identifier id, const generic_route_option_type &data);
         generic_route_option_type search_route_option(option_identifier id) const;
         void checksum(uint16_t new_check);

include/tins/ipv6.h

@@ -40,6 +40,10 @@
 #include "ipv6_address.h"
 
 namespace Tins {
+namespace Memory {
+class OutputMemoryStream;
+} // Memory
+
 class PacketSender;
 
 /**
@@ -299,7 +303,7 @@ class IPv6 : public PDU {
 private:
     void write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent);
     void set_last_next_header(uint8_t value);
-    static uint8_t *write_header(const ext_header &header, uint8_t *buffer);
+    static void write_header(const ext_header &header, Memory::OutputMemoryStream& stream);
     static bool is_extension_header(uint8_t header_id);
 
     TINS_BEGIN_PACK

include/tins/llc.h

@@ -383,6 +383,7 @@ namespace Tins {
         #endif
 
         typedef std::vector<uint8_t> field_type;
+        typedef std::list<field_type> field_list;
 
         void write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent);
 
@@ -395,7 +396,7 @@ namespace Tins {
         } control_field;
         Format _type;
         uint8_t information_field_length;
-        std::list<field_type> information_fields;
+        field_list information_fields;
     };
 
 }

include/tins/memory_helpers.h

@@ -3,6 +3,7 @@
 
 #include <stdint.h>
 #include <cstring>
+#include <algorithm>
 #include "exceptions.h"
 #include "ip_address.h"
 #include "ipv6_address.h"
@@ -20,6 +21,10 @@ void read_value(const uint8_t* buffer, T& value) {
     std::memcpy(&value, buffer, sizeof(value));
 }
 
+inline void write_data(uint8_t* buffer, const uint8_t* ptr, uint32_t size) {
+    std::memcpy(buffer, ptr, size);
+}
+
 template <typename T>
 void write_value(uint8_t* buffer, const T& value) {
     std::memcpy(buffer, &value, sizeof(value));
@@ -105,6 +110,73 @@ class InputMemoryStream {
     uint32_t size_;
 };
 
+class OutputMemoryStream {
+public:
+    OutputMemoryStream(uint8_t* buffer, uint32_t total_sz)
+    : buffer_(buffer), size_(total_sz) {
+    }
+
+    void skip(uint32_t size) {
+        buffer_ += size;
+        size_ -= size;
+    }
+
+    template <typename T>
+    void write(const T& value) {
+        if (TINS_UNLIKELY(size_ < sizeof(value))) {
+            throw serialization_error();
+        }
+        write_value(buffer_, value);
+        skip(sizeof(value));
+    }
+
+    template <typename ForwardIterator>
+    void write(ForwardIterator start, ForwardIterator end) {
+        const uint32_t length = std::distance(start, end); 
+        if (TINS_UNLIKELY(size_ < length)) {
+            throw serialization_error();
+        }
+        std::copy(start, end, buffer_);
+        skip(length);
+    }
+
+    void write(const uint8_t* ptr, uint32_t length) {
+        write(ptr, ptr + length);
+    }
+
+    void write(const IPv4Address& address) {
+        write(static_cast<uint32_t>(address));
+    }
+
+    void write(const IPv6Address& address) {
+        write(address.begin(), address.end());
+    }
+
+    template <size_t n>
+    void write(const HWAddress<n>& address) {
+        write(address.begin(), address.end());
+    }
+
+    void fill(uint32_t size, uint8_t value) {
+        if (TINS_UNLIKELY(size_ < size)) {
+            throw serialization_error();
+        }
+        std::fill(buffer_, buffer_ + size, value);
+        skip(size);
+    }
+
+    uint8_t* pointer() {
+        return buffer_;
+    }
+
+    uint32_t size() const {
+        return size_;
+    }
+private:
+    uint8_t* buffer_;
+    uint32_t size_;
+};
+
 } // Memory
 } // Tins
 

include/tins/tcp.h

@@ -44,6 +44,10 @@
 #include "cxxstd.h"
 
 namespace Tins {
+namespace Memory {
+class OutputMemoryStream;
+} // Memory
+
     /**
      * \class TCP
      * \brief Represents a TCP PDU.
@@ -583,7 +587,7 @@ namespace Tins {
         options_type::const_iterator search_option_iterator(OptionTypes type) const;
         options_type::iterator search_option_iterator(OptionTypes type);
 
-        uint8_t *write_option(const option &opt, uint8_t *buffer);
+        void write_option(const option &opt, Memory::OutputMemoryStream& stream);
 
         tcphdr _tcp;
         uint16_t _options_size, _total_options_size;

src/arp.cpp

@@ -28,7 +28,6 @@
  */
 
 #include <cstring>
-#include <cassert>
 #include <algorithm>
 #include "arp.h"
 #include "ip.h"
@@ -42,6 +41,7 @@
 using std::runtime_error;
 
 using Tins::Memory::InputMemoryStream;
+using Tins::Memory::OutputMemoryStream;
 
 namespace Tins {
 
@@ -109,10 +109,8 @@ uint32_t ARP::header_size() const {
 }
 
 void ARP::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *) {
-    #ifdef TINS_DEBUG
-    assert(total_sz >= sizeof(arphdr));
-    #endif
-    memcpy(buffer, &_arp, sizeof(arphdr));
+    OutputMemoryStream stream(buffer, total_sz);
+    stream.write(_arp);
 }
 
 bool ARP::matches_response(const uint8_t *ptr, uint32_t total_sz) const {

src/bootp.cpp

@@ -29,12 +29,12 @@
 
 #include <stdexcept>
 #include <cstring>
-#include <cassert>
 #include "bootp.h"
 #include "exceptions.h"
 #include "memory_helpers.h"
 
 using Tins::Memory::InputMemoryStream;
+using Tins::Memory::OutputMemoryStream;
 
 namespace Tins{
 
@@ -116,11 +116,9 @@ void BootP::vend(const vend_type &new_vend) {
 }
 
 void BootP::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent) {
-    #ifdef TINS_DEBUG
-    assert(total_sz >= sizeof(bootphdr) + _vend.size());
-    #endif
-    std::memcpy(buffer, &_bootp, sizeof(bootphdr));
-    std::copy(_vend.begin(), _vend.end(), buffer + sizeof(bootphdr));
+    OutputMemoryStream stream(buffer, total_sz);
+    stream.write(_bootp);
+    stream.write(_vend.begin(), _vend.end());
 }
 
 bool BootP::matches_response(const uint8_t *ptr, uint32_t total_sz) const {

src/dhcp.cpp

@@ -28,7 +28,6 @@
  */
 
 #include <stdexcept>
-#include <cassert>
 #include <cstring>
 #include "endianness.h"
 #include "dhcp.h"
@@ -43,6 +42,7 @@ using std::runtime_error;
 using std::find_if;
 
 using Tins::Memory::InputMemoryStream;
+using Tins::Memory::OutputMemoryStream;
 
 namespace Tins {
 
@@ -57,8 +57,8 @@ DHCP::DHCP()
 DHCP::DHCP(const uint8_t *buffer, uint32_t total_sz) 
 : BootP(buffer, total_sz, 0), _size(sizeof(uint32_t))
 {
-    const uint32_t bootp_size = BootP::header_size() - vend().size();
-    InputMemoryStream stream(buffer + bootp_size, total_sz - bootp_size);
+    InputMemoryStream stream(buffer, total_sz);
+    stream.skip(BootP::header_size() - vend().size());
     const uint32_t magic_number = stream.read<uint32_t>();
     if (magic_number != Endian::host_to_be<uint32_t>(0x63825363))
         throw malformed_packet();
@@ -238,22 +238,19 @@ uint32_t DHCP::header_size() const {
 }
 
 void DHCP::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *parent) {
-    #ifdef TINS_DEBUG
-    assert(total_sz >= header_size());
-    #endif
-    if(_size) {
-        vend_type &result(BootP::vend());
+    if (_size) {
+        vend_type &result = BootP::vend();
         result.resize(_size);
-        uint8_t *ptr = &result[0] + sizeof(uint32_t);
+        // Build a stream over the vend vector
+        OutputMemoryStream stream(&result[0], result.size());
         // Magic cookie
-        *((uint32_t*)&result[0]) = Endian::host_to_be<uint32_t>(0x63825363);
-        for(options_type::const_iterator it = _options.begin(); it != _options.end(); ++it) {
-            *(ptr++) = it->option();
-            *(ptr++) = static_cast<uint8_t>(it->length_field());
-            std::copy(it->data_ptr(), it->data_ptr() + it->data_size(), ptr);
-            ptr += it->data_size();
+        stream.write(Endian::host_to_be<uint32_t>(0x63825363));
+        for (options_type::const_iterator it = _options.begin(); it != _options.end(); ++it) {
+            stream.write(it->option());
+            stream.write<uint8_t>(it->length_field());
+            stream.write(it->data_ptr(), it->data_size());
         }
     }
     BootP::write_serialization(buffer, total_sz, parent);
 }
-}
+} // Tins