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 @@ public:
     }
 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 @@ public:
 /**
  * \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 @@ public:
     }
 };
 
+/**
+ * \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 @@ private:
     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 @@ public:
 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 @@ private:
     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;