Support nested prefixes in the LcTrie class

source/common/network/lc_trie.h

@@ -70,6 +70,9 @@ class LcTrie {
     // By shifting the value to the left by p bits(and back), the bits between 0 and p-1 are zero'd
     // out. Then to get the n bits, shift the IP back by the address_size minus the number of
     // desired bits.
+    if (n == 0) {
+      return IpType(0);
+    }
     return input << p >> (address_size - n);
   }
 
@@ -92,31 +95,6 @@ class LcTrie {
   typedef uint32_t Ipv4;
   typedef absl::uint128 Ipv6;
 
-  // Helper methods to retrieve the string representation of the address in IpPrefix using
-  // inet_ntop. These strings are used in the nested prefixes exception messages.
-  // TODO(ccaraman): Remove once nested prefixes are supported.
-  static std::string toString(const Ipv4& input) {
-    sockaddr_in addr4;
-    addr4.sin_family = AF_INET;
-    addr4.sin_addr.s_addr = htonl(input);
-    addr4.sin_port = htons(0);
-
-    Address::Ipv4Instance address(&addr4);
-    return address.ip()->addressAsString();
-  }
-
-  static std::string toString(const Ipv6& input) {
-    sockaddr_in6 addr6;
-    addr6.sin6_family = AF_INET6;
-    addr6.sin6_port = htons(0);
-
-    // Ipv6 stores the values in host byte order.
-    absl::uint128 ipv6 = Utility::Ip6htonl(input);
-    memcpy(&addr6.sin6_addr.s6_addr, &ipv6, sizeof(absl::uint128));
-    Address::Ipv6Instance address(addr6);
-    return address.ip()->addressAsString();
-  }
-
   /**
    * Structure to hold a CIDR range and the tag associated with it.
    */
@@ -150,20 +128,41 @@ class LcTrie {
      * @return true if other is a prefix of this.
      */
     bool isPrefix(const IpPrefix& other) {
-      return (length_ == 0 || (length_ <= other.length_ &&
-                               extractBits<IpType, address_size>(0, length_, ip_) ==
-                                   extractBits<IpType, address_size>(0, length_, other.ip_)));
+      return (length_ == 0 || (length_ <= other.length_ && contains(other.ip_)));
+    }
+
+    /**
+     * @param address supplies an IP address to check against this prefix.
+     * @return bool true if this prefix contains the address.
+     */
+    bool contains(const IpType& address) const {
+      return (extractBits<IpType, address_size>(0, length_, ip_) ==
+              extractBits<IpType, address_size>(0, length_, address));
     }
 
     std::string asString() { return fmt::format("{}/{}", toString(ip_), length_); }
 
+    void addNestedPrefix(const IpPrefix& other) {
+      if (nested_prefixes_ == nullptr) {
+        nested_prefixes_ = std::make_shared<std::vector<IpPrefix>>();
+      }
+      nested_prefixes_->push_back(other);
+    }
+
     // The address represented either in Ipv4(uint32_t) or Ipv6(asbl::uint128).
     IpType ip_{0};
     // Length of the cidr range.
     int length_;
     // TODO(ccaraman): Support more than one tag per entry.
     // Tag for this entry.
     std::string tag_;
+    // Other prefixes nested under this one. If an LC trie lookup matches on this
+    // prefix, the lookup will scan the nested prefixes to see if any of them match,
+    // too. This situation is rare, so to save memory in the common case the
+    // nested_prefixes field is a pointer to a vector rather than an inline vector.
+    // TODO(brian-pane) switch to a trie of nested prefixes, to ensure sublinear-time
+    // searching even in situations where there are a lot of nested prefixes.
+    std::shared_ptr<std::vector<IpPrefix>> nested_prefixes_;
   };
 
   /**
@@ -226,16 +225,12 @@ class LcTrie {
       std::sort(tag_data.begin(), tag_data.end());
       ip_prefixes_.push_back(tag_data[0]);
 
-      // Remove duplicate entries and check for nested prefixes.
+      // Set up ip_prefixes_, which should contain the supplied prefixes in sorted order,
+      // but with any nested prefixes encapsulated under their parent prefixes.
       for (size_t i = 1; i < tag_data.size(); ++i) {
-        // TODO(ccaraman): Add support for nested prefixes.
-        if (tag_data[i - 1].isPrefix(tag_data[i])) {
-          throw EnvoyException(fmt::format(
-              "LcTrie does not support nested prefixes. '{0}' is a nested prefix of '{1}'.",
-              tag_data[i].asString(), tag_data[i - 1].asString()));
-        }
-
-        if (tag_data[i - 1] != tag_data[i]) {
+        if (ip_prefixes_[ip_prefixes_.size() - 1].isPrefix(tag_data[i])) {
+          ip_prefixes_[ip_prefixes_.size() - 1].addNestedPrefix(tag_data[i]);
+        } else {
           ip_prefixes_.push_back(tag_data[i]);
         }
       }
@@ -531,21 +526,24 @@ LcTrie::LcTrieInternal<IpType, address_size>::getTags(const IpType& ip_address)
     address = node.address_;
   }
 
-  // /0 will match all IP addresses.
-  if (ip_prefixes_[address].length_ == 0) {
-    return_vector.push_back(ip_prefixes_[address].tag_);
-    // TODO(ccaraman): When nested prefixes are supported, should the /0 case be handled better?
-    return return_vector;
-  }
-
-  // Check the input IP address is within the CIDR range stored in ip_prefixes_[adr] by XOR'ing the
-  // two values and checking that up until the length of the CIDR range the value is 0.
-  IpType bitmask = ip_prefixes_[address].ip_ ^ ip_address;
-  if (extractBits<IpType, address_size>(0, ip_prefixes_[address].length_, bitmask) == 0) {
-    return_vector.push_back(ip_prefixes_[address].tag_);
+  // The prefix table entry ip_prefixes_[address] contains either a single prefix or
+  // a parent prefix with a set of additional prefixes nested under it. In the latter
+  // case, we compare the supplied ip_address against all the prefixes in the entry
+  // and return the union of all the matches' tags.
+  // TODO(ccaraman): determine whether there's a more optimal way to handle "/0" prefixes.
+  std::unordered_set<std::string> unique_tags;
+  const auto& prefix = ip_prefixes_[address];
+  if (prefix.contains(ip_address)) {
+    unique_tags.insert(prefix.tag_);
+    if (prefix.nested_prefixes_ != nullptr) {
+      for (const auto& nested_prefix : *prefix.nested_prefixes_) {
+        if (nested_prefix.contains(ip_address)) {
+          unique_tags.insert(nested_prefix.tag_);
+        }
+      }
+    }
   }
-  // TODO(ccaraman): Search through the nested prefix structure.
-  return return_vector;
+  return std::vector<std::string>(unique_tags.begin(), unique_tags.end());
 }
 
 } // namespace LcTrie

test/common/network/lc_trie_test.cc

@@ -35,15 +35,18 @@ class LcTrieTest : public testing::Test {
   void expectIPAndTags(
       const std::vector<std::pair<std::string, std::vector<std::string>>>& test_output) {
     for (const auto& kv : test_output) {
-      EXPECT_EQ(kv.second, trie_->getTags(Utility::parseInternetAddress(kv.first)));
+      std::vector<std::string> expected(kv.second);
+      std::sort(expected.begin(), expected.end());
+      std::vector<std::string> actual(trie_->getTags(Utility::parseInternetAddress(kv.first)));
+      std::sort(actual.begin(), actual.end());
+      EXPECT_EQ(expected, actual);
     }
   }
 
   std::unique_ptr<LcTrie> trie_;
 };
 
 // TODO(ccaraman): Add a performance and memory benchmark test.
-
 // Use the default constructor values.
 TEST_F(LcTrieTest, IPv4Defaults) {
   std::vector<std::vector<std::string>> cidr_range_strings = {
@@ -252,22 +255,50 @@ TEST_F(LcTrieTest, BothIpvVersions) {
 }
 
 TEST_F(LcTrieTest, NestedPrefixes) {
-  EXPECT_THROW_WITH_MESSAGE(setup({{"1.2.3.130/24", "1.2.3.255/32"}}), EnvoyException,
-                            "LcTrie does not support nested prefixes. '1.2.3.255/32' is a nested "
-                            "prefix of '1.2.3.0/24'.");
-
-  EXPECT_THROW_WITH_MESSAGE(
-      setup({{"0.0.0.0/0", "1.2.3.254/32"}}), EnvoyException,
-      "LcTrie does not support nested prefixes. '1.2.3.254/32' is a nested prefix of '0.0.0.0/0'.");
-
-  EXPECT_THROW_WITH_MESSAGE(setup({{"2406:da00:2000::/40", "2406:da00:2000::1/100"}}),
-                            EnvoyException,
-                            "LcTrie does not support nested prefixes. '2406:da00:2000::/100' is a "
-                            "nested prefix of '2406:da00:2000::/40'.");
-
-  EXPECT_THROW_WITH_MESSAGE(setup({{"::/0", "2406:da00:2000::1/100"}}), EnvoyException,
-                            "LcTrie does not support nested prefixes. '2406:da00:2000::/100' is a "
-                            "nested prefix of '::/0'.");
+  const std::vector<std::vector<std::string>> cidr_range_strings = {
+      {"203.0.113.0/24", "203.0.113.128/25"}, // tag_0
+      {"203.0.113.255/32"},                   // tag_1
+      {"198.51.100.0/24"},                    // tag_2
+      {"2001:db8::/96", "2001:db8::8000/97"}, // tag_3
+      {"2001:db8::ffff/128"},                 // tag_4
+      {"2001:db8:1::/48"},                    // tag_5
+      {"2001:db8:1::/48"}                     // tag_6
+  };
+  setup(cidr_range_strings);
+
+  const std::vector<std::pair<std::string, std::vector<std::string>>> test_case = {
+      {"203.0.113.0", {"tag_0"}},
+      {"203.0.113.192", {"tag_0"}},
+      {"203.0.113.255", {"tag_0", "tag_1"}},
+      {"198.51.100.1", {"tag_2"}},
+      {"2001:db8::ffff", {"tag_3", "tag_4"}},
+      {"2001:db8:1::ffff", {"tag_5", "tag_6"}}};
+  expectIPAndTags(test_case);
+}
+
+TEST_F(LcTrieTest, NestedPrefixesWithCatchAll) {
+  std::vector<std::vector<std::string>> cidr_range_strings = {
+      {"0.0.0.0/0"},                          // tag_0
+      {"203.0.113.0/24"},                     // tag_1
+      {"203.0.113.128/25"},                   // tag_2
+      {"198.51.100.0/24"},                    // tag_3
+      {"::0/0"},                              // tag_4
+      {"2001:db8::/96", "2001:db8::8000/97"}, // tag_5
+      {"2001:db8::ffff/128"},                 // tag_6
+      {"2001:db8:1::/48"}                     // tag_7
+  };
+  setup(cidr_range_strings);
+
+  std::vector<std::pair<std::string, std::vector<std::string>>> test_case = {
+      {"203.0.113.0", {"tag_0", "tag_1"}},
+      {"203.0.113.192", {"tag_0", "tag_1", "tag_2"}},
+      {"203.0.113.255", {"tag_0", "tag_1", "tag_2"}},
+      {"198.51.100.1", {"tag_0", "tag_3"}},
+      {"2001:db8::ffff", {"tag_4", "tag_5", "tag_6"}},
+      {"2001:db8:1::ffff", {"tag_4", "tag_7"}}
+
+  };
+  expectIPAndTags(test_case);
 }
 
 } // namespace LcTrie