contrib: add per-websocket-connection frame rate limit filter

CODEOWNERS

@@ -514,5 +514,6 @@
 /contrib/kae/ @Misakokoro @UNOWNED
 /contrib/istio @kyessenov @wbpcode @keithmattix @krinkinmu @zirain
 /contrib/reverse_tunnel_reporter @agrawroh @aakugan @basundhara-c
+/contrib/ws_local_ratelimit/ @Amila-Rukshan @UNOWNED
 
 /compat/openssl/ @tedjpoole @envoyproxy/envoy-openssl-sync

api/BUILD

@@ -84,6 +84,7 @@ proto_library(
         "//contrib/envoy/extensions/filters/http/language/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/http/peer_metadata/v3:pkg",
         "//contrib/envoy/extensions/filters/http/sxg/v3alpha:pkg",
+        "//contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/listener/postgres_inspector/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/network/client_ssl_auth/v3:pkg",
         "//contrib/envoy/extensions/filters/network/generic_proxy/codecs/kafka/v3:pkg",

api/contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha/BUILD

@@ -0,0 +1,13 @@
+# DO NOT EDIT. This file is generated by tools/proto_format/proto_sync.py.
+
+load("@envoy_api//bazel:api_build_system.bzl", "api_proto_package")
+
+licenses(["notice"])  # Apache 2
+
+api_proto_package(
+    deps = [
+        "//envoy/type/v3:pkg",
+        "@xds//udpa/annotations:pkg",
+        "@xds//xds/annotations/v3:pkg",
+    ],
+)

api/contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha/ws_local_ratelimit.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package envoy.extensions.filters.http.ws_local_ratelimit.v3alpha;
+
+import "envoy/type/v3/token_bucket.proto";
+
+import "xds/annotations/v3/status.proto";
+
+import "udpa/annotations/status.proto";
+import "validate/validate.proto";
+
+option java_package = "io.envoyproxy.envoy.extensions.filters.http.ws_local_ratelimit.v3alpha";
+option java_outer_classname = "WsLocalRatelimitProto";
+option java_multiple_files = true;
+option go_package = "github.com/envoyproxy/go-control-plane/contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha";
+option (udpa.annotations.file_status).package_version_status = ACTIVE;
+option (xds.annotations.v3.file_status).work_in_progress = true;
+
+// [#protodoc-title: WebSocket local rate limit HTTP filter]
+//
+// Rate limits WebSocket data frames (text, binary, continuation) on a per-WebSocket-session
+// basis. Unlike the generic local rate limit filter, this filter has no descriptor matching and
+// no shared/global bucket: the ``token_bucket`` configured here is instantiated fresh for every
+// WebSocket session, so each session is always rate limited independently - including multiple
+// sessions multiplexed as separate streams over one HTTP/2 or HTTP/3 downstream connection.
+// [#extension: envoy.filters.http.ws_local_ratelimit]
+
+message WsLocalRateLimit {
+  // The prefix to use when emitting statistics.
+  string stat_prefix = 1 [(validate.rules).string = {min_len: 1}];
+
+  // The token bucket applied to each WebSocket session. A new bucket instance is created per
+  // session; there is no bucket shared across sessions, even when multiple sessions are
+  // multiplexed over the same downstream connection.
+  type.v3.TokenBucket token_bucket = 2 [(validate.rules).message = {required: true}];
+
+  // Optional message sent back to the downstream client, in a WebSocket text frame, whenever a
+  // data frame is rejected by the rate limiter. If unset (the default), the rejected frame is
+  // simply dropped and no response is sent to the client.
+  string rejection_message = 3;
+}

api/versioning/BUILD

@@ -22,6 +22,7 @@ proto_library(
         "//contrib/envoy/extensions/filters/http/peak_ewma/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/http/peer_metadata/v3:pkg",
         "//contrib/envoy/extensions/filters/http/sxg/v3alpha:pkg",
+        "//contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/listener/postgres_inspector/v3alpha:pkg",
         "//contrib/envoy/extensions/filters/network/client_ssl_auth/v3:pkg",
         "//contrib/envoy/extensions/filters/network/generic_proxy/codecs/kafka/v3:pkg",

contrib/contrib_build_config.bzl

@@ -16,6 +16,7 @@ CONTRIB_EXTENSIONS = {
     "envoy.filters.http.language":                              "//contrib/language/filters/http/source:config_lib",
     "envoy.filters.http.peak_ewma":                             "//contrib/peak_ewma/filters/http/source:config",
     "envoy.filters.http.sxg":                                   "//contrib/sxg/filters/http/source:config",
+    "envoy.filters.http.ws_local_ratelimit":                    "//contrib/ws_local_ratelimit/filters/http/source:config",
     "envoy.filters.http.peer_metadata":                         "//contrib/istio/filters/http/peer_metadata/source:config",
     "envoy.filters.http.istio_stats":                           "//contrib/istio/filters/http/istio_stats/source:istio_stats",
     "envoy.filters.http.alpn":                                  "//contrib/istio/filters/http/alpn/source:config_lib",

contrib/extensions_metadata.yaml

@@ -206,6 +206,13 @@ envoy.filters.http.alpn:
   status: alpha
   type_urls:
   - istio.envoy.config.filter.http.alpn.v2alpha1.FilterConfig
+envoy.filters.http.ws_local_ratelimit:
+  categories:
+  - envoy.filters.http
+  security_posture: robust_to_untrusted_downstream
+  status: alpha
+  type_urls:
+  - envoy.extensions.filters.http.ws_local_ratelimit.v3alpha.WsLocalRateLimit
 envoy.load_balancing_policies.peak_ewma:
   categories:
   - envoy.load_balancing_policies

contrib/ws_local_ratelimit/filters/http/source/BUILD

@@ -0,0 +1,41 @@
+load(
+    "//bazel:envoy_build_system.bzl",
+    "envoy_cc_contrib_extension",
+    "envoy_cc_library",
+    "envoy_contrib_package",
+)
+
+licenses(["notice"])  # Apache 2
+
+envoy_contrib_package()
+
+# WebSocket local rate limit L7 HTTP filter
+
+envoy_cc_library(
+    name = "ws_local_ratelimit_filter_lib",
+    srcs = ["ws_local_ratelimit_filter.cc"],
+    hdrs = ["ws_local_ratelimit_filter.h"],
+    deps = [
+        "//envoy/http:filter_interface",
+        "//envoy/stream_info:filter_state_interface",
+        "//source/common/buffer:buffer_lib",
+        "//source/common/common:macros",
+        "//source/common/http:utility_lib",
+        "//source/common/websocket:codec_lib",
+        "//source/extensions/filters/common/local_ratelimit:local_ratelimit_lib",
+        "//source/extensions/filters/http/common:pass_through_filter_lib",
+        "@envoy_api//contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha:pkg_cc_proto",
+    ],
+)
+
+envoy_cc_contrib_extension(
+    name = "config",
+    srcs = ["config.cc"],
+    hdrs = ["config.h"],
+    deps = [
+        ":ws_local_ratelimit_filter_lib",
+        "//envoy/registry",
+        "//source/extensions/filters/http/common:factory_base_lib",
+        "@envoy_api//contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha:pkg_cc_proto",
+    ],
+)

contrib/ws_local_ratelimit/filters/http/source/config.cc

@@ -0,0 +1,36 @@
+#include "contrib/ws_local_ratelimit/filters/http/source/config.h"
+
+#include <string>
+
+#include "envoy/registry/registry.h"
+
+#include "contrib/ws_local_ratelimit/filters/http/source/ws_local_ratelimit_filter.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace HttpFilters {
+namespace WsLocalRateLimitFilter {
+
+absl::StatusOr<Http::FilterFactoryCb>
+WsLocalRateLimitFilterFactory::createHttpFilterFactoryFromProtoTyped(
+    const envoy::extensions::filters::http::ws_local_ratelimit::v3alpha::WsLocalRateLimit&
+        proto_config,
+    Server::Configuration::ServerFactoryContext& context,
+    Server::Configuration::ExtraFactoryContext& extra_context) {
+  WsLocalRateLimitConfigSharedPtr filter_config(
+      new WsLocalRateLimitConfig(proto_config, extra_context.scopeOr(context)));
+  return [filter_config](Http::FilterChainFactoryCallbacks& callbacks) -> void {
+    callbacks.addStreamFilter(std::make_shared<WsLocalRateLimitFilter>(filter_config));
+  };
+}
+
+/**
+ * Static registration for the WebSocket local rate limit filter. @see RegisterFactory.
+ */
+REGISTER_FACTORY(WsLocalRateLimitFilterFactory,
+                 Server::Configuration::NamedHttpFilterConfigFactory);
+
+} // namespace WsLocalRateLimitFilter
+} // namespace HttpFilters
+} // namespace Extensions
+} // namespace Envoy

contrib/ws_local_ratelimit/filters/http/source/config.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "source/extensions/filters/http/common/factory_base.h"
+
+#include "contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha/ws_local_ratelimit.pb.h"
+#include "contrib/envoy/extensions/filters/http/ws_local_ratelimit/v3alpha/ws_local_ratelimit.pb.validate.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace HttpFilters {
+namespace WsLocalRateLimitFilter {
+
+/**
+ * Config registration for the WebSocket local rate limit filter.
+ */
+class WsLocalRateLimitFilterFactory
+    : public Common::UnifiedFactoryBase<
+          envoy::extensions::filters::http::ws_local_ratelimit::v3alpha::WsLocalRateLimit> {
+public:
+  WsLocalRateLimitFilterFactory() : UnifiedFactoryBase("envoy.filters.http.ws_local_ratelimit") {}
+
+private:
+  absl::StatusOr<Http::FilterFactoryCb> createHttpFilterFactoryFromProtoTyped(
+      const envoy::extensions::filters::http::ws_local_ratelimit::v3alpha::WsLocalRateLimit&
+          proto_config,
+      Server::Configuration::ServerFactoryContext& context,
+      Server::Configuration::ExtraFactoryContext& extra_context) override;
+};
+
+DECLARE_FACTORY(WsLocalRateLimitFilterFactory);
+
+} // namespace WsLocalRateLimitFilter
+} // namespace HttpFilters
+} // namespace Extensions
+} // namespace Envoy

contrib/ws_local_ratelimit/filters/http/source/ws_local_ratelimit_filter.cc

@@ -0,0 +1,169 @@
+#include "contrib/ws_local_ratelimit/filters/http/source/ws_local_ratelimit_filter.h"
+
+#include <optional>
+
+#include "envoy/common/exception.h"
+#include "envoy/stats/stats.h"
+
+#include "source/common/buffer/buffer_impl.h"
+#include "source/common/common/macros.h"
+#include "source/common/http/utility.h"
+
+namespace Envoy {
+namespace Extensions {
+namespace HttpFilters {
+namespace WsLocalRateLimitFilter {
+
+WsLocalRateLimitConfig::WsLocalRateLimitConfig(
+    const envoy::extensions::filters::http::ws_local_ratelimit::v3alpha::WsLocalRateLimit& config,
+    Stats::Scope& scope)
+    : fill_interval_(std::chrono::milliseconds(
+          PROTOBUF_GET_MS_OR_DEFAULT(config.token_bucket(), fill_interval, 0))),
+      max_tokens_(config.token_bucket().max_tokens()),
+      tokens_per_fill_(PROTOBUF_GET_WRAPPED_OR_DEFAULT(config.token_bucket(), tokens_per_fill, 1)),
+      rejection_message_(config.rejection_message()),
+      encoded_rejection_frame_(encodeRejectionFrame(rejection_message_)),
+      stats_(generateStats(config.stat_prefix(), scope)) {
+  // Mirror the validation performed by LocalRateLimiterImpl so that an invalid fill interval is
+  // rejected at config-load time instead of throwing the first time a per-connection rate
+  // limiter is lazily constructed on the request path.
+  if (fill_interval_.count() > 0 && max_tokens_ != 0 &&
+      fill_interval_ < std::chrono::milliseconds(50)) {
+    throw EnvoyException("local rate limit token bucket fill timer must be >= 50ms");
+  }
+}
+
+WsLocalRateLimitStats WsLocalRateLimitConfig::generateStats(const std::string& prefix,
+                                                            Stats::Scope& scope) {
+  const std::string final_prefix = prefix + ".ws_local_rate_limit";
+  return {ALL_WS_LOCAL_RATE_LIMIT_STATS(POOL_COUNTER_PREFIX(scope, final_prefix))};
+}
+
+std::optional<std::string>
+WsLocalRateLimitConfig::encodeRejectionFrame(const std::string& rejection_message) {
+  if (rejection_message.empty()) {
+    return std::nullopt;
+  }
+
+  WebSocket::Frame frame;
+  frame.final_fragment_ = true;
+  frame.opcode_ = WebSocket::kFrameOpcodeText;
+  // Frames sent from the server to the client must not be masked.
+  frame.masking_key_ = std::nullopt;
+  frame.payload_length_ = rejection_message.size();
+
+  WebSocket::Encoder encoder;
+  std::optional<std::vector<uint8_t>> header = encoder.encodeFrameHeader(frame);
+  if (!header.has_value()) {
+    // Can't happen in practice (kFrameOpcodeText is always a valid opcode).
+    return std::nullopt;
+  }
+
+  std::string encoded(header->begin(), header->end());
+  encoded += rejection_message;
+  return encoded;
+}
+
+const std::string& PerStreamRateLimiter::key() {
+  CONSTRUCT_ON_FIRST_USE(std::string, "ws_local_ratelimit.per_stream_rate_limiter");
+}
+
+Http::FilterHeadersStatus WsLocalRateLimitFilter::decodeHeaders(Http::RequestHeaderMap& headers,
+                                                                bool) {
+  is_websocket_ = Http::Utility::isWebSocketUpgradeRequest(headers);
+  return Http::FilterHeadersStatus::Continue;
+}
+
+Http::FilterDataStatus WsLocalRateLimitFilter::decodeData(Buffer::Instance& data, bool) {
+  if (!is_websocket_) {
+    return Http::FilterDataStatus::Continue;
+  }
+
+  std::optional<std::vector<WebSocket::Frame>> frames = decoder_.decode(data);
+  data.drain(data.length());
+
+  if (frames.has_value()) {
+    auto& limiter = getPerStreamRateLimiter();
+    for (const auto& frame : frames.value()) {
+      if (frame.opcode_ == WebSocket::kFrameOpcodeText ||
+          frame.opcode_ == WebSocket::kFrameOpcodeBinary ||
+          frame.opcode_ == WebSocket::kFrameOpcodeContinuation) {
+        if (!limiter.requestAllowed({}).allowed) {
+          config_->stats().rate_limited_.inc();
+          if (!config_->rejectionMessage().empty()) {
+            ENVOY_LOG(debug, "WebSocket message rate limit exceeded, notifying downstream client");
+            sendRejectionFrame();
+          } else {
+            ENVOY_LOG(debug, "WebSocket message rate limit exceeded, dropping frame");
+          }
+          continue;
+        }
+        config_->stats().ok_.inc();
+      }
+
+      // encode the frame back to the outgoing data buffer, forwarding it upstream.
+      std::optional<std::vector<uint8_t>> encoded_frame_header = encoder_.encodeFrameHeader(frame);
+      if (encoded_frame_header.has_value()) {
+        data.add(encoded_frame_header->data(), encoded_frame_header->size());
+        if (frame.payload_ != nullptr) {
+          data.add(*frame.payload_);
+        }
+      }
+    }
+  }
+
+  return Http::FilterDataStatus::Continue;
+}
+
+Filters::Common::LocalRateLimit::LocalRateLimiterImpl&
+WsLocalRateLimitFilter::getPerStreamRateLimiter() {
+  auto typed_state =
+      decoder_callbacks_->streamInfo().filterState()->getDataReadOnly<PerStreamRateLimiter>(
+          PerStreamRateLimiter::key());
+
+  if (typed_state == nullptr) {
+    auto limiter = std::make_shared<PerStreamRateLimiter>(
+        config_->fillInterval(), config_->maxTokens(), config_->tokensPerFill(),
+        decoder_callbacks_->dispatcher());
+
+    // Request-level (not Connection-level) life span: on HTTP/2 and HTTP/3, multiple
+    // independent WebSocket sessions can be multiplexed over one downstream connection, and
+    // each must get its own bucket rather than sharing one across the whole connection.
+    decoder_callbacks_->streamInfo().filterState()->setData(
+        PerStreamRateLimiter::key(), limiter, StreamInfo::FilterState::LifeSpan::Request);
+    return limiter->value();
+  }
+
+  return const_cast<PerStreamRateLimiter&>(*typed_state).value();
+}
+
+void WsLocalRateLimitFilter::sendRejectionFrame() {
+  // encodeData() below injects data into the encode/response filter chain. That's only valid
+  // once response headers exist for this stream (e.g. the upstream's WebSocket upgrade response
+  // has already been received and forwarded) - calling it before that violates FilterManager's
+  // encode-iteration invariants (ASSERT(headers_continued_) in commonHandleAfterDataCallback)
+  // and crashes. This can happen for real: a client bursting frames fast enough that the budget
+  // is exceeded before the upstream's initial response headers arrive. Drop the notification in
+  // that case rather than risk it - the offending frame itself is still dropped either way.
+  if (!decoder_callbacks_->responseHeaders().has_value()) {
+    ENVOY_LOG(debug, "Cannot notify downstream of WebSocket rate limit rejection: response "
+                     "headers not sent yet");
+    return;
+  }
+
+  // The rejection frame's bytes never change once the filter is configured, so they're computed
+  // once in WsLocalRateLimitConfig rather than re-encoded on every single rejected frame.
+  const std::optional<std::string>& encoded_frame = config_->encodedRejectionFrame();
+  if (!encoded_frame.has_value()) {
+    return;
+  }
+  Buffer::OwnedImpl reply_buffer(encoded_frame->data(), encoded_frame->size());
+
+  // Send the frame directly to the downstream client on the encode path, bypassing upstream.
+  decoder_callbacks_->encodeData(reply_buffer, false);
+}
+
+} // namespace WsLocalRateLimitFilter
+} // namespace HttpFilters
+} // namespace Extensions
+} // namespace Envoy