From d73fcfd4a05684433bc3df2bf358a00b3676fbc0 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sun, 13 Nov 2022 10:25:59 -0800
Subject: [PATCH] Fix stream open futures sometimes waiting on the wrong
 notification

After a spurious wakeup, `Connection::open_uni` and `open_bi` would
switch to waiting for notifications about incoming streams, rather
than the intended notification about new stream ID flow control
budget. This was probably a copy-paste error from `poll_accept`.
---
 quinn/src/connection.rs |  5 ++++-
 quinn/src/tests.rs      | 45 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs
index 10fd15693..2d38e21bb 100644
--- a/quinn/src/connection.rs
+++ b/quinn/src/connection.rs
@@ -628,7 +628,7 @@ fn poll_open<'a>(
             // `state` lock ensures we didn't race with readiness
             Poll::Pending => return Poll::Pending,
             // Spurious wakeup, get a new future
-            Poll::Ready(()) => notify.set(conn.shared.stream_incoming[dir as usize].notified()),
+            Poll::Ready(()) => notify.set(conn.shared.stream_opening[dir as usize].notified()),
         }
     }
 }
@@ -813,7 +813,10 @@ pub struct ConnectionInner {
 
 #[derive(Debug, Default)]
 pub(crate) struct Shared {
+    /// Notified when new streams may be locally initiated due to an increase in stream ID flow
+    /// control budget
     stream_opening: [Notify; 2],
+    /// Notified when the peer has initiated a new stream
     stream_incoming: [Notify; 2],
     datagrams: Notify,
     closed: Notify,
diff --git a/quinn/src/tests.rs b/quinn/src/tests.rs
index b79e9d4ef..664dc6d75 100644
--- a/quinn/src/tests.rs
+++ b/quinn/src/tests.rs
@@ -224,10 +224,16 @@ async fn accept_after_close() {
 
 /// Construct an endpoint suitable for connecting to itself
 fn endpoint() -> Endpoint {
+    endpoint_with_config(TransportConfig::default())
+}
+
+fn endpoint_with_config(transport_config: TransportConfig) -> Endpoint {
     let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
     let key = rustls::PrivateKey(cert.serialize_private_key_der());
     let cert = rustls::Certificate(cert.serialize_der().unwrap());
-    let server_config = crate::ServerConfig::with_single_cert(vec![cert.clone()], key).unwrap();
+    let transport_config = Arc::new(transport_config);
+    let mut server_config = crate::ServerConfig::with_single_cert(vec![cert.clone()], key).unwrap();
+    server_config.transport_config(transport_config.clone());
 
     let mut roots = rustls::RootCertStore::empty();
     roots.add(&cert).unwrap();
@@ -236,7 +242,8 @@ fn endpoint() -> Endpoint {
         SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
     )
     .unwrap();
-    let client_config = ClientConfig::with_root_certificates(roots);
+    let mut client_config = ClientConfig::with_root_certificates(roots);
+    client_config.transport_config(transport_config);
     endpoint.set_default_client_config(client_config);
 
     endpoint
@@ -646,3 +653,37 @@ async fn rebind_recv() {
     assert_eq!(stream.read_to_end(MSG.len()).await.unwrap(), MSG);
     server.await.unwrap();
 }
+
+#[tokio::test]
+async fn stream_id_flow_control() {
+    let _guard = subscribe();
+    let mut cfg = TransportConfig::default();
+    cfg.max_concurrent_uni_streams(1u32.into());
+    let endpoint = endpoint_with_config(cfg);
+
+    let (client, server) = tokio::join!(
+        endpoint
+            .connect(endpoint.local_addr().unwrap(), "localhost")
+            .unwrap(),
+        async { endpoint.accept().await.unwrap().await }
+    );
+    let client = client.unwrap();
+    let server = server.unwrap();
+
+    // If `open_uni` doesn't get unblocked when the previous stream is dropped, this will time out.
+    tokio::join!(
+        async {
+            client.open_uni().await.unwrap();
+        },
+        async {
+            client.open_uni().await.unwrap();
+        },
+        async {
+            client.open_uni().await.unwrap();
+        },
+        async {
+            server.accept_uni().await.unwrap();
+            server.accept_uni().await.unwrap();
+        }
+    );
+}