From c4ff4232ae886cb650dc2fe4963bf12541b8d288 Mon Sep 17 00:00:00 2001
From: ckie <git-525ff67@ckie.dev>
Date: Mon, 21 Aug 2023 21:07:37 +0300
Subject: [PATCH] rkvm-{server,client}: Remove dead peers quickly

I use rkvm on a machine which hangs multiple times a day, leaving no
TCP FIN or RSTs behind.

Previously, the UX got very confusing once I rebooted the machine, because
it captured inputs without forwarding them anywhere live.

It's unfortunate that the TCP settings aren't shared between the client
& server because keepalive settings should be consistent to avoid one
peer disconnecting and the other waiting longer before timing out.
---
 Cargo.lock                | 14 +++++++++++++-
 rkvm-client/Cargo.toml    |  1 +
 rkvm-client/src/client.rs | 11 +++++++++++
 rkvm-server/Cargo.toml    |  1 +
 rkvm-server/src/server.rs |  8 ++++++++
 5 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index b4aeb2b..a78d9e4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -778,6 +778,7 @@ dependencies = [
  "rkvm-net",
  "rustls-pemfile",
  "serde",
+ "socket2 0.5.3",
  "thiserror",
  "tokio",
  "tokio-rustls",
@@ -829,6 +830,7 @@ dependencies = [
  "rustls-pemfile",
  "serde",
  "slab",
+ "socket2 0.5.3",
  "thiserror",
  "tokio",
  "tokio-rustls",
@@ -970,6 +972,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "socket2"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2538b18701741680e0322a2302176d3253a35388e2e62f172f64f4f16605f877"
+dependencies = [
+ "libc",
+ "windows-sys 0.48.0",
+]
+
 [[package]]
 name = "spin"
 version = "0.5.2"
@@ -1054,7 +1066,7 @@ dependencies = [
  "num_cpus",
  "pin-project-lite",
  "signal-hook-registry",
- "socket2",
+ "socket2 0.4.9",
  "tokio-macros",
  "windows-sys 0.48.0",
 ]
diff --git a/rkvm-client/Cargo.toml b/rkvm-client/Cargo.toml
index a88d79a..710eea4 100644
--- a/rkvm-client/Cargo.toml
+++ b/rkvm-client/Cargo.toml
@@ -19,6 +19,7 @@ clap = { version = "4.2.2", features = ["derive"] }
 thiserror = "1.0.40"
 tokio-rustls = "0.24.0"
 rustls-pemfile = "1.0.2"
+socket2 = { version = "0.5.3", features = ["all"] }
 
 [package.metadata.rpm]
 package = "rkvm-client"
diff --git a/rkvm-client/src/client.rs b/rkvm-client/src/client.rs
index 95e3464..fed04bd 100644
--- a/rkvm-client/src/client.rs
+++ b/rkvm-client/src/client.rs
@@ -3,9 +3,11 @@ use rkvm_net::auth::{AuthChallenge, AuthStatus};
 use rkvm_net::message::Message;
 use rkvm_net::version::Version;
 use rkvm_net::Update;
+use socket2::{SockRef, TcpKeepalive};
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::io;
+use std::time::Duration;
 use thiserror::Error;
 use tokio::io::{AsyncWriteExt, BufStream};
 use tokio::net::TcpStream;
@@ -43,6 +45,15 @@ pub async fn run(
     stream.set_linger(None).map_err(Error::Network)?;
     stream.set_nodelay(false).map_err(Error::Network)?;
 
+    SockRef::from(&stream)
+        .set_tcp_keepalive(
+            &TcpKeepalive::new()
+                .with_time(Duration::from_secs(1))
+                .with_interval(Duration::from_secs(10))
+                .with_retries(1),
+        )
+        .map_err(Error::Network)?;
+
     log::info!("Connected to server");
 
     let stream = connector
diff --git a/rkvm-server/Cargo.toml b/rkvm-server/Cargo.toml
index ad607a6..e5661f0 100644
--- a/rkvm-server/Cargo.toml
+++ b/rkvm-server/Cargo.toml
@@ -21,6 +21,7 @@ rustls-pemfile = "1.0.2"
 thiserror = "1.0.40"
 slab = "0.4.8"
 rand = "0.8.5"
+socket2 = { version = "0.5.3", features = ["all"] }
 
 [package.metadata.rpm]
 package = "rkvm-server"
diff --git a/rkvm-server/src/server.rs b/rkvm-server/src/server.rs
index 5e105e6..772d1db 100644
--- a/rkvm-server/src/server.rs
+++ b/rkvm-server/src/server.rs
@@ -8,6 +8,7 @@ use rkvm_net::message::Message;
 use rkvm_net::version::Version;
 use rkvm_net::Update;
 use slab::Slab;
+use socket2::{SockRef, TcpKeepalive};
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::ffi::CString;
 use std::io::{self, ErrorKind};
@@ -267,6 +268,13 @@ async fn client(
         stream.set_linger(None)?;
         stream.set_nodelay(false)?;
 
+        SockRef::from(&stream).set_tcp_keepalive(
+            &TcpKeepalive::new()
+                .with_time(Duration::from_secs(1))
+                .with_interval(Duration::from_secs(10))
+                .with_retries(1),
+        )?;
+
         let stream = acceptor.accept(stream).await?;
         log::info!("{}: TLS connected", addr);