diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs index 75998b2d4c..f088ec27ab 100644 --- a/quinn-proto/src/endpoint.rs +++ b/quinn-proto/src/endpoint.rs @@ -653,6 +653,13 @@ impl Endpoint { } } + /// Rejects new connections without affecting existing connections. + 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 449dd82e4c..42075da7de 100644 --- a/quinn-proto/src/tests/mod.rs +++ b/quinn-proto/src/tests/mod.rs @@ -2193,3 +2193,11 @@ fn stream_chunks(mut recv: RecvStream) -> Vec { buf } + +#[test] +fn reject_new_connections() { + let _guard = subscribe(); + let mut pair = Pair::default(); + pair.server.reject_new_connections(); + pair.assert_connection_refused(); +} diff --git a/quinn-proto/src/tests/util.rs b/quinn-proto/src/tests/util.rs index 97f306d18b..48bc953231 100644 --- a/quinn-proto/src/tests/util.rs +++ b/quinn-proto/src/tests/util.rs @@ -166,6 +166,16 @@ impl Pair { } } + /// Attempts to connect, asserting that the server doesn't accept the + /// connection and the client marks the connection as closed. + pub(super) fn assert_connection_refused(&mut self) { + dbg!(&self.client.connections); + let client_ch = self.begin_connect(client_config()); + self.drive(); + self.server.assert_no_accept(); + assert!(self.client.connections.get(&client_ch).unwrap().is_closed()); + } + pub(super) fn connect(&mut self) -> (ConnectionHandle, ConnectionHandle) { self.connect_with(client_config()) } diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index 03b9015951..5642047feb 100644 --- a/quinn/src/endpoint.rs +++ b/quinn/src/endpoint.rs @@ -236,6 +236,16 @@ impl Endpoint { self.inner.state.lock().unwrap().socket.local_addr() } + /// Rejects new connections without affecting existing connections. + 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.