From 1e07eb7b2ab5e315142140f5242b32a090ec864c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=A3rebe=20-=20Romain=20GERARD?= <erebe@erebe.eu>
Date: Mon, 29 Apr 2024 18:07:29 +0200
Subject: [PATCH] Add tcp-keepalive to help detect broken reverse tunnel

---
 src/tcp.rs | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/tcp.rs b/src/tcp.rs
index 786c6da..8f704f5 100644
--- a/src/tcp.rs
+++ b/src/tcp.rs
@@ -5,7 +5,9 @@ use crate::dns::DnsResolver;
 use base64::Engine;
 use bytes::BytesMut;
 use log::warn;
+use socket2::TcpKeepalive;
 use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
+
 use std::time::Duration;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpSocket, TcpStream};
@@ -16,22 +18,31 @@ use tracing::{debug, instrument};
 use url::{Host, Url};
 
 fn configure_socket(socket: &mut TcpSocket, so_mark: &Option<u32>) -> Result<(), anyhow::Error> {
+    let socket = socket2::SockRef::from(&socket);
     socket
         .set_nodelay(true)
-        .with_context(|| format!("cannot set no_delay on socket: {}", io::Error::last_os_error()))?;
+        .with_context(|| format!("cannot set no_delay on socket: {:?}", io::Error::last_os_error()))?;
+
+    #[cfg(not(target_os = "windows"))]
+    let tcp_keepalive = TcpKeepalive::new()
+        .with_time(Duration::from_secs(60))
+        .with_interval(Duration::from_secs(10))
+        .with_retries(3);
+
+    #[cfg(target_os = "windows")]
+    let tcp_keepalive = TcpKeepalive::new()
+        .with_time(Duration::from_secs(60))
+        .with_interval(Duration::from_secs(10));
+
+    socket
+        .set_tcp_keepalive(&tcp_keepalive)
+        .with_context(|| format!("cannot set tcp_keepalive on socket: {:?}", io::Error::last_os_error()))?;
 
     #[cfg(target_os = "linux")]
     if let Some(so_mark) = so_mark {
-        use std::os::fd::AsFd;
-
-        let ret = nix::sys::socket::setsockopt(&socket.as_fd(), nix::sys::socket::sockopt::Mark, so_mark);
-        if let Err(err) = ret {
-            return Err(anyhow!(
-                "Cannot set SO_MARK on the connection {:?} {:?}",
-                err,
-                io::Error::last_os_error()
-            ));
-        }
+        socket
+            .set_mark(*so_mark)
+            .with_context(|| format!("cannot set SO_MARK on socket: {:?}", io::Error::last_os_error()))?;
     }
 
     Ok(())