From 85347e12deb8bfde93aab955073955b986eedf3c Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Sun, 5 Sep 2021 19:23:50 +0900
Subject: [PATCH] nghttpx: Rate limit Stateless Reset transmission

---
 src/shrpx_quic.h                     |  1 +
 src/shrpx_quic_connection_handler.cc | 44 ++++++++++++++++++++++++++--
 src/shrpx_quic_connection_handler.h  |  4 +++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/shrpx_quic.h b/src/shrpx_quic.h
index 66254a91..39480446 100644
--- a/src/shrpx_quic.h
+++ b/src/shrpx_quic.h
@@ -64,6 +64,7 @@ constexpr size_t SHRPX_QUIC_STATELESS_RESET_SECRETLEN = 32;
 constexpr size_t SHRPX_QUIC_TOKEN_SECRETLEN = 32;
 constexpr size_t SHRPX_QUIC_TOKEN_RAND_DATALEN = 16;
 constexpr size_t SHRPX_QUIC_CONN_CLOSE_PKTLEN = 256;
+constexpr size_t SHRPX_QUIC_STATELESS_RESET_BURST = 100;
 
 // SHRPX_QUIC_RETRY_TOKEN_MAGIC is the magic byte of Retry token.
 // Sent in plaintext.
diff --git a/src/shrpx_quic_connection_handler.cc b/src/shrpx_quic_connection_handler.cc
index 1b21513d..05d85f28 100644
--- a/src/shrpx_quic_connection_handler.cc
+++ b/src/shrpx_quic_connection_handler.cc
@@ -37,10 +37,26 @@
 
 namespace shrpx {
 
-QUICConnectionHandler::QUICConnectionHandler(Worker *worker)
-    : worker_{worker} {}
+namespace {
+void stateless_reset_bucket_regen_timercb(struct ev_loop *loop, ev_timer *w,
+                                          int revents) {
+  auto quic_conn_handler = static_cast<QUICConnectionHandler *>(w->data);
 
-QUICConnectionHandler::~QUICConnectionHandler() {}
+  quic_conn_handler->on_stateless_reset_bucket_regen();
+}
+} // namespace
+
+QUICConnectionHandler::QUICConnectionHandler(Worker *worker)
+    : worker_{worker},
+      stateless_reset_bucket_{SHRPX_QUIC_STATELESS_RESET_BURST} {
+  ev_timer_init(&stateless_reset_bucket_regen_timer_,
+                stateless_reset_bucket_regen_timercb, 0., 1.);
+  stateless_reset_bucket_regen_timer_.data = this;
+}
+
+QUICConnectionHandler::~QUICConnectionHandler() {
+  ev_timer_stop(worker_->get_loop(), &stateless_reset_bucket_regen_timer_);
+}
 
 int QUICConnectionHandler::handle_packet(const UpstreamAddr *faddr,
                                          const Address &remote_addr,
@@ -378,6 +394,20 @@ int QUICConnectionHandler::send_stateless_reset(const UpstreamAddr *faddr,
                                                 size_t dcidlen,
                                                 const Address &remote_addr,
                                                 const Address &local_addr) {
+  if (stateless_reset_bucket_ == 0) {
+    if (LOG_ENABLED(INFO)) {
+      LOG(INFO) << "Stateless Reset bucket has been depleted";
+    }
+
+    return 0;
+  }
+
+  --stateless_reset_bucket_;
+
+  if (!ev_is_active(&stateless_reset_bucket_regen_timer_)) {
+    ev_timer_again(worker_->get_loop(), &stateless_reset_bucket_regen_timer_);
+  }
+
   int rv;
   std::array<uint8_t, NGTCP2_STATELESS_RESET_TOKENLEN> token;
   ngtcp2_cid cid;
@@ -469,6 +499,14 @@ void QUICConnectionHandler::remove_close_wait(const CloseWait *cw) {
   }
 }
 
+void QUICConnectionHandler::on_stateless_reset_bucket_regen() {
+  assert(stateless_reset_bucket_ < SHRPX_QUIC_STATELESS_RESET_BURST);
+
+  if (++stateless_reset_bucket_ == SHRPX_QUIC_STATELESS_RESET_BURST) {
+    ev_timer_stop(worker_->get_loop(), &stateless_reset_bucket_regen_timer_);
+  }
+}
+
 static void close_wait_timeoutcb(struct ev_loop *loop, ev_timer *w,
                                  int revents) {
   auto cw = static_cast<CloseWait *>(w->data);
diff --git a/src/shrpx_quic_connection_handler.h b/src/shrpx_quic_connection_handler.h
index 84530173..d6d46f05 100644
--- a/src/shrpx_quic_connection_handler.h
+++ b/src/shrpx_quic_connection_handler.h
@@ -126,10 +126,14 @@ public:
   void add_close_wait(CloseWait *cw);
   void remove_close_wait(const CloseWait *cw);
 
+  void on_stateless_reset_bucket_regen();
+
 private:
   Worker *worker_;
   std::unordered_map<ngtcp2_cid, ClientHandler *> connections_;
   std::unordered_map<ngtcp2_cid, CloseWait *> close_waits_;
+  ev_timer stateless_reset_bucket_regen_timer_;
+  size_t stateless_reset_bucket_;
 };
 
 } // namespace shrpx