From 8ce3f0232f969e8990493de2636fad19cf8df4cd Mon Sep 17 00:00:00 2001
From: Aman Sharma <amsharma@meta.com>
Date: Sat, 4 Nov 2023 00:30:54 -0700
Subject: [PATCH] Make a HeadersRateLimitFilter

Summary: I'm making separate RateLimitFilters for each of the events (control messages, headers, direct error handling, etc) that we rate limit. This is one step of that process.

Reviewed By: afrind

Differential Revision: D50698700

fbshipit-source-id: c6dedc97acb090b2aecb19b1680fe65beb72add7
---
 .../lib/http/codec/HeadersRateLimitFilter.h   | 55 +++++++++++++++++++
 proxygen/lib/http/codec/RateLimitFilter.cpp   |  7 +--
 2 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 proxygen/lib/http/codec/HeadersRateLimitFilter.h

diff --git a/proxygen/lib/http/codec/HeadersRateLimitFilter.h b/proxygen/lib/http/codec/HeadersRateLimitFilter.h
new file mode 100644
index 0000000000..72e1c384f6
--- /dev/null
+++ b/proxygen/lib/http/codec/HeadersRateLimitFilter.h
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <proxygen/lib/http/codec/HTTP2Constants.h>
+#include <proxygen/lib/http/codec/RateLimitFilter.h>
+
+namespace proxygen {
+
+class HeadersRateLimitFilter : public RateLimitFilter {
+ public:
+  static const uint32_t kDefaultMaxEventsPerInterval = 50000;
+  static const uint32_t kMaxEventsPerIntervalLowerBound = 100;
+  static constexpr std::chrono::milliseconds kDefaultTimeoutDuration{100};
+
+  explicit HeadersRateLimitFilter(folly::HHWheelTimer* timer,
+                                  HTTPSessionStats* httpSessionStats)
+      : RateLimitFilter(timer, httpSessionStats) {
+    maxEventsInInterval_ = kDefaultMaxEventsPerInterval;
+    timeoutDuration_ = kDefaultTimeoutDuration;
+  }
+
+  void onHeadersComplete(StreamID stream,
+                         std::unique_ptr<HTTPMessage> msg) override {
+    if (!incrementNumEventsInCurrentInterval()) {
+      callback_->onHeadersComplete(stream, std::move(msg));
+    } else {
+      callback_->onGoaway(http2::kMaxStreamID, ErrorCode::NO_ERROR);
+    }
+  }
+
+  void recordNumEventsInCurrentInterval(uint32_t numEvents) override {
+    if (httpSessionStats_) {
+      httpSessionStats_->recordHeadersInInterval(numEvents);
+    }
+  }
+
+  void recordRateLimitBreached() override {
+    if (httpSessionStats_) {
+      httpSessionStats_->recordHeadersRateLimited();
+    }
+  }
+
+  uint32_t getMaxEventsPerInvervalLowerBound() const override {
+    return kMaxEventsPerIntervalLowerBound;
+  }
+};
+
+} // namespace proxygen
diff --git a/proxygen/lib/http/codec/RateLimitFilter.cpp b/proxygen/lib/http/codec/RateLimitFilter.cpp
index 9f718cd10f..f586dbbbf5 100644
--- a/proxygen/lib/http/codec/RateLimitFilter.cpp
+++ b/proxygen/lib/http/codec/RateLimitFilter.cpp
@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <proxygen/lib/http/codec/HeadersRateLimitFilter.h>
 #include <proxygen/lib/http/codec/RateLimitFilter.h>
 
 namespace proxygen {
@@ -26,12 +27,10 @@ std::string_view RateLimitFilter::toStr(Type type) {
 }
 
 std::unique_ptr<RateLimitFilter> RateLimitFilter::createRateLimitFilter(
-    Type type,
-    folly::HHWheelTimer* /* timer */,
-    HTTPSessionStats* /* httpSessionStats */) {
+    Type type, folly::HHWheelTimer* timer, HTTPSessionStats* httpSessionStats) {
   switch (type) {
     case Type::HEADERS:
-      return nullptr;
+      return std::make_unique<HeadersRateLimitFilter>(timer, httpSessionStats);
     case Type::MISC_CONTROL_MSGS:
       return nullptr;
     case Type::RSTS: