facebook · mjoras · Oct 10, 2023 · Oct 10, 2023
diff --git a/proxygen/lib/http/codec/ControlMessageRateLimitFilter.h b/proxygen/lib/http/codec/ControlMessageRateLimitFilter.h
@@ -24,37 +24,61 @@ constexpr uint32_t kDefaultMaxDirectErrorHandlingPerInterval = 100;
 constexpr uint32_t kMaxDirectErrorHandlingPerIntervalLowerBound = 50;
 constexpr std::chrono::milliseconds kDefaultDirectErrorHandlingDuration{100};
 
+constexpr uint32_t kDefaultMaxHeadersPerInterval = 500;
+constexpr uint32_t kMaxHeadersPerIntervalLowerBound = 100;
+constexpr std::chrono::milliseconds kDefaultHeadersDuration{100};
+
+enum RateLimitTarget {
+  CONTROL_MSGS,
+  DIRECT_ERROR_HANDLING,
+  HEADERS,
+};
+
 /**
  * This class implements the rate limiting logic for control messages and
  * stream errors (that might produce HTTP error pages).  If a rate limit is
  * exeeded, the callback is converted to a session level error with
  * ProxygenError = kErrorDropped.  This is a signal to the codec callback that
  * the codec would like the connection dropped.
+ *
+ * TODO: Refactor this into separate filters, or group related parameters
+ * into structs.
  */
 class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
  public:
   explicit ControlMessageRateLimitFilter(folly::HHWheelTimer* timer,
                                          HTTPSessionStats* httpSessionStats)
       : resetControlMessages_(numControlMsgsInCurrentInterval_,
+                              RateLimitTarget::CONTROL_MSGS,
                               httpSessionStats),
+        resetDirectErrors_(numDirectErrorHandlingInCurrentInterval_,
+                           RateLimitTarget::DIRECT_ERROR_HANDLING,
+                           httpSessionStats),
+        resetHeaders_(numHeadersInCurrentInterval_,
+                      RateLimitTarget::HEADERS,
+                      httpSessionStats),
         timer_(timer),
         httpSessionStats_(httpSessionStats) {
   }
 
   void setSessionStats(HTTPSessionStats* httpSessionStats) {
     httpSessionStats_ = httpSessionStats;
     resetControlMessages_.httpSessionStats = httpSessionStats;
+    resetHeaders_.httpSessionStats = httpSessionStats;
   }
 
-  void setParams(
-      uint32_t maxControlMsgsPerInterval,
-      uint32_t maxDirectErrorHandlingPerInterval,
-      std::chrono::milliseconds controlMsgIntervalDuration,
-      std::chrono::milliseconds directErrorHandlingIntervalDuration) {
+  void setParams(uint32_t maxControlMsgsPerInterval,
+                 uint32_t maxDirectErrorHandlingPerInterval,
+                 uint32_t maxHeadersPerInterval,
+                 std::chrono::milliseconds controlMsgIntervalDuration,
+                 std::chrono::milliseconds directErrorHandlingIntervalDuration,
+                 std::chrono::milliseconds headersIntervalDuration) {
     maxControlMsgsPerInterval_ = maxControlMsgsPerInterval;
     maxDirectErrorHandlingPerInterval_ = maxDirectErrorHandlingPerInterval;
     controlMsgIntervalDuration_ = controlMsgIntervalDuration;
     directErrorHandlingIntervalDuration_ = directErrorHandlingIntervalDuration;
+    maxHeadersPerInterval_ = maxHeadersPerInterval;
+    headersIntervalDuration_ = headersIntervalDuration;
   }
 
   // Filter functions
@@ -86,6 +110,13 @@ class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
     }
   }
 
+  void onHeadersComplete(StreamID stream,
+                         std::unique_ptr<HTTPMessage> msg) override {
+    if (!incrementNumHeadersInCurInterval()) {
+      callback_->onHeadersComplete(stream, std::move(msg));
+    }
+  }
+
   void onError(HTTPCodec::StreamID streamID,
                const HTTPException& error,
                bool newTxn) override {
@@ -104,6 +135,7 @@ class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
   void detachThreadLocals() {
     resetControlMessages_.cancelTimeout();
     resetDirectErrors_.cancelTimeout();
+    resetHeaders_.cancelTimeout();
     timer_ = nullptr;
     // Free pass when switching threads
     numControlMsgsInCurrentInterval_ = 0;
@@ -140,6 +172,25 @@ class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
     return false;
   }
 
+  bool incrementNumHeadersInCurInterval() {
+    if (numHeadersInCurrentInterval_ == 0) {
+      // The first header (or first after a reset) schedules the next
+      // reset timer
+      CHECK(timer_);
+      timer_->scheduleTimeout(&resetHeaders_, headersIntervalDuration_);
+    }
+
+    if (++numHeadersInCurrentInterval_ > maxHeadersPerInterval_) {
+      if (httpSessionStats_) {
+        httpSessionStats_->recordHeadersRateLimited();
+      }
+      callback_->onGoaway(http2::kMaxStreamID, ErrorCode::NO_ERROR);
+      return true;
+    }
+
+    return false;
+  }
+
   bool incrementDirectErrorHandlingInCurInterval() {
     if (numDirectErrorHandlingInCurrentInterval_ == 0) {
       // The first control message (or first after a reset) schedules the next
@@ -168,20 +219,34 @@ class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
   class ResetCounterTimeout : public folly::HHWheelTimer::Callback {
    public:
     explicit ResetCounterTimeout(uint32_t& counterIn,
+                                 RateLimitTarget rateLimitTargetIn,
                                  HTTPSessionStats* httpSessionStatsIn = nullptr)
-        : counter(counterIn), httpSessionStats(httpSessionStatsIn) {
+        : counter(counterIn),
+          rateLimitTarget(rateLimitTargetIn),
+          httpSessionStats(httpSessionStatsIn) {
     }
 
     void timeoutExpired() noexcept override {
       if (counter > 0 && httpSessionStats) {
-        httpSessionStats->recordControlMsgsInInterval(counter);
+        switch (rateLimitTarget) {
+          case RateLimitTarget::CONTROL_MSGS:
+            httpSessionStats->recordControlMsgsInInterval(counter);
+            break;
+          case RateLimitTarget::DIRECT_ERROR_HANDLING:
+            // No stats for this one
+            break;
+          case RateLimitTarget::HEADERS:
+            httpSessionStats->recordHeadersInInterval(counter);
+            break;
+        }
       }
       counter = 0;
     }
     void callbackCanceled() noexcept override {
     }
 
     uint32_t& counter;
+    RateLimitTarget rateLimitTarget;
     HTTPSessionStats* httpSessionStats{nullptr};
   };
 
@@ -197,14 +262,18 @@ class ControlMessageRateLimitFilter : public PassThroughHTTPCodecFilter {
   uint32_t maxDirectErrorHandlingPerInterval_{
       kDefaultMaxDirectErrorHandlingPerInterval};
 
+  uint32_t numHeadersInCurrentInterval_{0};
+  uint32_t maxHeadersPerInterval_{kDefaultMaxHeadersPerInterval};
+
   std::chrono::milliseconds controlMsgIntervalDuration_{
       kDefaultControlMsgDuration};
   std::chrono::milliseconds directErrorHandlingIntervalDuration_{
       kDefaultDirectErrorHandlingDuration};
+  std::chrono::milliseconds headersIntervalDuration_{kDefaultHeadersDuration};
 
   ResetCounterTimeout resetControlMessages_;
-  ResetCounterTimeout resetDirectErrors_{
-      numDirectErrorHandlingInCurrentInterval_};
+  ResetCounterTimeout resetDirectErrors_;
+  ResetCounterTimeout resetHeaders_;
   folly::HHWheelTimer* timer_{nullptr};
   HTTPSessionStats* httpSessionStats_{nullptr};
 };

diff --git a/proxygen/lib/http/session/HTTPSession.cpp b/proxygen/lib/http/session/HTTPSession.cpp
@@ -222,7 +222,8 @@ void HTTPSession::setupCodec() {
     // existing flow control filter
   }
   if (codec_->supportsParallelRequests() && !controlMessageRateLimitFilter_ &&
-      sock_) {
+      sock_ &&
+      codec_->getTransportDirection() == TransportDirection::DOWNSTREAM) {
     controlMessageRateLimitFilter_ = new ControlMessageRateLimitFilter(
         &getEventBase()->timer(), sessionStats_);
     codec_.addFilters(std::unique_ptr<ControlMessageRateLimitFilter>(

diff --git a/proxygen/lib/http/session/HTTPSessionBase.cpp b/proxygen/lib/http/session/HTTPSessionBase.cpp
@@ -75,8 +75,11 @@ void HTTPSessionBase::setSessionStats(HTTPSessionStats* stats) {
 void HTTPSessionBase::setControlMessageRateLimitParams(
     uint32_t maxControlMsgsPerInterval,
     uint32_t maxDirectErrorHandlingPerInterval,
+    uint32_t maxHeadersPerInterval,
     std::chrono::milliseconds controlMsgIntervalDuration,
-    std::chrono::milliseconds directErrorHandlingIntervalDuration) {
+    std::chrono::milliseconds directErrorHandlingIntervalDuration,
+    std::chrono::milliseconds headersIntervalDuration) {
+
   if (maxControlMsgsPerInterval < kMaxControlMsgsPerIntervalLowerBound) {
     XLOG_EVERY_MS(WARNING, 60000)
         << "Invalid maxControlMsgsPerInterval: " << maxControlMsgsPerInterval;
@@ -92,12 +95,20 @@ void HTTPSessionBase::setControlMessageRateLimitParams(
         kMaxDirectErrorHandlingPerIntervalLowerBound;
   }
 
+  if (maxHeadersPerInterval < kMaxHeadersPerIntervalLowerBound) {
+    XLOG_EVERY_MS(WARNING, 60000)
+        << "Invalid maxHeadersPerInterval: " << maxHeadersPerInterval;
+    maxHeadersPerInterval = kMaxHeadersPerIntervalLowerBound;
+  }
+
   if (controlMessageRateLimitFilter_) {
     controlMessageRateLimitFilter_->setParams(
         maxControlMsgsPerInterval,
         maxDirectErrorHandlingPerInterval,
+        maxHeadersPerInterval,
         controlMsgIntervalDuration,
-        directErrorHandlingIntervalDuration);
+        directErrorHandlingIntervalDuration,
+        headersIntervalDuration);
   }
 }
 

diff --git a/proxygen/lib/http/session/HTTPSessionBase.h b/proxygen/lib/http/session/HTTPSessionBase.h
@@ -176,10 +176,13 @@ class HTTPSessionBase : public wangle::ManagedConnection {
       uint32_t maxControlMsgsPerInterval = kDefaultMaxControlMsgsPerInterval,
       uint32_t maxDirectErrorHandlingPerInterval =
           kDefaultMaxDirectErrorHandlingPerInterval,
+      uint32_t maxHeadersPerInterval = kDefaultMaxHeadersPerInterval,
       std::chrono::milliseconds controlMsgIntervalDuration =
           kDefaultControlMsgDuration,
       std::chrono::milliseconds directErrorHandlingIntervalDuration =
-          kDefaultDirectErrorHandlingDuration);
+          kDefaultDirectErrorHandlingDuration,
+      std::chrono::milliseconds headersIntervalDuration =
+          kDefaultHeadersDuration);
 
   InfoCallback* getInfoCallback() const {
     return infoCallback_;

diff --git a/proxygen/lib/http/session/HTTPSessionStats.h b/proxygen/lib/http/session/HTTPSessionStats.h
@@ -38,6 +38,10 @@ class HTTPSessionStats : public TTLBAStats {
   }
   virtual void recordControlMsgRateLimited() noexcept {
   }
+  virtual void recordHeadersInInterval(int64_t) noexcept {
+  }
+  virtual void recordHeadersRateLimited() noexcept {
+  }
 };
 
 } // namespace proxygen
diff --git a/proxygen/lib/http/session/test/HTTPDownstreamSessionTest.cpp b/proxygen/lib/http/session/test/HTTPDownstreamSessionTest.cpp
@@ -96,7 +96,7 @@ class HTTPDownstreamTest : public testing::Test {
 
     httpSession_->setFlowControl(
         flowControl[0], flowControl[1], flowControl[2]);
-    httpSession_->setEgressSettings({{SettingsId::MAX_CONCURRENT_STREAMS, 80},
+    httpSession_->setEgressSettings({{SettingsId::MAX_CONCURRENT_STREAMS, 200},
                                      {SettingsId::HEADER_TABLE_SIZE, 5555},
                                      {SettingsId::ENABLE_PUSH, 1},
                                      {SettingsId::ENABLE_EX_HEADERS, 1}});
@@ -3727,6 +3727,30 @@ TEST_F(HTTP2DownstreamSessionTest, TestPriorityFCBlocked) {
   this->eventBase_.loop();
 }
 
+TEST_F(HTTP2DownstreamSessionTest, TestHeadersRateLimitExceeded) {
+  httpSession_->setControlMessageRateLimitParams(
+      1000, 1000, 100, seconds(0), seconds(0), seconds(0));
+
+  std::vector<std::unique_ptr<testing::StrictMock<MockHTTPHandler>>> handlers;
+  for (int i = 0; i < 100; i++) {
+    auto handler = addSimpleStrictHandler();
+    auto rawHandler = handler.get();
+    handlers.push_back(std::move(handler));
+    rawHandler->expectHeaders();
+    rawHandler->expectEOM(
+        [rawHandler] { rawHandler->sendReplyWithBody(200, 100); });
+    sendRequest();
+  }
+  // Straw that breaks the camel's back
+  sendRequest();
+  for (int i = 0; i < 100; i++) {
+    handlers[i]->expectGoaway();
+    handlers[i]->expectDetachTransaction();
+  }
+  expectDetachSession();
+  flushRequestsAndLoopN(2);
+}
+
 TEST_F(HTTP2DownstreamSessionTest, TestControlMsgRateLimitExceeded) {
   auto streamid = clientCodec_->createStream();
 
@@ -3756,7 +3780,7 @@ TEST_F(HTTP2DownstreamSessionTest, TestControlMsgResetRateLimitTouched) {
 
   auto streamid = clientCodec_->createStream();
 
-  httpSession_->setControlMessageRateLimitParams(100, 100, milliseconds(0));
+  httpSession_->setControlMessageRateLimitParams(10, 100, 100, milliseconds(0));
 
   // Send 97 PRIORITY, 1 SETTINGS, and 2 PING frames. This doesn't exceed the
   // limit of 10.
@@ -3797,7 +3821,7 @@ TEST_F(HTTP2DownstreamSessionTest, TestControlMsgResetRateLimitTouched) {
 }
 
 TEST_F(HTTP2DownstreamSessionTest, DirectErrorHandlingLimitTouched) {
-  httpSession_->setControlMessageRateLimitParams(100, 50, milliseconds(0));
+  httpSession_->setControlMessageRateLimitParams(100, 10, 100, milliseconds(0));
 
   // Send 50 messages, each of which cause direct error handling. Since
   // this doesn't exceed the limit, this should not cause the connection
@@ -3829,7 +3853,7 @@ TEST_F(HTTP2DownstreamSessionTest, DirectErrorHandlingLimitTouched) {
 }
 
 TEST_F(HTTP2DownstreamSessionTest, DirectErrorHandlingLimitExceeded) {
-  httpSession_->setControlMessageRateLimitParams(100, 50, milliseconds(0));
+  httpSession_->setControlMessageRateLimitParams(100, 10, 100, milliseconds(0));
 
   // Send eleven messages, each of which causes direct error handling. Since
   // this exceeds the limit, the connection should be dropped.

diff --git a/proxygen/lib/http/session/test/HTTPUpstreamSessionTest.cpp b/proxygen/lib/http/session/test/HTTPUpstreamSessionTest.cpp
@@ -2931,7 +2931,7 @@ TEST_F(HTTP2UpstreamSessionTest, AttachDetach) {
     httpSession_->detachThreadLocals();
     httpSession_->attachThreadLocals(
         &base, nullptr, timerInstance, nullptr, fn, nullptr, nullptr);
-    EXPECT_EQ(filterCount, 3);
+    EXPECT_EQ(filterCount, 2);
     filterCount = 0;
     base.loopOnce();
   }
@@ -2997,7 +2997,7 @@ TEST_F(HTTP2UpstreamSessionTest, DetachFlowControlTimeout) {
     httpSession_->detachThreadLocals();
     httpSession_->attachThreadLocals(
         &base, nullptr, timerInstance, nullptr, fn, nullptr, nullptr);
-    EXPECT_EQ(filterCount, 3);
+    EXPECT_EQ(filterCount, 2);
     filterCount = 0;
     base.loopOnce();
   }

diff --git a/proxygen/lib/http/stats/ThreadLocalHTTPSessionStats.cpp b/proxygen/lib/http/stats/ThreadLocalHTTPSessionStats.cpp
@@ -39,6 +39,8 @@ TLHTTPSessionStats::TLHTTPSessionStats(const std::string& prefix)
       ttbtxExceedLimit(prefix + "_ttbtx_exceed_limit", facebook::fb303::SUM),
       ctrlMsgsRateLimited(prefix + "_ctrl_msgs_rate_limited",
                           facebook::fb303::SUM),
+      headersRateLimited(prefix + "_headers_rate_limited",
+                         facebook::fb303::SUM),
       txnsPerSession(prefix + "_txn_per_session",
                      1,
                      0,
@@ -64,6 +66,15 @@ TLHTTPSessionStats::TLHTTPSessionStats(const std::string& prefix)
           facebook::fb303::AVG,
           50,
           99,
+          100),
+      headersInInterval(
+          prefix + "_headers_in_interval",
+          1 /* bucketWidth */,
+          0 /* min */,
+          500 /* max, keep in sync with kDefaultMaxHeadersMsgsPerInterval */,
+          facebook::fb303::AVG,
+          50,
+          99,
           100) {
 }
 
@@ -172,4 +183,12 @@ void TLHTTPSessionStats::recordControlMsgRateLimited() noexcept {
   ctrlMsgsRateLimited.add(1);
 }
 
+void TLHTTPSessionStats::recordHeadersInInterval(int64_t quantity) noexcept {
+  headersInInterval.add(quantity);
+}
+
+void TLHTTPSessionStats::recordHeadersRateLimited() noexcept {
+  headersRateLimited.add(1);
+}
+
 } // namespace proxygen
diff --git a/proxygen/lib/http/stats/ThreadLocalHTTPSessionStats.h b/proxygen/lib/http/stats/ThreadLocalHTTPSessionStats.h
@@ -44,6 +44,8 @@ class TLHTTPSessionStats : public HTTPSessionStats {
 
   void recordControlMsgsInInterval(int64_t quantity) noexcept override;
   void recordControlMsgRateLimited() noexcept override;
+  void recordHeadersInInterval(int64_t quantity) noexcept override;
+  void recordHeadersRateLimited() noexcept override;
 
   BaseStats::TLCounter txnsOpen;
   BaseStats::TLCounter pendingBufferedReadBytes;
@@ -68,9 +70,11 @@ class TLHTTPSessionStats : public HTTPSessionStats {
   BaseStats::TLTimeseries ttbtxNotFound;
   BaseStats::TLTimeseries ttbtxExceedLimit;
   BaseStats::TLTimeseries ctrlMsgsRateLimited;
+  BaseStats::TLTimeseries headersRateLimited;
   BaseStats::TLHistogram txnsPerSession;
   BaseStats::TLHistogram sessionIdleTime;
   BaseStats::TLHistogram ctrlMsgsInInterval;
+  BaseStats::TLHistogram headersInInterval;
 };
 
 } // namespace proxygen