dns: adding limits to the key value store. (envoyproxy#20698)

Risk Level: low Testing: unit tests Docs Changes: inline Release Notes: [Optional Runtime guard:] Signed-off-by: Alyssa Wilk <[email protected]>
AppMonet · Apr 13, 2022 · 888c603 · 888c603
1 parent 8f13d85
commit 888c603
Show file tree

Hide file tree

Showing 7 changed files with 68 additions and 45 deletions.
diff --git a/api/envoy/extensions/key_value/file_based/v3/config.proto b/api/envoy/extensions/key_value/file_based/v3/config.proto
@@ -3,6 +3,7 @@ syntax = "proto3";
 package envoy.extensions.key_value.file_based.v3;
 
 import "google/protobuf/duration.proto";
+import "google/protobuf/wrappers.proto";
 
 import "xds/annotations/v3/status.proto";
 
@@ -28,4 +29,8 @@ message FileBasedKeyValueStoreConfig {
 
   // The interval at which the key value store should be flushed to the file.
   google.protobuf.Duration flush_interval = 2;
+
+  // The maximum number of entries to cache, or 0 to allow for unlimited entries.
+  // Defaults to 1000 if not present.
+  google.protobuf.UInt32Value max_entries = 3;
 }
diff --git a/source/common/common/BUILD b/source/common/common/BUILD
@@ -155,11 +155,13 @@ envoy_cc_library(
     hdrs = ["key_value_store_base.h"],
     external_deps = [
         "abseil_cleanup",
+        "quiche_quic_platform",
     ],
     deps = [
         "//envoy/common:key_value_store_interface",
         "//envoy/event:dispatcher_interface",
         "//envoy/filesystem:filesystem_interface",
+        "@com_github_google_quiche//:quiche_common_lib",
     ],
 )
 

diff --git a/source/common/common/key_value_store_base.cc b/source/common/common/key_value_store_base.cc
@@ -31,8 +31,8 @@ absl::optional<absl::string_view> getToken(absl::string_view& contents, std::str
 } // namespace
 
 KeyValueStoreBase::KeyValueStoreBase(Event::Dispatcher& dispatcher,
-                                     std::chrono::milliseconds flush_interval)
-    : flush_timer_(dispatcher.createTimer([this, flush_interval]() {
+                                     std::chrono::milliseconds flush_interval, uint32_t max_entries)
+    : max_entries_(max_entries), flush_timer_(dispatcher.createTimer([this, flush_interval]() {
         flush();
         flush_timer_->enableTimer(flush_interval);
       })) {
@@ -41,13 +41,7 @@ KeyValueStoreBase::KeyValueStoreBase(Event::Dispatcher& dispatcher,
   }
 }
 
-// Assuming |contents| is in the format
-// [length]\n[key]\n[length]\n[value]
-// parses contents into the provided store.
-// This is best effort, and will return false on failure without clearing
-// partially parsed data.
-bool KeyValueStoreBase::parseContents(absl::string_view contents,
-                                      absl::flat_hash_map<std::string, std::string>& store) const {
+bool KeyValueStoreBase::parseContents(absl::string_view contents) {
   std::string error;
   while (!contents.empty()) {
     absl::optional<absl::string_view> key = getToken(contents, error);
@@ -59,43 +53,49 @@ bool KeyValueStoreBase::parseContents(absl::string_view contents,
       ENVOY_LOG(warn, error);
       return false;
     }
-    store.emplace(std::string(key.value()), std::string(value.value()));
+    addOrUpdate(key.value(), value.value());
   }
   return true;
 }
 
-void KeyValueStoreBase::addOrUpdate(absl::string_view key, absl::string_view value) {
-  store_.insert_or_assign(key, std::string(value));
+void KeyValueStoreBase::addOrUpdate(absl::string_view key_view, absl::string_view value_view) {
+  ENVOY_BUG(!under_iterate_, "addOrUpdate under the stack of iterate");
+  std::string key(key_view);
+  std::string value(value_view);
+  // Attempt to insert the entry into the store. If it already exists, remove
+  // the old entry and insert the new one so it will be in the proper place in
+  // the linked list.
+  if (!store_.emplace(key, value).second) {
+    store_.erase(key);
+    store_.emplace(key, value);
+  }
+  if (max_entries_ && store_.size() > max_entries_) {
+    store_.pop_front();
+  }
   if (!flush_timer_->enabled()) {
     flush();
   }
 }
 
 void KeyValueStoreBase::remove(absl::string_view key) {
-  store_.erase(key);
+  ENVOY_BUG(!under_iterate_, "remove under the stack of iterate");
+  store_.erase(std::string(key));
   if (!flush_timer_->enabled()) {
     flush();
   }
 }
 
 absl::optional<absl::string_view> KeyValueStoreBase::get(absl::string_view key) {
-  auto it = store_.find(key);
+  auto it = store_.find(std::string(key));
   if (it == store_.end()) {
     return {};
   }
   return it->second;
 }
 
 void KeyValueStoreBase::iterate(ConstIterateCb cb) const {
-#ifndef NDEBUG
-  // When running in debug mode, verify we don't modify the underlying store
-  // while iterating.
-  absl::flat_hash_map<std::string, std::string> store_before_iteration = store_;
-  absl::Cleanup verify_store_is_not_modified = [this, &store_before_iteration] {
-    ASSERT(store_ == store_before_iteration,
-           "Expected iterate to not modify the underlying store.");
-  };
-#endif
+  under_iterate_ = true;
+  absl::Cleanup restore_under_iterate = [this] { under_iterate_ = false; };
 
   for (const auto& [key, value] : store_) {
     Iterate ret = cb(key, value);

diff --git a/source/common/common/key_value_store_base.h b/source/common/common/key_value_store_base.h
@@ -6,27 +6,25 @@
 
 #include "source/common/common/logger.h"
 
-#include "absl/container/flat_hash_map.h"
+#include "quiche/common/quiche_linked_hash_map.h"
 
 namespace Envoy {
 
 // This is the base implementation of the KeyValueStore. It handles the various
 // functions other than flush(), which will be implemented by subclasses.
-//
-// Note this implementation HAS UNBOUNDED SIZE.
-// It is assumed the callers manage the number of entries. Use with care.
 class KeyValueStoreBase : public KeyValueStore,
                           public Logger::Loggable<Logger::Id::key_value_store> {
 public:
   // Sets up flush() for the configured interval.
-  KeyValueStoreBase(Event::Dispatcher& dispatcher, std::chrono::milliseconds flush_interval);
+  KeyValueStoreBase(Event::Dispatcher& dispatcher, std::chrono::milliseconds flush_interval,
+                    uint32_t max_entries);
 
   // If |contents| is in the form of
   // [length]\n[key][length]\n[value]
-  // parses key value pairs from |contents| into the store provided.
+  // parses key value pairs from |contents| and inserts into store_.
   // Returns true on success and false on failure.
-  bool parseContents(absl::string_view contents,
-                     absl::flat_hash_map<std::string, std::string>& store) const;
+  bool parseContents(absl::string_view contents);
+
   std::string error;
   // KeyValueStore
   void addOrUpdate(absl::string_view key, absl::string_view value) override;
@@ -35,8 +33,11 @@ class KeyValueStoreBase : public KeyValueStore,
   void iterate(ConstIterateCb cb) const override;
 
 protected:
+  const uint32_t max_entries_;
   const Event::TimerPtr flush_timer_;
-  absl::flat_hash_map<std::string, std::string> store_;
+  quiche::QuicheLinkedHashMap<std::string, std::string> store_;
+  // Used for validation only.
+  mutable bool under_iterate_{};
 };
 
 } // namespace Envoy
diff --git a/source/extensions/key_value/file_based/config.cc b/source/extensions/key_value/file_based/config.cc
@@ -9,15 +9,15 @@ namespace KeyValue {
 FileBasedKeyValueStore::FileBasedKeyValueStore(Event::Dispatcher& dispatcher,
                                                std::chrono::milliseconds flush_interval,
                                                Filesystem::Instance& file_system,
-                                               const std::string& filename)
-    : KeyValueStoreBase(dispatcher, flush_interval), file_system_(file_system),
+                                               const std::string& filename, uint32_t max_entries)
+    : KeyValueStoreBase(dispatcher, flush_interval, max_entries), file_system_(file_system),
       filename_(filename) {
   if (!file_system_.fileExists(filename_)) {
     ENVOY_LOG(info, "File for key value store does not yet exist: {}", filename);
     return;
   }
   const std::string contents = file_system_.fileReadToEnd(filename_);
-  if (!parseContents(contents, store_)) {
+  if (!parseContents(contents)) {
     ENVOY_LOG(warn, "Failed to parse key value store file {}", filename);
   }
 }
@@ -48,10 +48,11 @@ KeyValueStorePtr FileBasedKeyValueStoreFactory::createStore(
   const auto file_config = MessageUtil::anyConvertAndValidate<
       envoy::extensions::key_value::file_based::v3::FileBasedKeyValueStoreConfig>(
       typed_config.config().typed_config(), validation_visitor);
-  auto milliseconds =
+  const auto milliseconds =
       std::chrono::milliseconds(DurationUtil::durationToMilliseconds(file_config.flush_interval()));
+  const uint32_t max_entries = PROTOBUF_GET_WRAPPED_OR_DEFAULT(file_config, max_entries, 1000);
   return std::make_unique<FileBasedKeyValueStore>(dispatcher, milliseconds, file_system,
-                                                  file_config.filename());
+                                                  file_config.filename(), max_entries);
 }
 
 REGISTER_FACTORY(FileBasedKeyValueStoreFactory, KeyValueStoreFactory);

diff --git a/source/extensions/key_value/file_based/config.h b/source/extensions/key_value/file_based/config.h
@@ -18,7 +18,8 @@ namespace KeyValue {
 class FileBasedKeyValueStore : public KeyValueStoreBase {
 public:
   FileBasedKeyValueStore(Event::Dispatcher& dispatcher, std::chrono::milliseconds flush_interval,
-                         Filesystem::Instance& file_system, const std::string& filename);
+                         Filesystem::Instance& file_system, const std::string& filename,
+                         uint32_t max_entries);
   // KeyValueStore
   void flush() override;
 

diff --git a/test/extensions/key_value/file_based/key_value_store_test.cc b/test/extensions/key_value/file_based/key_value_store_test.cc
@@ -25,10 +25,10 @@ class KeyValueStoreTest : public testing::Test {
     createStore();
   }
 
-  void createStore() {
+  void createStore(uint32_t max_entries = 0) {
     flush_timer_ = new NiceMock<Event::MockTimer>(&dispatcher_);
-    store_ = std::make_unique<FileBasedKeyValueStore>(dispatcher_, flush_interval_,
-                                                      Filesystem::fileSystemForTest(), filename_);
+    store_ = std::make_unique<FileBasedKeyValueStore>(
+        dispatcher_, flush_interval_, Filesystem::fileSystemForTest(), filename_, max_entries);
   }
   NiceMock<Event::MockDispatcher> dispatcher_;
   std::string filename_;
@@ -47,6 +47,21 @@ TEST_F(KeyValueStoreTest, Basic) {
   EXPECT_EQ(absl::nullopt, store_->get("foo"));
 }
 
+TEST_F(KeyValueStoreTest, MaxEntries) {
+  createStore(2);
+  // Add 2 entries.
+  store_->addOrUpdate("1", "a");
+  store_->addOrUpdate("2", "b");
+  EXPECT_EQ("a", store_->get("1").value());
+  EXPECT_EQ("b", store_->get("2").value());
+
+  // Adding '3' should evict '1'.
+  store_->addOrUpdate("3", "c");
+  EXPECT_EQ("c", store_->get("3").value());
+  EXPECT_EQ("b", store_->get("2").value());
+  EXPECT_EQ(absl::nullopt, store_->get("1"));
+}
+
 TEST_F(KeyValueStoreTest, Persist) {
   store_->addOrUpdate("foo", "bar");
   store_->addOrUpdate("ba\nz", "ee\np");
@@ -114,8 +129,7 @@ TEST_F(KeyValueStoreTest, ShouldCrashIfIterateCallbackAddsOrUpdatesStore) {
     return KeyValueStore::Iterate::Continue;
   };
 
-  EXPECT_DEATH(store_->iterate(update_value_callback),
-               "Expected iterate to not modify the underlying store.");
+  EXPECT_DEATH(store_->iterate(update_value_callback), "addOrUpdate under the stack of iterate");
 
   KeyValueStore::ConstIterateCb add_key_callback = [this](const std::string& key,
                                                           const std::string&) {
@@ -125,8 +139,7 @@ TEST_F(KeyValueStoreTest, ShouldCrashIfIterateCallbackAddsOrUpdatesStore) {
     }
     return KeyValueStore::Iterate::Continue;
   };
-  EXPECT_DEATH(store_->iterate(add_key_callback),
-               "Expected iterate to not modify the underlying store.");
+  EXPECT_DEATH(store_->iterate(add_key_callback), "addOrUpdate under the stack of iterate");
 }
 #endif