diff --git a/include/tins/ip.h b/include/tins/ip.h index 236882c..6da8390 100644 --- a/include/tins/ip.h +++ b/include/tins/ip.h @@ -523,7 +523,6 @@ public: * \param opt The option to be added. */ void add_option(option &&opt) { - internal_add_option(opt); options_.push_back(std::move(opt)); } @@ -538,7 +537,6 @@ public: template void add_option(Args&&... args) { options_.emplace_back(std::forward(args)...); - internal_add_option(options_.back()); } #endif @@ -756,7 +754,8 @@ private: void tot_len(uint16_t new_tot_len); void prepare_for_serialize(); - void internal_add_option(const option& option); + uint32_t calculate_options_size() const; + uint32_t pad_options_size(uint32_t size) const; void init_ip_fields(); void write_serialization(uint8_t* buffer, uint32_t total_sz); void write_option(const option& opt, Memory::OutputMemoryStream& stream); @@ -765,11 +764,9 @@ private: void checksum(uint16_t new_check); options_type::const_iterator search_option_iterator(option_identifier id) const; options_type::iterator search_option_iterator(option_identifier id); - void update_padded_options_size(); - ip_header header_; - uint16_t options_size_, padded_options_size_; options_type options_; + ip_header header_; }; } // Tins diff --git a/include/tins/memory_helpers.h b/include/tins/memory_helpers.h index ec63eac..a07ad43 100644 --- a/include/tins/memory_helpers.h +++ b/include/tins/memory_helpers.h @@ -68,8 +68,12 @@ void write_value(uint8_t* buffer, const T& value) { class InputMemoryStream { public: - InputMemoryStream(const uint8_t* buffer, size_t total_sz); - InputMemoryStream(const std::vector& data); + InputMemoryStream(const uint8_t* buffer, size_t total_sz) + : buffer_(buffer), size_(total_sz) { + } + + InputMemoryStream(const std::vector& data) : buffer_(&data[0]), size_(data.size()) { + } template T read() { @@ -97,17 +101,46 @@ public: skip(sizeof(value)); } - void skip(size_t size); - bool can_read(size_t byte_count) const; + void skip(size_t size) { + if (TINS_UNLIKELY(size > size_)) { + throw malformed_packet(); + } + buffer_ += size; + size_ -= size; + } + + bool can_read(size_t byte_count) const { + return TINS_LIKELY(size_ >= byte_count); + } + + void read(void* output_buffer, size_t output_buffer_size) { + if (!can_read(output_buffer_size)) { + throw malformed_packet(); + } + read_data(buffer_, (uint8_t*)output_buffer, output_buffer_size); + skip(output_buffer_size); + } + + const uint8_t* pointer() const { + return buffer_; + } + + size_t size() const { + return size_; + } + + void size(size_t new_size) { + size_ = new_size; + } + + operator bool() const { + return size_ > 0; + } + void read(std::vector& value, size_t count); void read(HWAddress<6>& address); void read(IPv4Address& address); void read(IPv6Address& address); - void read(void* output_buffer, size_t output_buffer_size); - const uint8_t* pointer() const; - size_t size() const; - void size(size_t new_size); - operator bool() const; private: const uint8_t* buffer_; size_t size_; @@ -115,8 +148,13 @@ private: class OutputMemoryStream { public: - OutputMemoryStream(uint8_t* buffer, size_t total_sz); - OutputMemoryStream(std::vector& buffer); + OutputMemoryStream(uint8_t* buffer, size_t total_sz) + : buffer_(buffer), size_(total_sz) { + } + + OutputMemoryStream(std::vector& buffer) + : buffer_(&buffer[0]), size_(buffer.size()) { + } template void write(const T& value) { @@ -151,14 +189,37 @@ public: skip(length); } - void skip(size_t size); - void write(const uint8_t* ptr, size_t length); + void skip(size_t size) { + if (TINS_UNLIKELY(size > size_)) { + throw malformed_packet(); + } + buffer_ += size; + size_ -= size; + } + + void write(const uint8_t* ptr, size_t length) { + write(ptr, ptr + length); + } + + void fill(size_t size, uint8_t value) { + if (TINS_UNLIKELY(size_ < size)) { + throw serialization_error(); + } + std::memset(buffer_, value, size); + skip(size); + } + + uint8_t* pointer() { + return buffer_; + } + + size_t size() const { + return size_; + } + void write(const HWAddress<6>& address); void write(const IPv4Address& address); void write(const IPv6Address& address); - void fill(size_t size, uint8_t value); - uint8_t* pointer(); - size_t size() const; private: uint8_t* buffer_; size_t size_; diff --git a/include/tins/pdu_option.h b/include/tins/pdu_option.h index c614828..4724e56 100644 --- a/include/tins/pdu_option.h +++ b/include/tins/pdu_option.h @@ -226,7 +226,7 @@ public: * \brief Move constructor. * \param rhs The PDUOption to be moved. */ - PDUOption(PDUOption&& rhs) { + PDUOption(PDUOption&& rhs) TINS_NOEXCEPT { real_size_ = 0; *this = std::move(rhs); } @@ -235,7 +235,7 @@ public: * \brief Move assignment operator. * \param rhs The PDUOption to be moved. */ - PDUOption& operator=(PDUOption&& rhs) { + PDUOption& operator=(PDUOption&& rhs) TINS_NOEXCEPT { option_ = rhs.option_; size_ = rhs.size_; if (real_size_ > small_buffer_size) { @@ -250,7 +250,7 @@ public: else { std::memcpy(payload_.small_buffer, rhs.data_ptr(), rhs.data_size()); } - return* this; + return *this; } #endif // TINS_IS_CXX11 diff --git a/include/tins/tcp.h b/include/tins/tcp.h index 8794ca7..9ba2019 100644 --- a/include/tins/tcp.h +++ b/include/tins/tcp.h @@ -475,7 +475,6 @@ public: * \param option The option to be added. */ void add_option(option &&opt) { - internal_add_option(opt); options_.push_back(std::move(opt)); } @@ -490,7 +489,6 @@ public: template void add_option(Args&&... args) { options_.emplace_back(std::forward(args)...); - internal_add_option(options_.back()); } #endif @@ -606,18 +604,17 @@ private: return opt->to(); } - void internal_add_option(const option& option); void write_serialization(uint8_t* buffer, uint32_t total_sz); void checksum(uint16_t new_check); - void update_options_size(); + uint32_t calculate_options_size() const; + uint32_t pad_options_size(uint32_t size) const; options_type::const_iterator search_option_iterator(OptionTypes type) const; options_type::iterator search_option_iterator(OptionTypes type); void write_option(const option& opt, Memory::OutputMemoryStream& stream); - tcp_header header_; - uint16_t options_size_, total_options_size_; options_type options_; + tcp_header header_; }; } // Tins diff --git a/src/ip.cpp b/src/ip.cpp index ec4bd94..2380126 100644 --- a/src/ip.cpp +++ b/src/ip.cpp @@ -73,14 +73,13 @@ IP::IP(address_type ip_dst, address_type ip_src) { this->src_addr(ip_src); } -IP::IP(const uint8_t* buffer, uint32_t total_sz) -: options_size_(0) { +IP::IP(const uint8_t* buffer, uint32_t total_sz) { InputMemoryStream stream(buffer, total_sz); stream.read(header_); // Make sure we have enough size for options and not less than we should - if (head_len() * sizeof(uint32_t) > total_sz || - head_len() * sizeof(uint32_t) < sizeof(header_)) { + if (TINS_UNLIKELY(head_len() * sizeof(uint32_t) > total_sz || + head_len() * sizeof(uint32_t) < sizeof(header_))) { throw malformed_packet(); } const uint8_t* options_end = buffer + head_len() * sizeof(uint32_t); @@ -108,7 +107,6 @@ IP::IP(const uint8_t* buffer, uint32_t total_sz) else { options_.push_back(option(opt_type)); } - options_size_ += option_size; } else if (opt_type == END) { // If the end option found, we're done @@ -120,10 +118,8 @@ IP::IP(const uint8_t* buffer, uint32_t total_sz) } else { options_.push_back(option(opt_type)); - options_size_++; } } - update_padded_options_size(); if (stream) { // Don't avoid consuming more than we should if tot_len is 0, // since this is the case when using TCP segmentation offload @@ -171,8 +167,6 @@ void IP::init_ip_fields() { header_.version = 4; ttl(DEFAULT_TTL); id(1); - options_size_ = 0; - padded_options_size_ = 0; } bool IP::is_fragmented() const { @@ -321,18 +315,20 @@ uint16_t IP::stream_identifier() const { } void IP::add_option(const option& opt) { - internal_add_option(opt); options_.push_back(opt); } -void IP::update_padded_options_size() { - uint8_t padding = options_size_ % 4; - padded_options_size_ = padding ? (options_size_ - padding + 4) : options_size_; +uint32_t IP::calculate_options_size() const { + uint32_t options_size = 0; + for (options_type::const_iterator iter = options_.begin(); iter != options_.end(); ++iter) { + options_size += 1 + iter->data_size(); + } + return options_size; } -void IP::internal_add_option(const option& opt) { - options_size_ += static_cast(1 + opt.data_size()); - update_padded_options_size(); +uint32_t IP::pad_options_size(uint32_t size) const { + uint8_t padding = size % 4; + return padding ? (size - padding + 4) : size; } bool IP::remove_option(option_identifier id) { @@ -340,9 +336,7 @@ bool IP::remove_option(option_identifier id) { if (iter == options_.end()) { return false; } - options_size_ -= static_cast(1 + iter->data_size()); options_.erase(iter); - update_padded_options_size(); return true; } @@ -375,7 +369,7 @@ void IP::write_option(const option& opt, OutputMemoryStream& stream) { // Virtual method overriding uint32_t IP::header_size() const { - return sizeof(header_) + padded_options_size_; + return sizeof(header_) + pad_options_size(calculate_options_size()); } PacketSender::SocketType pdu_type_to_sender_type(PDU::PDUType type) { @@ -459,8 +453,10 @@ void IP::write_serialization(uint8_t* buffer, uint32_t total_sz) { for (options_type::const_iterator it = options_.begin(); it != options_.end(); ++it) { write_option(*it, stream); } + const uint32_t options_size = calculate_options_size(); + const uint32_t padded_options_size = pad_options_size(options_size); // Add option padding - stream.fill(padded_options_size_ - options_size_, 0); + stream.fill(padded_options_size - options_size, 0); uint32_t check = Utils::do_checksum(buffer, stream.pointer()); while (check >> 16) { diff --git a/src/memory_helpers.cpp b/src/memory_helpers.cpp index 9d6d242..a547e1f 100644 --- a/src/memory_helpers.cpp +++ b/src/memory_helpers.cpp @@ -39,26 +39,6 @@ namespace Memory { // InputMemoryStream -InputMemoryStream::InputMemoryStream(const uint8_t* buffer, size_t total_sz) -: buffer_(buffer), size_(total_sz) { -} - -InputMemoryStream::InputMemoryStream(const vector& data) -: buffer_(&data[0]), size_(data.size()) { -} - -void InputMemoryStream::skip(size_t size) { - if (TINS_UNLIKELY(size > size_)) { - throw malformed_packet(); - } - buffer_ += size; - size_ -= size; -} - -bool InputMemoryStream::can_read(size_t byte_count) const { - return TINS_LIKELY(size_ >= byte_count); -} - void InputMemoryStream::read(vector& value, size_t count) { if (!can_read(count)) { throw malformed_packet(); @@ -87,52 +67,8 @@ void InputMemoryStream::read(IPv6Address& address) { skip(IPv6Address::address_size); } -void InputMemoryStream::read(void* output_buffer, size_t output_buffer_size) { - if (!can_read(output_buffer_size)) { - throw malformed_packet(); - } - read_data(buffer_, (uint8_t*)output_buffer, output_buffer_size); - skip(output_buffer_size); -} - -const uint8_t* InputMemoryStream::pointer() const { - return buffer_; -} - -size_t InputMemoryStream::size() const { - return size_; -} - -void InputMemoryStream::size(size_t new_size) { - size_ = new_size; -} - -InputMemoryStream::operator bool() const { - return size_ > 0; -} - // OutputMemoryStream -OutputMemoryStream::OutputMemoryStream(uint8_t* buffer, size_t total_sz) -: buffer_(buffer), size_(total_sz) { -} - -OutputMemoryStream::OutputMemoryStream(vector& buffer) -: buffer_(&buffer[0]), size_(buffer.size()) { -} - -void OutputMemoryStream::skip(size_t size) { - if (TINS_UNLIKELY(size > size_)) { - throw malformed_packet(); - } - buffer_ += size; - size_ -= size; -} - -void OutputMemoryStream::write(const uint8_t* ptr, size_t length) { - write(ptr, ptr + length); -} - void OutputMemoryStream::write(const HWAddress<6>& address) { write(address.begin(), address.end()); } @@ -145,21 +81,5 @@ void OutputMemoryStream::write(const IPv6Address& address) { write(address.begin(), address.end()); } -void OutputMemoryStream::fill(size_t size, uint8_t value) { - if (TINS_UNLIKELY(size_ < size)) { - throw serialization_error(); - } - std::memset(buffer_, value, size); - skip(size); -} - -uint8_t* OutputMemoryStream::pointer() { - return buffer_; -} - -size_t OutputMemoryStream::size() const { - return size_; -} - } // Memory } // Tins diff --git a/src/tcp.cpp b/src/tcp.cpp index bbafe4c..2906302 100644 --- a/src/tcp.cpp +++ b/src/tcp.cpp @@ -56,15 +56,14 @@ PDU::metadata TCP::extract_metadata(const uint8_t *buffer, uint32_t total_sz) { } TCP::TCP(uint16_t dport, uint16_t sport) -: header_(), options_size_(0), total_options_size_(0) { +: header_() { this->dport(dport); this->sport(sport); data_offset(sizeof(tcp_header) / sizeof(uint32_t)); window(DEFAULT_WINDOW); } -TCP::TCP(const uint8_t* buffer, uint32_t total_sz) -: options_size_(0), total_options_size_(0) { +TCP::TCP(const uint8_t* buffer, uint32_t total_sz) { InputMemoryStream stream(buffer, total_sz); stream.read(header_); // Check that we have at least the amount of bytes we need and not less @@ -74,6 +73,12 @@ TCP::TCP(const uint8_t* buffer, uint32_t total_sz) } const uint8_t* header_end = buffer + (data_offset() * sizeof(uint32_t)); + if (stream.pointer() < header_end) { + // Estimate about 4 bytes per option and reserver that so we avoid doing + // multiple reallocations on the vector + options_.reserve((header_end - stream.pointer()) / sizeof(uint32_t)); + } + while (stream.pointer() < header_end) { const OptionTypes option_type = (OptionTypes)stream.read(); if (option_type <= NOP) { @@ -284,25 +289,26 @@ void TCP::flags(small_uint<12> value) { void TCP::add_option(const option& opt) { options_.push_back(opt); - internal_add_option(opt); } uint32_t TCP::header_size() const { - return sizeof(header_) + total_options_size_; + return sizeof(header_) + pad_options_size(calculate_options_size()); } void TCP::write_serialization(uint8_t* buffer, uint32_t total_sz) { OutputMemoryStream stream(buffer, total_sz); + const uint32_t options_size = calculate_options_size(); + const uint32_t total_options_size = pad_options_size(options_size); // Set checksum to 0, we'll calculate it at the end checksum(0); - header_.doff = (sizeof(tcp_header) + total_options_size_) / sizeof(uint32_t); + header_.doff = (sizeof(tcp_header) + total_options_size) / sizeof(uint32_t); stream.write(header_); for (options_type::const_iterator it = options_.begin(); it != options_.end(); ++it) { write_option(*it, stream); } - if (options_size_ < total_options_size_) { - const uint16_t padding = total_options_size_ - options_size_; + if (options_size < total_options_size) { + const uint16_t padding = total_options_size - options_size; stream.fill(padding, 1); } @@ -366,19 +372,23 @@ void TCP::write_option(const option& opt, OutputMemoryStream& stream) { } } -void TCP::update_options_size() { - uint8_t padding = options_size_ & 3; - total_options_size_ = (padding) ? (options_size_ - padding + 4) : options_size_; +uint32_t TCP::calculate_options_size() const { + uint32_t options_size = 0; + for (options_type::const_iterator iter = options_.begin(); iter != options_.end(); ++iter) { + const option& opt = *iter; + 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(opt.data_size()); + } + } + return options_size; } -void TCP::internal_add_option(const option& opt) { - 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(opt.data_size()); - } - update_options_size(); +uint32_t TCP::pad_options_size(uint32_t size) const { + uint8_t padding = size & 3; + return padding ? (size - padding + 4) : size; } bool TCP::remove_option(OptionTypes type) { @@ -386,14 +396,7 @@ bool TCP::remove_option(OptionTypes 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(iter->data_size()); - } options_.erase(iter); - update_options_size(); return true; }