Avoid heap allocations in IPAddress

net/base/ip_address.cc

 // Copyright (c) 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "net/base/ip_address.h"
 
 #include <limits.h>
 
 #include "base/strings/string_piece.h"
 #include "base/strings/string_split.h"
 #include "base/strings/stringprintf.h"
 #include "net/base/parse_number.h"
 #include "url/gurl.h"
 #include "url/url_canon_ip.h"
 
 namespace {
 
 // The prefix for IPv6 mapped IPv4 addresses.
 // https://tools.ietf.org/html/rfc4291#section-2.5.5.2
 const uint8_t kIPv4MappedPrefix[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF, 0xFF};
 
 // Note that this function assumes:
 // * |ip_address| is at least |prefix_length_in_bits| (bits) long;
 // * |ip_prefix| is at least |prefix_length_in_bits| (bits) long.
-bool IPAddressPrefixCheck(const std::vector<uint8_t>& ip_address,
+bool IPAddressPrefixCheck(const net::IPAddress::IPAddressBytes& ip_address,
                           const uint8_t* ip_prefix,
                           size_t prefix_length_in_bits) {
   // Compare all the bytes that fall entirely within the prefix.
   size_t num_entire_bytes_in_prefix = prefix_length_in_bits / 8;
   for (size_t i = 0; i < num_entire_bytes_in_prefix; ++i) {
     if (ip_address[i] != ip_prefix[i])
       return false;
   }
 
   // In case the prefix was not a multiple of 8, there will be 1 byte
   // which is only partially masked.
   size_t remaining_bits = prefix_length_in_bits % 8;
   if (remaining_bits != 0) {
     uint8_t mask = 0xFF << (8 - remaining_bits);
     size_t i = num_entire_bytes_in_prefix;
     if ((ip_address[i] & mask) != (ip_prefix[i] & mask))
       return false;
   }
   return true;
 }
 
 // Returns true if |ip_address| matches any of the reserved IPv4 ranges. This
 // method operates on a blacklist of reserved IPv4 ranges. Some ranges are
 // consolidated.
 // Sources for info:
 // www.iana.org/assignments/ipv4-address-space/ipv4-address-space.xhtml
 // www.iana.org/assignments/iana-ipv4-special-registry/
 // iana-ipv4-special-registry.xhtml
-bool IsReservedIPv4(const std::vector<uint8_t>& ip_address) {
+bool IsReservedIPv4(const net::IPAddress::IPAddressBytes& ip_address) {
   // Different IP versions have different range reservations.
   DCHECK_EQ(net::IPAddress::kIPv4AddressSize, ip_address.size());
   struct {
     const uint8_t address[4];
     size_t prefix_length_in_bits;
   } static const kReservedIPv4Ranges[] = {
       {{0, 0, 0, 0}, 8},     {{10, 0, 0, 0}, 8},      {{100, 64, 0, 0}, 10},
       {{127, 0, 0, 0}, 8},   {{169, 254, 0, 0}, 16},  {{172, 16, 0, 0}, 12},
       {{192, 0, 2, 0}, 24},  {{192, 88, 99, 0}, 24},  {{192, 168, 0, 0}, 16},
       {{198, 18, 0, 0}, 15}, {{198, 51, 100, 0}, 24}, {{203, 0, 113, 0}, 24},
       {{224, 0, 0, 0}, 3}};
 
   for (const auto& range : kReservedIPv4Ranges) {
     if (IPAddressPrefixCheck(ip_address, range.address,
                              range.prefix_length_in_bits)) {
       return true;
     }
   }
 
   return false;
 }
 
 // Returns true if |ip_address| matches any of the reserved IPv6 ranges. This
 // method operates on a whitelist of non-reserved IPv6 ranges. All IPv6
 // addresses outside these ranges are reserved.
 // Sources for info:
 // www.iana.org/assignments/ipv6-address-space/ipv6-address-space.xhtml
-bool IsReservedIPv6(const std::vector<uint8_t>& ip_address) {
+bool IsReservedIPv6(const net::IPAddress::IPAddressBytes& ip_address) {
   // Different IP versions have different range reservations.
   DCHECK_EQ(net::IPAddress::kIPv6AddressSize, ip_address.size());
   struct {
     const uint8_t address_prefix[2];
     size_t prefix_length_in_bits;
   } static const kPublicIPv6Ranges[] = {
       // 2000::/3  -- Global Unicast
       {{0x20, 0}, 3},
       // ff00::/8  -- Multicast
       {{0xff, 0}, 8},
   };
 
   for (const auto& range : kPublicIPv6Ranges) {
     if (IPAddressPrefixCheck(ip_address, range.address_prefix,
                              range.prefix_length_in_bits)) {
       return false;
     }
   }
 
   return true;
 }
 
 bool ParseIPLiteralToBytes(const base::StringPiece& ip_literal,
-                           std::vector<uint8_t>* bytes) {
+                           net::IPAddress::IPAddressBytes* bytes) {
   // |ip_literal| could be either an IPv4 or an IPv6 literal. If it contains
   // a colon however, it must be an IPv6 address.
   if (ip_literal.find(':') != base::StringPiece::npos) {
     // GURL expects IPv6 hostnames to be surrounded with brackets.
     std::string host_brackets = "[";
     ip_literal.AppendToString(&host_brackets);
     host_brackets.push_back(']');
     url::Component host_comp(0, host_brackets.size());
 
     // Try parsing the hostname as an IPv6 literal.
     bytes->resize(16);  // 128 bits.
     return url::IPv6AddressToNumber(host_brackets.data(), host_comp,
                                     bytes->data());
   }
 
   // Otherwise the string is an IPv4 address.
   bytes->resize(4);  // 32 bits.
   url::Component host_comp(0, ip_literal.size());
   int num_components;
   url::CanonHostInfo::Family family = url::IPv4AddressToNumber(
       ip_literal.data(), host_comp, bytes->data(), &num_components);
   return family == url::CanonHostInfo::IPV4;
 }
 
 }  // namespace
 
 namespace net {
 
+IPAddress::IPAddressBytes::IPAddressBytes() : size_(0) {}
+
+net::IPAddress::IPAddressBytes::~IPAddressBytes() {}
+net::IPAddress::IPAddressBytes::IPAddressBytes(
+    net::IPAddress::IPAddressBytes const& other) = default;
+
+bool operator<(IPAddress::IPAddressBytes lhs, IPAddress::IPAddressBytes rhs) {

[eroman, 2017/05/12 23:25:04]

std::array<> already defines these operators (as lexicographic compares, which
is what we want), so you can just forward to them. As in:

  return lhs < rhs;
  return lhs > rhs;
  return lhs == rhs;

[Ryan Hamilton, 2017/05/13 13:20:46]

Ah! So I implemented this initially, and while it compiles, it's subtly wrong.
(There's even a test which fails). The std::array operators look at all the
elements in the array, but IPAddressByte operators only look at the elements in
[0, size_). 

[eroman, 2017/05/15 22:15:23]

Glad that there is test coverage to catch this!

+  if (lhs.size_ < rhs.size_)
+    return true;
+  if (lhs.size_ > rhs.size_)
+    return false;
+  for (size_t i = 0; i < lhs.size_; ++i) {
+    if (lhs.bytes_[i] < rhs.bytes_[i])
+      return true;
+  }
+  return false;
+}
+
+bool operator>(IPAddress::IPAddressBytes lhs, IPAddress::IPAddressBytes rhs) {
+  if (lhs.size_ > rhs.size_)
+    return true;
+  if (lhs.size_ < rhs.size_)
+    return false;
+  for (size_t i = 0; i < lhs.size_; ++i) {
+    if (lhs.bytes_[i] > rhs.bytes_[i])
+      return true;
+  }
+  return false;
+}
+
+bool operator==(IPAddress::IPAddressBytes lhs, IPAddress::IPAddressBytes rhs) {
+  if (lhs.size_ != rhs.size_)
+    return false;
+  for (size_t i = 0; i < lhs.size_; ++i) {
+    if (lhs.bytes_[i] != rhs.bytes_[i])
+      return false;
+  }
+  return true;
+}
+
+bool operator!=(IPAddress::IPAddressBytes lhs, IPAddress::IPAddressBytes rhs) {
+  return !(lhs == rhs);
+}
+
 IPAddress::IPAddress() {}
 
-IPAddress::IPAddress(const std::vector<uint8_t>& address)
-    : ip_address_(address) {}
+IPAddress::IPAddress(const std::vector<uint8_t>& address) {

[eroman, 2017/05/12 23:25:05]

I am a bit worried that we have codepaths that may pass variable sized
addresses, that will now result in memory errors.

Am I being overly paranoid? Should we throw in a CHECK() at the entry points
which take variable sized data? (i.e. here and the other ctor).  We can leave
things like push_back() as DCHECK().

[Ryan Hamilton, 2017/05/13 13:20:47]

Fair enough! I added a CHECK() in the IPAddressBytes constructor which both
these call into now.

+  for (uint8_t val : address)
+    ip_address_.push_back(val);

[eroman, 2017/05/12 23:25:05]

may want to do a single memcpy(). Or perhaps code generator is smart enough to
optimize the loop.

[Ryan Hamilton, 2017/05/13 13:20:46]

Done.

+}
 
 IPAddress::IPAddress(const IPAddress& other) = default;
 
-IPAddress::IPAddress(const uint8_t* address, size_t address_len)
-    : ip_address_(address, address + address_len) {}
+IPAddress::IPAddress(const uint8_t* address, size_t address_len) {
+  for (size_t i = 0; i < address_len; ++i)
+    ip_address_.push_back(address[i]);

[eroman, 2017/05/12 23:25:04]

Same comments as above.

[Ryan Hamilton, 2017/05/13 13:20:47]

Done.

+}
 
 IPAddress::IPAddress(uint8_t b0, uint8_t b1, uint8_t b2, uint8_t b3) {
-  ip_address_.reserve(4);
   ip_address_.push_back(b0);
   ip_address_.push_back(b1);
   ip_address_.push_back(b2);
   ip_address_.push_back(b3);
 }
 
 IPAddress::IPAddress(uint8_t b0,
                      uint8_t b1,
                      uint8_t b2,
                      uint8_t b3,
                      uint8_t b4,
                      uint8_t b5,
                      uint8_t b6,
                      uint8_t b7,
                      uint8_t b8,
                      uint8_t b9,
                      uint8_t b10,
                      uint8_t b11,
                      uint8_t b12,
                      uint8_t b13,
                      uint8_t b14,
                      uint8_t b15) {
-  const uint8_t address[] = {b0, b1, b2,  b3,  b4,  b5,  b6,  b7,
-                             b8, b9, b10, b11, b12, b13, b14, b15};
-  ip_address_ = std::vector<uint8_t>(std::begin(address), std::end(address));
+  ip_address_.push_back(b0);
+  ip_address_.push_back(b1);
+  ip_address_.push_back(b2);
+  ip_address_.push_back(b3);
+  ip_address_.push_back(b4);
+  ip_address_.push_back(b5);
+  ip_address_.push_back(b6);
+  ip_address_.push_back(b7);
+  ip_address_.push_back(b8);
+  ip_address_.push_back(b9);
+  ip_address_.push_back(b10);
+  ip_address_.push_back(b11);
+  ip_address_.push_back(b12);
+  ip_address_.push_back(b13);
+  ip_address_.push_back(b14);
+  ip_address_.push_back(b15);
 }
 
 IPAddress::~IPAddress() {}
 
 bool IPAddress::IsIPv4() const {
   return ip_address_.size() == kIPv4AddressSize;
 }
 
 bool IPAddress::IsIPv6() const {
   return ip_address_.size() == kIPv6AddressSize;
 }
 
 bool IPAddress::IsValid() const {
   return IsIPv4() || IsIPv6();
 }
 
 bool IPAddress::IsReserved() const {
   if (IsIPv4()) {
     return IsReservedIPv4(ip_address_);
   } else if (IsIPv6()) {
     return IsReservedIPv6(ip_address_);
   }
   return false;
 }
 
 bool IPAddress::IsZero() const {
   for (auto x : ip_address_) {

[eroman, 2017/05/12 23:25:04]

Fancy. So this works because of the begin() and end() methods?

[Ryan Hamilton, 2017/05/13 13:20:46]

Right?! Every once in a while, C++ doesn't suck. There are also a TON of places
where IPAddress.bytes() is used in a similar context and works magically.

     if (x != 0)
       return false;
   }
 
   return !empty();
 }
 
 bool IPAddress::IsIPv4MappedIPv6() const {
   return IsIPv6() && IPAddressStartsWith(*this, kIPv4MappedPrefix);
 }
 
 bool IPAddress::AssignFromIPLiteral(const base::StringPiece& ip_literal) {
-  std::vector<uint8_t> number;
+  IPAddressBytes number;
 
   if (!ParseIPLiteralToBytes(ip_literal, &number))
     return false;
 
   std::swap(number, ip_address_);

[eroman, 2017/05/12 23:25:05]

Should probably just change this to assignment now:
  ip_address_ = number;

std::swap() is worse now, since we need to copy ip_address_ back into number_
too.

On a separate note, we should probably just change the contract so ip_address_
is cleared on failure, to avoid needing this temporary at all.

[Ryan Hamilton, 2017/05/13 13:20:46]

Done.
TODO added.

   return true;
 }
 
+std::vector<uint8_t> IPAddress::BytesAsVector() const {
+  return std::vector<uint8_t>(ip_address_.begin(), ip_address_.end());
+}
+
 // static
 IPAddress IPAddress::IPv4Localhost() {
   static const uint8_t kLocalhostIPv4[] = {127, 0, 0, 1};
   return IPAddress(kLocalhostIPv4);
 }
 
 // static
 IPAddress IPAddress::IPv6Localhost() {
   static const uint8_t kLocalhostIPv6[] = {0, 0, 0, 0, 0, 0, 0, 0,
                                            0, 0, 0, 0, 0, 0, 0, 1};
@@ skipping 161 matching lines @@
       if (diff & (1 << (CHAR_BIT - 1)))
         return i * CHAR_BIT + j;
       diff <<= 1;
     }
     NOTREACHED();
   }
   return a1.size() * CHAR_BIT;
 }
 
 unsigned MaskPrefixLength(const IPAddress& mask) {
   std::vector<uint8_t> all_ones(mask.size(), 0xFF);

[eroman, 2017/05/12 23:25:05]

If IPAddress were mutable (i.e. merge IPAddressBytes into it), patterns like
this will be a bit simpler. (Can create a stack-based IPAddress and manipulate
it directly, without the temporary vector<>).

[Ryan Hamilton, 2017/05/13 13:20:47]

How 'bout we do the flattening in a follow up since I suspect that's going to
change quite a number of places.

[eroman, 2017/05/15 22:15:23]

sgtm

   return CommonPrefixLength(mask, IPAddress(all_ones));
 }
 
 }  // namespace net