Bound federation handshakes, and finish abandoned .well-known lookups

A shared .well-known lookup only ran while a caller awaited it. If every
caller gave up, the next request resumed it with its fetch timeout long
expired, and cached a failure for an hour. Run it as its own task.

The TLS handshake and a proxy's CONNECT answer had no timeout. hyper
finishes some connections in a background task that no request timeout
reaches, so a peer that stalled after accepting leaked the task and its
socket. Give each 30 seconds.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01R5b7Ruy4cgt27Mjq6Ptxbb

rust/src/federation/agent/connect.rs

@@ -55,7 +55,10 @@ use tokio::{
     io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
     net::TcpStream,
 };
-use tokio_rustls::TlsConnector;
+use tokio_rustls::{
+    rustls::{pki_types::ServerName, ClientConfig},
+    TlsConnector,
+};
 
 use super::{
     policy::{IpPolicy, Proxy},
@@ -75,6 +78,12 @@ pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
 /// Twisted's `HostnameEndpoint` does.
 const ATTEMPT_DELAY: Duration = Duration::from_millis(300);
 
+/// How long the TLS handshake, and a proxy's answer to `CONNECT`, may each
+/// take. Twisted bounds neither, relying on the request's own timeout. We
+/// need a bound here too: hyper finishes some connections in a background
+/// task that no request timeout reaches.
+pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
+
 /// The longest `CONNECT` response head we read from a proxy.
 const MAX_PROXY_RESPONSE: usize = 16 * 1024;
 
@@ -144,14 +153,13 @@ impl ConnectorInner {
         let mut first_error = None;
         for server in servers {
             match self.connect_to(&server).await {
-                Ok(tcp) => {
-                    let connector = TlsConnector::from(tls_config.clone());
-                    return match connector.connect(tls_name.clone(), tcp).await {
-                        Ok(tls) => Ok(FederationStream::new(Box::new(tls))),
-                        Err(err) => Err(TransportError::Tls(
-                            anyhow!(err).context(format!("TLS handshake with {host}")),
-                        )),
-                    };
+                Ok(stream) => {
+                    let tls = tls_handshake(tls_config.clone(), tls_name.clone(), stream)
+                        .await
+                        .map_err(|err| {
+                            TransportError::Tls(err.context(format!("TLS handshake with {host}")))
+                        })?;
+                    return Ok(FederationStream::new(tls));
                 }
                 Err(err) => {
                     info!(
@@ -216,18 +224,26 @@ impl ConnectorInner {
 
         let mut stream: Box<dyn Io> = if proxy.tls {
             let name = server_name(&proxy.host).map_err(TransportError::Tls)?;
-            let connector = TlsConnector::from(self.tls.config_for(&proxy.host));
-            let tls = connector.connect(name, tcp).await.map_err(|err| {
-                TransportError::Tls(anyhow!(err).context(format!("TLS to proxy {}", proxy.host)))
-            })?;
-            Box::new(tls)
+            tls_handshake(self.tls.config_for(&proxy.host), name, Box::new(tcp))
+                .await
+                .map_err(|err| {
+                    TransportError::Tls(err.context(format!("TLS to proxy {}", proxy.host)))
+                })?
         } else {
             Box::new(tcp)
         };
 
-        proxy_connect(&mut stream, host, port, proxy.authorization.as_deref())
-            .await
-            .map_err(TransportError::Connect)?;
+        tokio::time::timeout(
+            HANDSHAKE_TIMEOUT,
+            proxy_connect(&mut stream, host, port, proxy.authorization.as_deref()),
+        )
+        .await
+        .unwrap_or_else(|_| {
+            Err(anyhow!(
+                "proxy did not answer CONNECT in {HANDSHAKE_TIMEOUT:?}"
+            ))
+        })
+        .map_err(TransportError::Connect)?;
 
         Ok(stream)
     }
@@ -289,6 +305,20 @@ async fn connect_tcp(ips: Vec<IpAddr>, port: u16) -> Result<TcpStream, Transport
     }
 }
 
+/// Speak TLS over `stream`, verifying the server as `name`.
+async fn tls_handshake(
+    config: Arc<ClientConfig>,
+    name: ServerName<'static>,
+    stream: Box<dyn Io>,
+) -> anyhow::Result<Box<dyn Io>> {
+    let connector = TlsConnector::from(config);
+    match tokio::time::timeout(HANDSHAKE_TIMEOUT, connector.connect(name, stream)).await {
+        Ok(Ok(tls)) => Ok(Box::new(tls)),
+        Ok(Err(err)) => Err(anyhow!(err)),
+        Err(_) => bail!("timed out after {HANDSHAKE_TIMEOUT:?}"),
+    }
+}
+
 /// Ask an HTTP proxy for a tunnel to `host:port`.
 async fn proxy_connect<S: AsyncRead + AsyncWrite + Unpin + ?Sized>(
     stream: &mut S,
@@ -444,6 +474,30 @@ mod tests {
         assert!(err.to_string().contains("403"), "{err}");
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn stalled_tls_handshake_times_out() {
+        // The peer never answers the ClientHello.
+        let (client, _server) = duplex(4096);
+        let config = TlsPolicy::new(&crate::config::TlsConfig {
+            federation_verify_certificates: true,
+            federation_client_minimum_tls_version: "1.2".to_owned(),
+            federation_certificate_verification_whitelist: Vec::new(),
+            federation_custom_ca_pems: None,
+        })
+        .unwrap()
+        .config_for("remote.example");
+
+        let err = tls_handshake(
+            config,
+            server_name("remote.example").unwrap(),
+            Box::new(client),
+        )
+        .await
+        .err()
+        .expect("handshake should fail");
+        assert!(err.to_string().contains("timed out"), "{err}");
+    }
+
     #[tokio::test]
     async fn tcp_tries_each_address() {
         // Nothing listens on port 1 of the first address, so the second wins.

rust/src/federation/agent/well_known.rs

@@ -17,7 +17,9 @@
 //!
 //! The caching rules are Python's, constant for constant. Two things are
 //! new: each fetch attempt has a timeout (Python's has none), and
-//! concurrent lookups of one server name share a single fetch.
+//! concurrent lookups of one server name share a single fetch. That fetch
+//! runs as its own task, so it finishes and caches its result even if every
+//! request waiting on it gives up.
 
 use std::{
     collections::HashMap,
@@ -184,14 +186,22 @@ impl WellKnownResolver {
                 .or_insert_with(|| {
                     let this = self.clone();
                     let server_name = server_name.to_owned();
-                    async move {
-                        let result = this.lookup(&server_name, prev).await;
-                        this.in_flight
-                            .lock()
-                            .expect("poisoned")
-                            .remove(&server_name);
-                        result
-                    }
+                    // A future only its callers poll would stall when they
+                    // all go away, and resume on the next request with its
+                    // fetch timeout long expired.
+                    let task = tokio::spawn(async move {
+                        let _done = InFlightGuard {
+                            resolver: &this,
+                            server_name: &server_name,
+                        };
+                        this.lookup(&server_name, prev).await
+                    });
+                    task.map(|joined| {
+                        joined.unwrap_or_else(|err| {
+                            warn!("Well-known lookup failed: {err}");
+                            None
+                        })
+                    })
                     .boxed()
                     .shared()
                 })
@@ -354,6 +364,20 @@ impl WellKnownResolver {
     }
 }
 
+/// Removes a lookup from `in_flight` when it finishes, or panics.
+struct InFlightGuard<'a> {
+    resolver: &'a WellKnownResolver,
+    server_name: &'a str,
+}
+
+impl Drop for InFlightGuard<'_> {
+    fn drop(&mut self) {
+        if let Ok(mut in_flight) = self.resolver.in_flight.lock() {
+            in_flight.remove(self.server_name);
+        }
+    }
+}
+
 fn jitter() -> f64 {
     rand::rng().random_range((1.0 - CACHE_PERIOD_JITTER)..=(1.0 + CACHE_PERIOD_JITTER))
 }
@@ -802,4 +826,26 @@ mod tests {
         assert_eq!(fetcher.calls(), 1);
         assert!(resolver.in_flight.lock().unwrap().is_empty());
     }
+
+    #[tokio::test(start_paused = true)]
+    async fn abandoned_lookup_still_completes() {
+        let fetcher = FakeFetcher::new(delegate("target"));
+        let (resolver, _) = make_resolver(fetcher.clone());
+
+        // Start a lookup, then give up on it before the fetch returns.
+        {
+            let lookup = std::pin::pin!(resolver.get_well_known("example.com"));
+            assert!(futures::poll!(lookup).is_pending());
+        }
+
+        // The fetch carries on without us, and its result is cached.
+        tokio::time::sleep(Duration::from_secs(1)).await;
+        assert!(resolver.in_flight.lock().unwrap().is_empty());
+        assert_eq!(fetcher.calls(), 1);
+        assert_eq!(
+            resolver.get_well_known("example.com").await.as_deref(),
+            Some("target")
+        );
+        assert_eq!(fetcher.calls(), 0);
+    }
 }