From 1320d7efab8a29452787266a29ea14d006469c6e Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Fri, 17 Sep 2021 18:30:16 +0900
Subject: [PATCH] nghttpx: Do not accept new connection during graceful
 shutdown period

---
 src/shrpx_http3_upstream.cc          | 20 ++++++++++++++++++--
 src/shrpx_http3_upstream.h           |  1 +
 src/shrpx_quic_connection_handler.cc | 17 +++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/shrpx_http3_upstream.cc b/src/shrpx_http3_upstream.cc
index 9ed1a8cf..b3c509f5 100644
--- a/src/shrpx_http3_upstream.cc
+++ b/src/shrpx_http3_upstream.cc
@@ -124,7 +124,8 @@ Http3Upstream::Http3Upstream(ClientHandler *handler)
       httpconn_{nullptr},
       downstream_queue_{downstream_queue_size(handler->get_worker()),
                         !get_config()->http2_proxy},
-      idle_close_{false} {
+      idle_close_{false},
+      retry_close_{false} {
   ev_timer_init(&timer_, timeoutcb, 0., 0.);
   timer_.data = this;
 
@@ -1338,7 +1339,7 @@ void Http3Upstream::on_handler_delete() {
     quic_conn_handler->remove_connection_id(&cid);
   }
 
-  if (idle_close_) {
+  if (idle_close_ || retry_close_) {
     return;
   }
 
@@ -1624,6 +1625,21 @@ int Http3Upstream::on_read(const UpstreamAddr *faddr,
         return -1;
       }
 
+      if (worker->get_graceful_shutdown()) {
+        ngtcp2_cid ini_dcid, ini_scid;
+
+        ngtcp2_cid_init(&ini_dcid, dcid, dcidlen);
+        ngtcp2_cid_init(&ini_scid, scid, scidlen);
+
+        quic_conn_handler->send_connection_close(
+            faddr, version, &ini_dcid, &ini_scid, remote_addr, local_addr,
+            NGTCP2_CONNECTION_REFUSED);
+
+        return -1;
+      }
+
+      retry_close_ = true;
+
       quic_conn_handler->send_retry(handler_->get_upstream_addr(), version,
                                     dcid, dcidlen, scid, scidlen, remote_addr,
                                     local_addr);
diff --git a/src/shrpx_http3_upstream.h b/src/shrpx_http3_upstream.h
index 0778a3f5..b811a4bb 100644
--- a/src/shrpx_http3_upstream.h
+++ b/src/shrpx_http3_upstream.h
@@ -166,6 +166,7 @@ private:
   nghttp3_conn *httpconn_;
   DownstreamQueue downstream_queue_;
   bool idle_close_;
+  bool retry_close_;
   std::vector<uint8_t> conn_close_;
 };
 
diff --git a/src/shrpx_quic_connection_handler.cc b/src/shrpx_quic_connection_handler.cc
index 85be3762..28b6c17c 100644
--- a/src/shrpx_quic_connection_handler.cc
+++ b/src/shrpx_quic_connection_handler.cc
@@ -180,6 +180,12 @@ int QUICConnectionHandler::handle_packet(const UpstreamAddr *faddr,
         return 0;
       }
 
+      if (worker_->get_graceful_shutdown()) {
+        send_connection_close(faddr, version, &hd.dcid, &hd.scid, remote_addr,
+                              local_addr, NGTCP2_CONNECTION_REFUSED);
+        return 0;
+      }
+
       if (hd.token.len == 0) {
         if (quicconf.upstream.require_token) {
           send_retry(faddr, version, dcid, dcidlen, scid, scidlen, remote_addr,
@@ -262,6 +268,12 @@ int QUICConnectionHandler::handle_packet(const UpstreamAddr *faddr,
       break;
     }
     case NGTCP2_ERR_RETRY:
+      if (worker_->get_graceful_shutdown()) {
+        send_connection_close(faddr, version, &hd.dcid, &hd.scid, remote_addr,
+                              local_addr, NGTCP2_CONNECTION_REFUSED);
+        return 0;
+      }
+
       send_retry(faddr, version, dcid, dcidlen, scid, scidlen, remote_addr,
                  local_addr);
       return 0;
@@ -456,6 +468,11 @@ int QUICConnectionHandler::send_retry(
   auto d =
       static_cast<ev_tstamp>(NGTCP2_DEFAULT_INITIAL_RTT * 3) / NGTCP2_SECONDS;
 
+  if (LOG_ENABLED(INFO)) {
+    LOG(INFO) << "Enter close-wait period " << d << "s with " << buf.size()
+              << " bytes sentinel packet";
+  }
+
   auto cw = std::make_unique<CloseWait>(worker_, std::vector<ngtcp2_cid>{idcid},
                                         std::move(buf), d);