http: Adding flexible remove function to HeaderMap (#3186)

Signed-off-by: Alyssa Wilk <[email protected]>
envoyproxy · Apr 24, 2018 · d8d089a · d8d089a
1 parent a87d265
commit d8d089a
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 0 deletions.
diff --git a/include/envoy/http/header_map.h b/include/envoy/http/header_map.h
@@ -465,6 +465,12 @@ class HeaderMap {
    */
   virtual void remove(const LowerCaseString& key) PURE;
 
+  /**
+   * Remove all instances of headers where the key begins with the supplied prefix.
+   * @param prefix supplies the prefix to match header keys against.
+   */
+  virtual void removePrefix(const LowerCaseString& prefix) PURE;
+
   /**
    * @return the number of headers in the map.
    */

diff --git a/source/common/http/header_map_impl.cc b/source/common/http/header_map_impl.cc
@@ -9,6 +9,8 @@
 #include "common/common/utility.h"
 #include "common/singleton/const_singleton.h"
 
+#include "absl/strings/match.h"
+
 namespace Envoy {
 namespace Http {
 
@@ -464,6 +466,25 @@ void HeaderMapImpl::remove(const LowerCaseString& key) {
   }
 }
 
+void HeaderMapImpl::removePrefix(const LowerCaseString& prefix) {
+  headers_.remove_if([&](const HeaderEntryImpl& entry) {
+    bool to_remove = absl::StartsWith(entry.key().getStringView(), prefix.get());
+    if (to_remove) {
+      // If this header should be removed, make sure any references in the
+      // static lookup table are cleared as well.
+      StaticLookupEntry::EntryCb cb =
+          ConstSingleton<StaticLookupTable>::get().find(entry.key().c_str());
+      if (cb) {
+        StaticLookupResponse ref_lookup_response = cb(*this);
+        if (ref_lookup_response.entry_) {
+          *ref_lookup_response.entry_ = nullptr;
+        }
+      }
+    }
+    return to_remove;
+  });
+}
+
 HeaderMapImpl::HeaderEntryImpl& HeaderMapImpl::maybeCreateInline(HeaderEntryImpl** entry,
                                                                  const LowerCaseString& key) {
   if (*entry) {

diff --git a/source/common/http/header_map_impl.h b/source/common/http/header_map_impl.h
@@ -68,6 +68,7 @@ class HeaderMapImpl : public HeaderMap {
   void iterateReverse(ConstIterateCb cb, void* context) const override;
   Lookup lookup(const LowerCaseString& key, const HeaderEntry** entry) const override;
   void remove(const LowerCaseString& key) override;
+  void removePrefix(const LowerCaseString& key) override;
   size_t size() const override { return headers_.size(); }
 
 protected:

diff --git a/test/common/http/header_map_impl_test.cc b/test/common/http/header_map_impl_test.cc
@@ -349,6 +349,43 @@ TEST(HeaderMapImplTest, Remove) {
   EXPECT_EQ(0UL, headers.size());
 }
 
+TEST(HeaderMapImplTest, RemoveRegex) {
+  // These will match.
+  LowerCaseString key1 = LowerCaseString("X-prefix-foo");
+  LowerCaseString key3 = LowerCaseString("X-Prefix-");
+  LowerCaseString key5 = LowerCaseString("x-prefix-eep");
+  // These will not.
+  LowerCaseString key2 = LowerCaseString(" x-prefix-foo");
+  LowerCaseString key4 = LowerCaseString("y-x-prefix-foo");
+
+  HeaderMapImpl headers;
+  headers.addReference(key1, "value");
+  headers.addReference(key2, "value");
+  headers.addReference(key3, "value");
+  headers.addReference(key4, "value");
+  headers.addReference(key5, "value");
+
+  // Test removing the first header, middle headers, and the end header.
+  headers.removePrefix(LowerCaseString("x-prefix-"));
+  EXPECT_EQ(nullptr, headers.get(key1));
+  EXPECT_NE(nullptr, headers.get(key2));
+  EXPECT_EQ(nullptr, headers.get(key3));
+  EXPECT_NE(nullptr, headers.get(key4));
+  EXPECT_EQ(nullptr, headers.get(key5));
+
+  // Remove all headers.
+  headers.removePrefix(LowerCaseString(""));
+  EXPECT_EQ(nullptr, headers.get(key2));
+  EXPECT_EQ(nullptr, headers.get(key4));
+
+  // Add inline and remove by regex
+  headers.insertContentLength().value(5);
+  EXPECT_STREQ("5", headers.ContentLength()->value().c_str());
+  EXPECT_EQ(1UL, headers.size());
+  headers.removePrefix(LowerCaseString("content"));
+  EXPECT_EQ(nullptr, headers.ContentLength());
+}
+
 TEST(HeaderMapImplTest, SetRemovesAllValues) {
   HeaderMapImpl headers;