client: add automatic tunnel re-registration on connection failures

When the remote TCP connection fails repeatedly (10 consecutive failures),
the client now automatically re-registers with the server to get a fresh
endpoint port, then restarts all connections. This handles the case where
the server has cleaned up the client entry or the network path changed.

Re-registration uses exponential backoff (2s to 60s) and respects the
shutdown signal throughout.

The public API (open_tunnel/ClientConfig) is unchanged.

Author: gdesmott

client/src/lib.rs

@@ -1,12 +1,15 @@
-use std::sync::Arc;
-use tokio::sync::Semaphore;
+use std::sync::{
+    atomic::{AtomicU32, Ordering},
+    Arc,
+};
 
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
 use socket2::{SockRef, TcpKeepalive};
 use tokio::io;
 use tokio::net::TcpStream;
 pub use tokio::sync::broadcast;
+use tokio::sync::{mpsc, Semaphore};
 use tokio::time::{sleep, Duration};
 
 pub const PROXY_SERVER: &str = "https://your-domain.com";
@@ -18,6 +21,9 @@ const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(10);
 #[cfg(not(target_os = "windows"))]
 const TCP_KEEPALIVE_RETRIES: u32 = 5;
 
+/// Number of consecutive remote connection failures before re-registering the tunnel endpoint.
+const MAX_CONSECUTIVE_FAILURES: u32 = 10;
+
 #[derive(Debug, Serialize, Deserialize)]
 struct ProxyResponse {
     id: String,
@@ -45,7 +51,11 @@ pub struct ClientConfig {
     pub credential: Option<String>,
 }
 
-/// Open tunnels directly between server and localhost
+/// Open tunnels directly between server and localhost.
+///
+/// Registers a tunnel endpoint, then spawns a supervisor that maintains
+/// connections and automatically re-registers if the endpoint becomes
+/// unreachable.
 pub async fn open_tunnel(config: ClientConfig) -> Result<String> {
     let ClientConfig {
         server,
@@ -56,65 +66,128 @@ pub async fn open_tunnel(config: ClientConfig) -> Result<String> {
         max_conn,
         credential,
     } = config;
-    let tunnel_info = get_tunnel_endpoint(server, subdomain, credential).await?;
+    let tunnel_info =
+        get_tunnel_endpoint(server.clone(), subdomain.clone(), credential.clone()).await?;
+    let url = tunnel_info.url.clone();
 
-    // TODO check the connect is failed and restart the proxy.
-    tunnel_to_endpoint(
-        tunnel_info.clone(),
+    let supervisor_config = SupervisorConfig {
+        server,
+        subdomain,
+        credential,
         local_host,
         local_port,
         shutdown_signal,
         max_conn,
-    )
-    .await;
+    };
+    tokio::spawn(tunnel_supervisor(supervisor_config, tunnel_info));
 
-    Ok(tunnel_info.url)
+    Ok(url)
 }
 
-async fn get_tunnel_endpoint(
+struct SupervisorConfig {
     server: Option<String>,
     subdomain: Option<String>,
     credential: Option<String>,
-) -> Result<TunnelServerInfo> {
-    let server = server.as_deref().unwrap_or(PROXY_SERVER);
-    let assigned_domain = subdomain.as_deref().unwrap_or("?new");
-    let mut uri = format!("{}/{}", server, assigned_domain);
-    if let Some(credential) = credential {
-        uri = format!("{}?credential={}", uri, credential);
-    }
-    log::info!("Request for assign domain: {}", uri);
+    local_host: Option<String>,
+    local_port: u16,
+    shutdown_signal: broadcast::Sender<()>,
+    max_conn: u8,
+}
 
-    let resp = reqwest::get(uri).await?.json::<ProxyResponse>().await?;
-    log::info!("Response from server: {:#?}", resp);
+// Runs the register → connect → detect-failures → re-register cycle.
+//
+// Each iteration ("round") spawns a pool of connections to the current endpoint.
+// If too many consecutive remote-connect failures occur (the server cleaned up
+// our listener port, or the network path changed), connection tasks signal via
+// `reregister_tx` and the supervisor requests a fresh endpoint from the API
+// server—using the same subdomain so the public tunnel URL stays stable.
+async fn tunnel_supervisor(config: SupervisorConfig, initial_info: TunnelServerInfo) {
+    let mut current_info = initial_info;
+    let mut shutdown_rx = config.shutdown_signal.subscribe();
 
-    let parts = resp.url.split("//").collect::<Vec<&str>>();
-    let mut host = parts[1].split(':').collect::<Vec<&str>>()[0];
-    host = match host.split_once('.') {
-        Some((_, base)) => base,
-        None => host,
-    };
+    loop {
+        log::info!("Starting tunnel connections to {:?}", current_info);
 
-    let tunnel_info = TunnelServerInfo {
-        host: host.to_string(),
-        port: resp.port,
-        max_conn_count: resp.max_conn_count,
-        url: resp.url,
-    };
+        // Per-round shutdown channel: lets us stop this round's connections
+        // without tearing down the whole supervisor.
+        let (round_stop_tx, _) = broadcast::channel::<()>(1);
+        let consecutive_failures = Arc::new(AtomicU32::new(0));
+        let (reregister_tx, mut reregister_rx) = mpsc::channel::<()>(1);
 
-    Ok(tunnel_info)
+        start_tunnel_connections(
+            &current_info,
+            config.local_host.clone(),
+            config.local_port,
+            round_stop_tx.clone(),
+            config.max_conn,
+            consecutive_failures,
+            reregister_tx,
+        );
+
+        // Block until either the connections ask for re-registration or we
+        // are told to shut down entirely.
+        tokio::select! {
+            _ = reregister_rx.recv() => {
+                log::warn!(
+                    "Re-registering tunnel after {} consecutive remote connection failures",
+                    MAX_CONSECUTIVE_FAILURES
+                );
+                let _ = round_stop_tx.send(());
+                sleep(Duration::from_millis(500)).await;
+            }
+            _ = shutdown_rx.recv() => {
+                let _ = round_stop_tx.send(());
+                return;
+            }
+        }
+
+        // Re-register with exponential backoff (2 s → 4 s → … → 60 s cap).
+        // The same subdomain is requested so the public URL doesn't change;
+        // only the internal listener port is refreshed.
+        let mut backoff = Duration::from_secs(2);
+        loop {
+            match get_tunnel_endpoint(
+                config.server.clone(),
+                config.subdomain.clone(),
+                config.credential.clone(),
+            )
+            .await
+            {
+                Ok(info) => {
+                    log::info!("Re-registered tunnel endpoint: {:?}", info);
+                    current_info = info;
+                    break;
+                }
+                Err(err) => {
+                    log::error!(
+                        "Re-registration failed: {:?}, retrying in {:?}",
+                        err,
+                        backoff
+                    );
+                    tokio::select! {
+                        _ = sleep(backoff) => {
+                            backoff = (backoff * 2).min(Duration::from_secs(60));
+                        }
+                        _ = shutdown_rx.recv() => return,
+                    }
+                }
+            }
+        }
+    }
 }
 
-async fn tunnel_to_endpoint(
-    server: TunnelServerInfo,
+fn start_tunnel_connections(
+    server: &TunnelServerInfo,
     local_host: Option<String>,
     local_port: u16,
     shutdown_signal: broadcast::Sender<()>,
     max_conn: u8,
+    consecutive_failures: Arc<AtomicU32>,
+    reregister_tx: mpsc::Sender<()>,
 ) {
-    log::info!("Tunnel server info: {:?}", server);
-    let server_host = server.host;
+    let server_host = server.host.clone();
     let server_port = server.port;
-    let local_host = local_host.unwrap_or(LOCAL_HOST.to_string());
+    let local_host = local_host.unwrap_or_else(|| LOCAL_HOST.to_string());
 
     let count = std::cmp::min(server.max_conn_count, max_conn);
     log::info!("Max connection count: {}", count);
@@ -135,50 +208,82 @@ async fn tunnel_to_endpoint(
                     };
                     let server_host = server_host.clone();
                     let local_host = local_host.clone();
-
+                    let consecutive_failures = consecutive_failures.clone();
+                    let reregister_tx = reregister_tx.clone();
                     let mut shutdown_receiver = shutdown_signal.subscribe();
 
                     tokio::spawn(async move {
-                        log::info!("Create a new proxy connection.");
                         tokio::select! {
-                            res = handle_connection(server_host, server_port, local_host, local_port) => {
-                                match res {
-                                    Ok(_) => log::info!("Connection result: {:?}", res),
-                                    Err(err) => {
-                                        log::error!("Failed to connect to proxy or local server: {:?}", err);
-                                        sleep(Duration::from_secs(10)).await;
-                                    }
-                                }
-                            }
+                            _ = tunnel_one_connection(
+                                &server_host, server_port,
+                                &local_host, local_port,
+                                &consecutive_failures, &reregister_tx,
+                            ) => {}
                             _ = shutdown_receiver.recv() => {
-                                log::info!("Shutting down the connection immediately");
+                                log::info!("Shutting down connection");
                             }
                         }
 
                         drop(permit);
                     });
                 }
                 _ = shutdown_receiver.recv() => {
-                    log::info!("Shuttign down the loop immediately");
+                    log::info!("Shutting down the loop");
                     return;
                 }
             };
         }
     });
 }
 
-async fn handle_connection(
-    remote_host: String,
-    remote_port: u16,
-    local_host: String,
+// Handle one tunnel connection attempt.
+//
+// Only *remote* TCP-connect failures count toward re-registration: they
+// indicate the server's listener port is gone (cleanup, restart, etc.).
+// Local-connect or proxy errors (e.g. local server restarting) do not bump
+// the counter, avoiding spurious re-registration loops.
+//
+// A single successful remote connect resets the counter to zero, so
+// intermittent failures won't accumulate across otherwise-healthy rounds.
+async fn tunnel_one_connection(
+    server_host: &str,
+    server_port: u16,
+    local_host: &str,
+    local_port: u16,
+    consecutive_failures: &AtomicU32,
+    reregister_tx: &mpsc::Sender<()>,
+) {
+    log::debug!("Connecting to remote: {}:{}", server_host, server_port);
+    let remote_stream = match TcpStream::connect(format!("{server_host}:{server_port}")).await {
+        Ok(stream) => {
+            consecutive_failures.store(0, Ordering::Relaxed);
+            stream
+        }
+        Err(err) => {
+            let failures = consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1;
+            log::error!("Remote connect failed ({failures} consecutive): {:?}", err);
+            if failures >= MAX_CONSECUTIVE_FAILURES {
+                let _ = reregister_tx.try_send(());
+            }
+            sleep(Duration::from_secs(10)).await;
+            return;
+        }
+    };
+
+    if let Err(err) = proxy_through(remote_stream, local_host, local_port).await {
+        log::error!("Proxy error: {:?}", err);
+        sleep(Duration::from_secs(10)).await;
+    }
+}
+
+async fn proxy_through(
+    mut remote_stream: TcpStream,
+    local_host: &str,
     local_port: u16,
 ) -> Result<()> {
-    log::debug!("Connect to remote: {}, {}", remote_host, remote_port);
-    let mut remote_stream = TcpStream::connect(format!("{}:{}", remote_host, remote_port)).await?;
-    log::debug!("Connect to local: {}, {}", local_host, local_port);
-    let mut local_stream = TcpStream::connect(format!("{}:{}", local_host, local_port)).await?;
+    log::debug!("Connecting to local: {}:{}", local_host, local_port);
+    let mut local_stream = TcpStream::connect(format!("{local_host}:{local_port}")).await?;
 
-    // configure keepalive on remote socket to early detect network issues and attempt to re-establish the connection.
     let ka = TcpKeepalive::new()
         .with_time(TCP_KEEPALIVE_TIME)
         .with_interval(TCP_KEEPALIVE_INTERVAL);
@@ -190,3 +295,36 @@ async fn handle_connection(
     io::copy_bidirectional(&mut remote_stream, &mut local_stream).await?;
     Ok(())
 }
+
+async fn get_tunnel_endpoint(
+    server: Option<String>,
+    subdomain: Option<String>,
+    credential: Option<String>,
+) -> Result<TunnelServerInfo> {
+    let server = server.as_deref().unwrap_or(PROXY_SERVER);
+    let assigned_domain = subdomain.as_deref().unwrap_or("?new");
+    let mut uri = format!("{}/{}", server, assigned_domain);
+    if let Some(credential) = credential {
+        uri = format!("{}?credential={}", uri, credential);
+    }
+    log::info!("Request for assign domain: {}", uri);
+
+    let resp = reqwest::get(uri).await?.json::<ProxyResponse>().await?;
+    log::info!("Response from server: {:#?}", resp);
+
+    let parts = resp.url.split("//").collect::<Vec<&str>>();
+    let mut host = parts[1].split(':').collect::<Vec<&str>>()[0];
+    host = match host.split_once('.') {
+        Some((_, base)) => base,
+        None => host,
+    };
+
+    let tunnel_info = TunnelServerInfo {
+        host: host.to_string(),
+        port: resp.port,
+        max_conn_count: resp.max_conn_count,
+        url: resp.url,
+    };
+
+    Ok(tunnel_info)
+}