move exit node ping to module

exitnode/exitnode.go

@@ -0,0 +1,176 @@
+// Package exitnode implements the exit-node ping dance run before
+// registering with the server: request the candidate exit nodes, ping each
+// one over HTTP, and report the results so the server can pick the best one.
+// It is shared between newt and olm, which both register the same way.
+package exitnode
+
+import (
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/fosrl/newt/logger"
+)
+
+// ExitNodeData is the payload the server sends in response to a
+// "*/ping/request" message.
+type ExitNodeData struct {
+	ExitNodes []ExitNode `json:"exitNodes"`
+	ChainId   string     `json:"chainId"`
+}
+
+// ExitNode is a candidate exit node offered by the server for ping selection.
+type ExitNode struct {
+	ID                     int     `json:"exitNodeId"`
+	Name                   string  `json:"exitNodeName"`
+	Endpoint               string  `json:"endpoint"`
+	Weight                 float64 `json:"weight"`
+	WasPreviouslyConnected bool    `json:"wasPreviouslyConnected"`
+}
+
+// ExitNodePingResult is the measured latency (or error) for one exit node,
+// sent back to the server in the "*/wg/register" message's pingResults field.
+type ExitNodePingResult struct {
+	ExitNodeID             int     `json:"exitNodeId"`
+	LatencyMs              int64   `json:"latencyMs"`
+	Weight                 float64 `json:"weight"`
+	Error                  string  `json:"error,omitempty"`
+	Name                   string  `json:"exitNodeName"`
+	Endpoint               string  `json:"endpoint"`
+	WasPreviouslyConnected bool    `json:"wasPreviouslyConnected"`
+}
+
+// PingExitNodes pings the given exit nodes over HTTP and returns a per-node
+// ExitNodePingResult suitable for inclusion in a wg/register message's
+// pingResults field, so the server can select the best exit node.
+//
+// If there's only one exit node, or preferEndpoint names one of them, the
+// matching node is returned immediately with LatencyMs 0 and no pinging is
+// done. Otherwise every node is pinged pingAttempts times over HTTP GET
+// <endpoint>/ping and the average latency of successful attempts is used.
+//
+// When alreadyConnected is true, a node flagged WasPreviouslyConnected is
+// excluded from the results as long as at least one other healthy node is
+// available, biasing reconnects toward switching away from a possibly
+// degraded node.
+func PingExitNodes(exitNodes []ExitNode, preferEndpoint string, alreadyConnected bool) []ExitNodePingResult {
+	if len(exitNodes) == 0 {
+		return nil
+	}
+
+	if len(exitNodes) == 1 || preferEndpoint != "" {
+		selected := exitNodes[0]
+		if preferEndpoint != "" {
+			for _, node := range exitNodes {
+				if node.Endpoint == preferEndpoint {
+					selected = node
+					break
+				}
+			}
+		}
+
+		logger.Debug("Only one exit node available, using it directly: %s", selected.Endpoint)
+
+		return []ExitNodePingResult{
+			{
+				ExitNodeID:             selected.ID,
+				LatencyMs:              0,
+				Weight:                 selected.Weight,
+				Error:                  "",
+				Name:                   selected.Name,
+				Endpoint:               selected.Endpoint,
+				WasPreviouslyConnected: selected.WasPreviouslyConnected,
+			},
+		}
+	}
+
+	type nodeResult struct {
+		Node    ExitNode
+		Latency time.Duration
+		Err     error
+	}
+
+	results := make([]nodeResult, len(exitNodes))
+	const pingAttempts = 3
+	for i, node := range exitNodes {
+		var totalLatency time.Duration
+		var lastErr error
+		successes := 0
+		httpClient := &http.Client{
+			Timeout: 5 * time.Second,
+		}
+		url := node.Endpoint
+		if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
+			url = "http://" + url
+		}
+		if !strings.HasSuffix(url, "/ping") {
+			url = strings.TrimRight(url, "/") + "/ping"
+		}
+		for j := 0; j < pingAttempts; j++ {
+			start := time.Now()
+			resp, err := httpClient.Get(url)
+			latency := time.Since(start)
+			if err != nil {
+				lastErr = err
+				logger.Warn("Failed to ping exit node %d (%s) attempt %d: %v", node.ID, url, j+1, err)
+				continue
+			}
+			resp.Body.Close()
+			totalLatency += latency
+			successes++
+		}
+		var avgLatency time.Duration
+		if successes > 0 {
+			avgLatency = totalLatency / time.Duration(successes)
+		}
+		if successes == 0 {
+			results[i] = nodeResult{Node: node, Latency: 0, Err: lastErr}
+		} else {
+			results[i] = nodeResult{Node: node, Latency: avgLatency, Err: nil}
+		}
+	}
+
+	var pingResults []ExitNodePingResult
+	for _, res := range results {
+		errMsg := ""
+		if res.Err != nil {
+			errMsg = res.Err.Error()
+		}
+		pingResults = append(pingResults, ExitNodePingResult{
+			ExitNodeID:             res.Node.ID,
+			LatencyMs:              res.Latency.Milliseconds(),
+			Weight:                 res.Node.Weight,
+			Error:                  errMsg,
+			Name:                   res.Node.Name,
+			Endpoint:               res.Node.Endpoint,
+			WasPreviouslyConnected: res.Node.WasPreviouslyConnected,
+		})
+	}
+
+	if alreadyConnected {
+		var filteredPingResults []ExitNodePingResult
+		previouslyConnectedNodeIdx := -1
+		for i, res := range pingResults {
+			if res.WasPreviouslyConnected {
+				previouslyConnectedNodeIdx = i
+			}
+		}
+		goodNodeCount := 0
+		for i, res := range pingResults {
+			if i != previouslyConnectedNodeIdx && res.LatencyMs > 0 && res.Error == "" {
+				goodNodeCount++
+			}
+		}
+		if previouslyConnectedNodeIdx != -1 && goodNodeCount > 0 {
+			for i, res := range pingResults {
+				if i != previouslyConnectedNodeIdx {
+					filteredPingResults = append(filteredPingResults, res)
+				}
+			}
+			pingResults = filteredPingResults
+			logger.Info("Excluding previously connected exit node from ping results due to other available nodes")
+		}
+	}
+
+	return pingResults
+}