Merge pull request #350 from Dstack-TEE/wg-reconnect

cvm: Auto reconnect when wg get stucked

Author: kvinwang

basefiles/wg-checker.service

@@ -1,6 +1,6 @@
 [Unit]
 Description=WireGuard Endpoint Checker Service
-After=network-online.target tboot.service
+After=network-online.target dstack-prepare.service
 Wants=network-online.target
 
 [Service]

basefiles/wg-checker.sh

@@ -4,27 +4,73 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-get_conf_endpoint() {
-    grep "Endpoint" /etc/wireguard/wg0.conf | awk "{print \$3}"
+HANDSHAKE_TIMEOUT=180
+LAST_REFRESH=0
+STALE_SINCE=0
+DSTACK_WORK_DIR=${DSTACK_WORK_DIR:-/dstack}
+IFNAME=dstack-wg0
+
+get_latest_handshake() {
+    wg show $IFNAME latest-handshakes 2>/dev/null | awk 'BEGIN { max = 0 } NF >= 2 { if ($2 > max) max = $2 } END { print max }'
 }
 
-get_current_endpoint() {
-    wg show wg0 endpoints | awk "{print \$2}"
+maybe_refresh() {
+    now=$1
+
+    if [ "$LAST_REFRESH" -ne 0 ] && [ $((now - LAST_REFRESH)) -lt $HANDSHAKE_TIMEOUT ]; then
+        return
+    fi
+
+    if ! command -v dstack-util >/dev/null 2>&1; then
+        printf 'dstack-util not found; cannot refresh gateway.\n' >&2
+        LAST_REFRESH=$now
+        return
+    fi
+
+    printf 'WireGuard handshake stale; refreshing dstack gateway...\n'
+    if dstack-util gateway-refresh --work-dir "$DSTACK_WORK_DIR"; then
+        printf 'dstack gateway refresh succeeded.\n'
+    else
+        printf 'dstack gateway refresh failed.\n' >&2
+    fi
+
+    LAST_REFRESH=$now
+    STALE_SINCE=$now
 }
 
-check_endpoint() {
-    CONF_ENDPOINT=$(get_conf_endpoint)
-    CURRENT_ENDPOINT=$(get_current_endpoint)
+check_handshake() {
+    if ! command -v wg >/dev/null 2>&1; then
+        return
+    fi
+
+    now=$(date +%s)
+    latest=$(get_latest_handshake)
+
+    if [ -z "$latest" ]; then
+        latest=0
+    fi
 
-    if [ "$CURRENT_ENDPOINT" != "$CONF_ENDPOINT" ]; then
-        echo "Wg endpoint changed from $CONF_ENDPOINT to $CURRENT_ENDPOINT."
-        wg syncconf wg0 <(wg-quick strip wg0)
+    if [ "$latest" -gt 0 ]; then
+        if [ $((now - latest)) -ge $HANDSHAKE_TIMEOUT ]; then
+            maybe_refresh "$now"
+        else
+            STALE_SINCE=0
+        fi
+    else
+        if [ "$STALE_SINCE" -eq 0 ]; then
+            STALE_SINCE=$now
+        fi
+        if [ $((now - STALE_SINCE)) -ge $HANDSHAKE_TIMEOUT ]; then
+            maybe_refresh "$now"
+        fi
     fi
 }
 
 while true; do
-    if [ -f /etc/wireguard/wg0.conf ]; then
-        check_endpoint
+    if [ -f /etc/wireguard/$IFNAME.conf ]; then
+        check_handshake
+    else
+        STALE_SINCE=0
     fi
     sleep 10
 done

dstack-util/src/main.rs

@@ -24,7 +24,7 @@ use std::{
     io::{self, Read, Write},
     path::PathBuf,
 };
-use system_setup::{cmd_sys_setup, SetupArgs};
+use system_setup::{cmd_gateway_refresh, cmd_sys_setup, GatewayRefreshArgs, SetupArgs};
 use tdx_attest as att;
 use utils::AppKeys;
 
@@ -64,6 +64,8 @@ enum Commands {
     Rand(RandArgs),
     /// Prepare dstack system.
     Setup(SetupArgs),
+    /// Refresh the dstack gateway configuration
+    GatewayRefresh(GatewayRefreshArgs),
     /// Notify the host about the dstack app
     NotifyHost(HostNotifyArgs),
     /// Remove orphaned containers
@@ -533,6 +535,9 @@ async fn main() -> Result<()> {
         Commands::Setup(args) => {
             cmd_sys_setup(args).await?;
         }
+        Commands::GatewayRefresh(args) => {
+            cmd_gateway_refresh(args).await?;
+        }
         Commands::NotifyHost(args) => {
             cmd_notify_host(args).await?;
         }