From 98f5fe2a3fabb9ff991f8c831e8d43de76985ff3 Mon Sep 17 00:00:00 2001
From: Jonathan Johnson <jon@khonsulabs.com>
Date: Wed, 7 Jun 2023 01:15:52 -0700
Subject: [PATCH] Added Endpoint::reject_new_connections (#1585)

This commit adds a new function that refuses new connections without
impacting existing connections. Internally, this just sets the
connection limit to 0, which causes incoming connections to be rejected.
This is the same approach that was taken in 0.8.5 when `Incoming` was
dropped.
---
 quinn-proto/src/endpoint.rs  | 12 ++++++++++++
 quinn-proto/src/tests/mod.rs | 13 +++++++++++++
 quinn/src/endpoint.rs        | 15 +++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index 75998b2d4..79e91871a 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -653,6 +653,18 @@ impl Endpoint {
         }
     }
 
+    /// Reject new incoming connections without affecting existing connections
+    ///
+    /// Convenience short-hand for using
+    /// [`set_server_config`](Self::set_server_config) to update
+    /// [`concurrent_connections`](ServerConfig::concurrent_connections) to
+    /// zero.
+    pub fn reject_new_connections(&mut self) {
+        if let Some(config) = self.server_config.as_mut() {
+            Arc::make_mut(config).concurrent_connections(0);
+        }
+    }
+
     /// Access the configuration used by this endpoint
     pub fn config(&self) -> &EndpointConfig {
         &self.config
diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
index 449dd82e4..24a7c6cb2 100644
--- a/quinn-proto/src/tests/mod.rs
+++ b/quinn-proto/src/tests/mod.rs
@@ -2193,3 +2193,16 @@ fn stream_chunks(mut recv: RecvStream) -> Vec<u8> {
 
     buf
 }
+
+#[test]
+fn reject_new_connections() {
+    let _guard = subscribe();
+    let mut pair = Pair::default();
+    pair.server.reject_new_connections();
+
+    // The server should now reject incoming connections.
+    let client_ch = pair.begin_connect(client_config());
+    pair.drive();
+    pair.server.assert_no_accept();
+    assert!(pair.client.connections.get(&client_ch).unwrap().is_closed());
+}
diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs
index 03b901595..274a501bf 100644
--- a/quinn/src/endpoint.rs
+++ b/quinn/src/endpoint.rs
@@ -236,6 +236,21 @@ impl Endpoint {
         self.inner.state.lock().unwrap().socket.local_addr()
     }
 
+    /// Reject new incoming connections without affecting existing connections
+    ///
+    /// Convenience short-hand for using
+    /// [`set_server_config`](Self::set_server_config) to update
+    /// [`concurrent_connections`](ServerConfig::concurrent_connections) to
+    /// zero.
+    pub fn reject_new_connections(&self) {
+        self.inner
+            .state
+            .lock()
+            .unwrap()
+            .inner
+            .reject_new_connections();
+    }
+
     /// Close all of this endpoint's connections immediately and cease accepting new connections.
     ///
     /// See [`Connection::close()`] for details.