algesten · sosthene-nitrokey · Feb 3, 2025
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -31,6 +31,7 @@ gzip = ["dep:flate2"]
 brotli = ["dep:brotli-decompressor"]
 charset = ["dep:encoding_rs"]
 json = ["dep:serde", "dep:serde_json", "cookie_store?/serde_json"]
+keepalive = ["dep:socket2", "socket2/all"]
 
 ######## UNSTABLE FEATURES.
 # Might be removed or changed in a minor version.
@@ -90,6 +91,7 @@ encoding_rs = { version = "0.8.34", optional = true }
 
 serde = { version = "1.0.138", optional = true, default-features = false, features = ["std"] }
 serde_json = { version = "1.0.120", optional = true, default-features = false, features = ["std"] }
+socket2 = { version = "0.5.8", optional = true }
 
 [dev-dependencies]
 env_logger = "0.11.6"

diff --git a/src/config.rs b/src/config.rs
@@ -156,6 +156,8 @@ pub struct Config {
     max_idle_connections: usize,
     max_idle_connections_per_host: usize,
     max_idle_age: Duration,
+    #[cfg(feature = "keepalive")]
+    tcp_keepalive: Option<socket2::TcpKeepalive>,
 
     // Chain built for middleware.
     pub(crate) middleware: MiddlewareChain,
@@ -394,6 +396,12 @@ impl Config {
     pub fn max_idle_age(&self) -> Duration {
         self.max_idle_age
     }
+
+    #[cfg(feature = "keepalive")]
+    /// TCP Keepalive configuration
+    pub fn tcp_keepalive(&self) -> &Option<socket2::TcpKeepalive> {
+        &self.tcp_keepalive
+    }
 }
 
 /// Builder of [`Config`]
@@ -712,6 +720,17 @@ impl<Scope: private::ConfigScope> ConfigBuilder<Scope> {
         self.config().timeouts.recv_body = v;
         self
     }
+
+    /// Configuration of TCP keepalive
+    ///
+    /// If `None`, keepalive is not configured for the connections
+    ///
+    /// Defaults to `None`
+    #[cfg(feature = "keepalive")]
+    pub fn tcp_keepalive(mut self, v: Option<socket2::TcpKeepalive>) -> Self {
+        self.config().tcp_keepalive = v;
+        self
+    }
 }
 
 /// Possible config values for headers.
@@ -849,6 +868,8 @@ impl Default for Config {
             max_idle_connections_per_host: 3,
             max_idle_age: Duration::from_secs(15),
             middleware: MiddlewareChain::default(),
+            #[cfg(feature = "keepalive")]
+            tcp_keepalive: None,
             force_send_body: false,
         }
     }

diff --git a/src/unversioned/transport/tcp.rs b/src/unversioned/transport/tcp.rs
@@ -33,6 +33,8 @@ impl<In: Transport> Connector<In> for TcpConnector {
 
         let config = &details.config;
         let stream = try_connect(&details.addrs, details.timeout, config)?;
+        #[cfg(feature = "keepalive")]
+        let stream = set_keepalive(stream, details.config.tcp_keepalive().as_ref())?;
 
         let buffers = LazyBuffers::new(config.input_buffer_size(), config.output_buffer_size());
         let transport = TcpTransport::new(stream, buffers);
@@ -41,6 +43,20 @@ impl<In: Transport> Connector<In> for TcpConnector {
     }
 }
 
+#[cfg(feature = "keepalive")]
+fn set_keepalive(
+    stream: TcpStream,
+    keepalive: Option<&socket2::TcpKeepalive>,
+) -> Result<TcpStream, crate::Error> {
+    use socket2::Socket;
+    let socket: Socket = stream.into();
+    if let Some(keepalive) = keepalive {
+        socket.set_keepalive(true)?;
+        socket.set_tcp_keepalive(keepalive)?;
+    }
+    Ok(socket.into())
+}
+
 fn try_connect(
     addrs: &ResolvedSocketAddrs,
     timeout: NextTimeout,