Add way to stop initial ping

This commit is contained in:
Owen
2025-06-22 12:50:55 -04:00
parent e642983b88
commit ac67df63fa
2 changed files with 66 additions and 40 deletions

38
main.go
View File

@@ -292,6 +292,8 @@ func main() {
}
}
var pingWithRetryStopChan chan struct{}
// Register handlers for different message types
client.RegisterHandler("newt/wg/connect", func(msg websocket.WSMessage) {
logger.Info("Received registration message")
@@ -388,7 +390,12 @@ persistent_keepalive_interval=5`, fixKey(privateKey.String()), fixKey(wgData.Pub
logger.Info("WireGuard device created. Lets ping the server now...")
// Even if pingWithRetry returns an error, it will continue trying in the background
_ = pingWithRetry(tnet, wgData.ServerIP, pingTimeout)
if pingWithRetryStopChan != nil {
// Stop the previous pingWithRetry if it exists
close(pingWithRetryStopChan)
pingWithRetryStopChan = nil
}
pingWithRetryStopChan, _ = pingWithRetry(tnet, wgData.ServerIP, pingTimeout)
// Always mark as connected and start the proxy manager regardless of initial ping result
// as the pings will continue in the background
@@ -703,6 +710,10 @@ persistent_keepalive_interval=5`, fixKey(privateKey.String()), fixKey(wgData.Pub
err = client.SendMessage("newt/socket/containers", map[string]interface{}{
"containers": containers,
})
if err != nil {
logger.Error("Failed to send registration message: %v", err)
}
logger.Info("Sent registration message")
if err != nil {
logger.Error("Failed to send Docker container list: %v", err)
} else {
@@ -718,23 +729,24 @@ persistent_keepalive_interval=5`, fixKey(privateKey.String()), fixKey(wgData.Pub
// request from the server the list of nodes to ping at newt/ping/request
stopFunc = client.SendMessageInterval("newt/ping/request", map[string]interface{}{}, 3*time.Second)
// Send registration message to the server for backward compatibility
err := client.SendMessage("newt/wg/register", map[string]interface{}{
"publicKey": publicKey.String(),
"newtVersion": newtVersion,
"backwardsCompatible": true,
})
if err != nil {
logger.Error("Failed to send registration message: %v", err)
return err
}
logger.Info("Sent registration message")
if wgService != nil {
wgService.LoadRemoteConfig()
}
}
// Send registration message to the server for backward compatibility
err := client.SendMessage("newt/wg/register", map[string]interface{}{
"publicKey": publicKey.String(),
"newtVersion": newtVersion,
"backwardsCompatible": true,
})
if err != nil {
logger.Error("Failed to send registration message: %v", err)
return err
}
logger.Info("Sent registration message")
return nil
})

68
util.go
View File

@@ -1,4 +1,4 @@
package util
package main
import (
"bytes"
@@ -90,13 +90,14 @@ func ping(tnet *netstack.Net, dst string, timeout time.Duration) (time.Duration,
return latency, nil
}
func pingWithRetry(tnet *netstack.Net, dst string, timeout time.Duration) error {
func pingWithRetry(tnet *netstack.Net, dst string, timeout time.Duration) (stopChan chan struct{}, err error) {
const (
initialMaxAttempts = 5
initialRetryDelay = 2 * time.Second
maxRetryDelay = 60 * time.Second // Cap the maximum delay
)
stopChan = make(chan struct{})
attempt := 1
retryDelay := initialRetryDelay
@@ -105,9 +106,8 @@ func pingWithRetry(tnet *netstack.Net, dst string, timeout time.Duration) error
if latency, err := ping(tnet, dst, timeout); err == nil {
// Successful ping
logger.Info("Ping latency: %v", latency)
logger.Info("Tunnel connection to server established successfully!")
return nil
return stopChan, nil
} else {
logger.Warn("Ping attempt %d failed: %v", attempt, err)
}
@@ -117,34 +117,40 @@ func pingWithRetry(tnet *netstack.Net, dst string, timeout time.Duration) error
attempt = 2 // Continue from attempt 2
for {
logger.Info("Ping attempt %d", attempt)
if latency, err := ping(tnet, dst, timeout); err != nil {
logger.Warn("Ping attempt %d failed: %v", attempt, err)
// Increase delay after certain thresholds but cap it
if attempt%5 == 0 && retryDelay < maxRetryDelay {
retryDelay = time.Duration(float64(retryDelay) * 1.5)
if retryDelay > maxRetryDelay {
retryDelay = maxRetryDelay
}
logger.Info("Increasing ping retry delay to %v", retryDelay)
}
time.Sleep(retryDelay)
attempt++
} else {
// Successful ping
logger.Info("Ping succeeded after %d attempts", attempt)
logger.Info("Ping latency: %v", latency)
logger.Info("Tunnel connection to server established successfully!")
select {
case <-stopChan:
logger.Info("Stopping pingWithRetry goroutine")
return
default:
logger.Info("Ping attempt %d", attempt)
if latency, err := ping(tnet, dst, timeout); err != nil {
logger.Warn("Ping attempt %d failed: %v", attempt, err)
// Increase delay after certain thresholds but cap it
if attempt%5 == 0 && retryDelay < maxRetryDelay {
retryDelay = time.Duration(float64(retryDelay) * 1.5)
if retryDelay > maxRetryDelay {
retryDelay = maxRetryDelay
}
logger.Info("Increasing ping retry delay to %v", retryDelay)
}
time.Sleep(retryDelay)
attempt++
} else {
// Successful ping
logger.Info("Ping succeeded after %d attempts", attempt)
logger.Info("Ping latency: %v", latency)
logger.Info("Tunnel connection to server established successfully!")
return
}
}
}
}()
// Return an error for the first batch of attempts (to maintain compatibility with existing code)
return fmt.Errorf("initial ping attempts failed, continuing in background")
return stopChan, fmt.Errorf("initial ping attempts failed, continuing in background")
}
func startPingCheck(tnet *netstack.Net, serverIP string, client *websocket.Client) chan struct{} {
@@ -171,6 +177,14 @@ func startPingCheck(tnet *netstack.Net, serverIP string, client *websocket.Clien
connectionLost = true
logger.Warn("Connection to server lost. Continuous reconnection attempts will be made.")
stopFunc = client.SendMessageInterval("newt/ping/request", map[string]interface{}{}, 3*time.Second)
// Send registration message to the server for backward compatibility
err := client.SendMessage("newt/wg/register", map[string]interface{}{
"publicKey": publicKey.String(),
"backwardsCompatible": true,
})
if err != nil {
logger.Error("Failed to send registration message: %v", err)
}
}
currentInterval = time.Duration(float64(currentInterval) * 1.5)
if currentInterval > maxInterval {
@@ -369,7 +383,7 @@ func updateTargets(pm *proxy.ProxyManager, action string, tunnelIP string, proto
return nil
}
func executeUpdownScript(action, proto, target string, updownScript string) (string, error) {
func executeUpdownScript(action, proto, target string) (string, error) {
if updownScript == "" {
return target, nil
}