From fe362576bceacd27e992b0b318d24977fae7869d Mon Sep 17 00:00:00 2001
From: Miguel Guarniz <mi9uel9@gmail.com>
Date: Mon, 15 Aug 2022 12:58:41 -0400
Subject: [PATCH] Prevent multiple QUIC connections to the same host

Signed-off-by: Miguel Guarniz <mi9uel9@gmail.com>
---
 src/async_impl/h3_client/mod.rs  |  1 +
 src/async_impl/h3_client/pool.rs | 18 ++++++++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/async_impl/h3_client/mod.rs b/src/async_impl/h3_client/mod.rs
index 59e3a7667..919e13c0a 100644
--- a/src/async_impl/h3_client/mod.rs
+++ b/src/async_impl/h3_client/mod.rs
@@ -40,6 +40,7 @@ impl H3Client {
         trace!("did not find connection {:?} in pool so connecting...", key);
 
         let dest = pool::domain_as_uri(key.clone());
+        self.pool.connecting(key.clone())?;
         let (driver, tx) = self.connector.connect(dest).await?;
         Ok(self.pool.new_connection(key, driver, tx))
     }
diff --git a/src/async_impl/h3_client/pool.rs b/src/async_impl/h3_client/pool.rs
index c31102aab..6fcb8e719 100644
--- a/src/async_impl/h3_client/pool.rs
+++ b/src/async_impl/h3_client/pool.rs
@@ -1,5 +1,5 @@
 use bytes::Bytes;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::mpsc::{Receiver, TryRecvError};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
@@ -27,12 +27,21 @@ impl Pool {
     pub fn new(timeout: Option<Duration>) -> Self {
         Self {
             inner: Arc::new(Mutex::new(PoolInner {
+                connecting: HashSet::new(),
                 idle_conns: HashMap::new(),
                 timeout,
             })),
         }
     }
 
+    pub fn connecting(&self, key: Key) -> Result<(), BoxError> {
+        let mut inner = self.inner.lock().unwrap();
+        if !inner.connecting.insert(key.clone()) {
+            return Err(format!("HTTP/3 connecting already in progress for {:?}", key).into());
+        }
+        return Ok(());
+    }
+
     pub fn try_pool(&self, key: &Key) -> Option<PoolClient> {
         let mut inner = self.inner.lock().unwrap();
         let timeout = inner.timeout;
@@ -77,13 +86,18 @@ impl Pool {
 
         let client = PoolClient::new(tx);
         let conn = PoolConnection::new(client.clone(), close_rx);
-        inner.insert(key, conn);
+        inner.insert(key.clone(), conn);
+
+        // We clean up "connecting" here so we don't have to acquire the lock again.
+        let existed = inner.connecting.remove(&key);
+        debug_assert!(existed, "key not in connecting set");
 
         client
     }
 }
 
 struct PoolInner {
+    connecting: HashSet<Key>,
     idle_conns: HashMap<Key, PoolConnection>,
     timeout: Option<Duration>,
 }