From a206a90b364898a4fadfec7b9fbb1513b3e41216 Mon Sep 17 00:00:00 2001
From: Systemcluster <me@systemcluster.me>
Date: Mon, 22 Jul 2024 20:41:24 +0200
Subject: [PATCH] Add timeout to tcp connections

---
 src/client.rs           | 17 ++++++++++++++++-
 src/dist/client_auth.rs |  4 ++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 51a6a1fe3..f131a695f 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -19,6 +19,18 @@ use byteorder::{BigEndian, ByteOrder};
 use retry::{delay::Fixed, retry};
 use std::io::{self, BufReader, BufWriter, Read};
 use std::net::TcpStream;
+use std::time::Duration;
+use std::{env, net::ToSocketAddrs};
+
+const DEFAULT_CONNECTION_TIMEOUT: u64 = 20;
+
+/// Get the time clients should try to connect to the server before continuing, in seconds.
+pub(crate) fn get_connection_timeout() -> u64 {
+    env::var("SCCACHE_CONNECTION_TIMEOUT")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(DEFAULT_CONNECTION_TIMEOUT)
+}
 
 /// A connection to an sccache server.
 pub struct ServerConnection {
@@ -65,7 +77,10 @@ impl ServerConnection {
 /// Establish a TCP connection to an sccache server listening on `port`.
 pub fn connect_to_server(port: u16) -> io::Result<ServerConnection> {
     trace!("connect_to_server({})", port);
-    let stream = TcpStream::connect(("127.0.0.1", port))?;
+    let stream = TcpStream::connect_timeout(
+        &("127.0.0.1", port).to_socket_addrs()?.next().unwrap(),
+        Duration::from_secs(get_connection_timeout()),
+    )?;
     ServerConnection::new(stream)
 }
 
diff --git a/src/dist/client_auth.rs b/src/dist/client_auth.rs
index 5555411c0..2aa267c44 100644
--- a/src/dist/client_auth.rs
+++ b/src/dist/client_auth.rs
@@ -15,7 +15,7 @@ use tokio::runtime::Runtime;
 use url::Url;
 use uuid::Uuid;
 
-use crate::errors::*;
+use crate::{client::get_connection_timeout, errors::*};
 
 // These (arbitrary) ports need to be registered as valid redirect urls in the oauth provider you're using
 pub const VALID_PORTS: &[u16] = &[12731, 32492, 56909];
@@ -503,7 +503,7 @@ async fn try_bind() -> Result<HyperBuilderWrap> {
             .expect("Expected at least one address in parsed socket address");
 
         // Hyper binds with reuseaddr and reuseport so binding won't fail as you'd expect on Linux
-        match TcpStream::connect(addr) {
+        match TcpStream::connect_timeout(&addr, Duration::from_secs(get_connection_timeout())) {
             // Already open
             Ok(_) => continue,
             // Doesn't seem to be open