filter out magic packets in the middle device to prevent flapping

device/middle_device.go

@@ -8,6 +8,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/fosrl/newt/bind"
 	"github.com/fosrl/newt/logger"
 	"golang.zx2c4.com/wireguard/tun"
 )
@@ -423,6 +424,33 @@ func extractDestIP(packet []byte) (netip.Addr, bool) {
 	return netip.Addr{}, false
 }
 
+// extractUDPPayload returns the UDP payload of packet, if packet is a well-formed
+// IPv4 or IPv6 UDP datagram (ignoring IPv6 extension headers).
+func extractUDPPayload(packet []byte) ([]byte, bool) {
+	if len(packet) < 20 {
+		return nil, false
+	}
+
+	const udpProtocol = 17
+
+	switch packet[0] >> 4 {
+	case 4:
+		ihl := int(packet[0]&0x0f) * 4
+		if ihl < 20 || len(packet) < ihl+8 || packet[9] != udpProtocol {
+			return nil, false
+		}
+		return packet[ihl+8:], true
+	case 6:
+		const ipv6HeaderLen = 40
+		if len(packet) < ipv6HeaderLen+8 || packet[6] != udpProtocol {
+			return nil, false
+		}
+		return packet[ipv6HeaderLen+8:], true
+	}
+
+	return nil, false
+}
+
 // Read intercepts packets going UP from the TUN device (towards WireGuard)
 func (d *MiddleDevice) Read(bufs [][]byte, sizes []int, offset int) (n int, err error) {
 	for {
@@ -497,17 +525,19 @@ func (d *MiddleDevice) Read(bufs [][]byte, sizes []int, offset int) (n int, err
 		rules := d.rules
 		d.rulesMutex.RUnlock()
 
-		if len(rules) == 0 {
-			return n, nil
-		}
-
-		// Process packets and filter out handled ones
+		// Process packets and filter out handled ones. This always runs (even with
+		// no per-IP rules registered) so magic connectivity-test packets can be
+		// dropped before they reach WireGuard - see isLeakedMagicPacket.
 		writeIdx := 0
 		for readIdx := 0; readIdx < n; readIdx++ {
 			packet := bufs[readIdx][offset : offset+sizes[readIdx]]
 
+			if isLeakedMagicPacket(packet) {
+				continue
+			}
+
 			destIP, ok := extractDestIP(packet)
-			if !ok {
+			if !ok || len(rules) == 0 {
 				if writeIdx != readIdx {
 					bufs[writeIdx] = bufs[readIdx]
 					sizes[writeIdx] = sizes[readIdx]
@@ -539,6 +569,57 @@ func (d *MiddleDevice) Read(bufs [][]byte, sizes []int, offset int) (n int, err
 	}
 }
 
+// isLeakedMagicPacket reports whether packet carries one of our UDP connectivity-test
+// magic payloads (see bind.IsMagicPacket). These packets are sent directly between
+// physical UDP sockets by the local-endpoint holepunch tester and must never be
+// encapsulated by WireGuard: if OS routing sends one into this TUN interface instead
+// of out the real network interface (e.g. because the destination falls inside a
+// routed tunnel subnet), tunneling and echoing it back would make a LAN-local
+// endpoint falsely appear directly reachable. Dropping it here makes the test
+// correctly time out instead.
+func isLeakedMagicPacket(packet []byte) bool {
+	payload, ok := extractUDPPayload(packet)
+	return ok && bind.IsMagicPacket(payload)
+}
+
+// filterDownstreamBufs drops packets going DOWN to the TUN device (from WireGuard)
+// that are handled by a per-IP rule or are a leaked magic connectivity-test packet
+// (see isLeakedMagicPacket) - always checked, even with no rules registered. It
+// returns bufs unchanged (no allocation) unless a packet actually needs to be
+// dropped, at which point it switches to an owned copy of the buffers kept so far.
+func filterDownstreamBufs(bufs [][]byte, rules []FilterRule, offset int) [][]byte {
+	filtered := bufs
+	for i, buf := range bufs {
+		drop := len(buf) <= offset
+		if !drop {
+			packet := buf[offset:]
+			if isLeakedMagicPacket(packet) {
+				drop = true
+			} else if destIP, ok := extractDestIP(packet); ok && len(rules) > 0 {
+				for _, rule := range rules {
+					if rule.DestIP == destIP && rule.Handler(packet) {
+						drop = true
+						break
+					}
+				}
+			}
+		}
+
+		if drop {
+			if len(filtered) == len(bufs) {
+				// First drop: switch to an owned, growable copy of everything kept so far.
+				filtered = append([][]byte(nil), bufs[:i]...)
+			}
+			continue
+		}
+
+		if len(filtered) != len(bufs) {
+			filtered = append(filtered, buf)
+		}
+	}
+	return filtered
+}
+
 // Write intercepts packets going DOWN to the TUN device (from WireGuard)
 func (d *MiddleDevice) Write(bufs [][]byte, offset int) (int, error) {
 	for {
@@ -558,38 +639,7 @@ func (d *MiddleDevice) Write(bufs [][]byte, offset int) (int, error) {
 		rules := d.rules
 		d.rulesMutex.RUnlock()
 
-		var filteredBufs [][]byte
-		if len(rules) == 0 {
-			filteredBufs = bufs
-		} else {
-			filteredBufs = make([][]byte, 0, len(bufs))
-			for _, buf := range bufs {
-				if len(buf) <= offset {
-					continue
-				}
-
-				packet := buf[offset:]
-				destIP, ok := extractDestIP(packet)
-				if !ok {
-					filteredBufs = append(filteredBufs, buf)
-					continue
-				}
-
-				handled := false
-				for _, rule := range rules {
-					if rule.DestIP == destIP {
-						if rule.Handler(packet) {
-							handled = true
-							break
-						}
-					}
-				}
-
-				if !handled {
-					filteredBufs = append(filteredBufs, buf)
-				}
-			}
-		}
+		filteredBufs := filterDownstreamBufs(bufs, rules, offset)
 
 		if len(filteredBufs) == 0 {
 			return len(bufs), nil

device/middle_device_test.go

@@ -4,9 +4,22 @@ import (
 	"net/netip"
 	"testing"
 
+	"github.com/fosrl/newt/bind"
 	"github.com/fosrl/newt/util"
 )
 
+// buildIPv4UDPPacket builds a minimal IPv4/UDP packet (no options) carrying payload.
+func buildIPv4UDPPacket(payload []byte) []byte {
+	const ipHeaderLen = 20
+	const udpHeaderLen = 8
+
+	packet := make([]byte, ipHeaderLen+udpHeaderLen+len(payload))
+	packet[0] = 0x45 // version 4, IHL 5
+	packet[9] = 17    // protocol: UDP
+	copy(packet[ipHeaderLen+udpHeaderLen:], payload)
+	return packet
+}
+
 func TestExtractDestIP(t *testing.T) {
 	tests := []struct {
 		name   string
@@ -88,6 +101,49 @@ func TestGetProtocol(t *testing.T) {
 	}
 }
 
+func TestIsLeakedMagicPacket(t *testing.T) {
+	request := make([]byte, bind.MagicTestRequestLen)
+	copy(request, bind.MagicTestRequest)
+
+	response := make([]byte, bind.MagicTestResponseLen)
+	copy(response, bind.MagicTestResponse)
+
+	tests := []struct {
+		name   string
+		packet []byte
+		want   bool
+	}{
+		{
+			name:   "magic test request leaked into tunnel",
+			packet: buildIPv4UDPPacket(request),
+			want:   true,
+		},
+		{
+			name:   "magic test response leaked into tunnel",
+			packet: buildIPv4UDPPacket(response),
+			want:   true,
+		},
+		{
+			name:   "ordinary UDP payload",
+			packet: buildIPv4UDPPacket([]byte("just some ordinary application data")),
+			want:   false,
+		},
+		{
+			name:   "too short to be a packet",
+			packet: []byte{0x45, 0x00},
+			want:   false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := isLeakedMagicPacket(tt.packet); got != tt.want {
+				t.Errorf("isLeakedMagicPacket() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func BenchmarkExtractDestIP(b *testing.B) {
 	packet := []byte{
 		0x45, 0x00, 0x00, 0x54, 0x00, 0x00, 0x40, 0x00,
@@ -100,3 +156,61 @@ func BenchmarkExtractDestIP(b *testing.B) {
 		extractDestIP(packet)
 	}
 }
+
+func TestFilterDownstreamBufsNoDropIsAllocFree(t *testing.T) {
+	bufs := make([][]byte, 128)
+	for i := range bufs {
+		bufs[i] = buildIPv4UDPPacket(make([]byte, 1372))
+	}
+
+	allocs := testing.AllocsPerRun(1000, func() {
+		out := filterDownstreamBufs(bufs, nil, 0)
+		if len(out) != len(bufs) {
+			t.Fatalf("expected no packets dropped, got %d/%d", len(out), len(bufs))
+		}
+	})
+
+	if allocs != 0 {
+		t.Errorf("filterDownstreamBufs() with nothing to drop allocated %v times per call, want 0", allocs)
+	}
+}
+
+func TestFilterDownstreamBufsDropsMagicPacket(t *testing.T) {
+	request := make([]byte, bind.MagicTestRequestLen)
+	copy(request, bind.MagicTestRequest)
+
+	bufs := [][]byte{
+		buildIPv4UDPPacket([]byte("ordinary payload one")),
+		buildIPv4UDPPacket(request),
+		buildIPv4UDPPacket([]byte("ordinary payload two")),
+	}
+
+	out := filterDownstreamBufs(bufs, nil, 0)
+	if len(out) != 2 {
+		t.Fatalf("expected 1 packet dropped, got %d remaining", len(out))
+	}
+}
+
+func BenchmarkFilterDownstreamBufsNoDrop(b *testing.B) {
+	bufs := make([][]byte, 128)
+	for i := range bufs {
+		bufs[i] = buildIPv4UDPPacket(make([]byte, 1372))
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		filterDownstreamBufs(bufs, nil, 0)
+	}
+}
+
+func BenchmarkIsLeakedMagicPacket(b *testing.B) {
+	// A typical ~1400 byte ordinary application payload (the common case on the
+	// hot path - almost every real packet should look like this).
+	ordinary := buildIPv4UDPPacket(make([]byte, 1372))
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		isLeakedMagicPacket(ordinary)
+	}
+}

go.mod

@@ -32,4 +32,4 @@ require (
 )
 
 // To be used ONLY for local development
-// replace github.com/fosrl/newt => ../newt
+replace github.com/fosrl/newt => ../newt

go.sum

@@ -1,7 +1,5 @@
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
-github.com/fosrl/newt v1.15.0 h1:WpL0whZM1FMjUe2Vy5jSH1bgbxm1O9k1qCyF/mqZT+s=
-github.com/fosrl/newt v1.15.0/go.mod h1:l6kWoZPSaXT+ZRUjiyPgwflRqZWYaXpUj9oQ0sOPh4o=
 github.com/godbus/dbus/v5 v5.2.2 h1:TUR3TgtSVDmjiXOgAAyaZbYmIeP3DPkld3jgKGV8mXQ=
 github.com/godbus/dbus/v5 v5.2.2/go.mod h1:3AAv2+hPq5rdnr5txxxRwiGjPXamgoIHgz9FPBfOp3c=
 github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=

olm/connect.go

@@ -51,6 +51,11 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) {
 		o.updateRegister = nil
 	}
 
+	if o.stopPingRequest != nil {
+		o.stopPingRequest()
+		o.stopPingRequest = nil
+	}
+
 	// if there is an existing tunnel then close it
 	if o.dev != nil {
 		logger.Info("Got new message. Closing existing tunnel!")
@@ -257,6 +262,14 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) {
 		network.SetDNSServers([]string{o.dnsProxy.GetProxyIP().String()})
 	}
 
+	if wgData.ExitNode != nil && wgData.ExitNode.Connect {
+		if err := o.connectExitNode(*wgData.ExitNode); err != nil {
+			logger.Error("Failed to connect to exit node: %v", err)
+		}
+	} else {
+		logger.Debug("No exit node to connect to (not provided, or connect flag is false)")
+	}
+
 	o.apiServer.SetRegistered(true)
 
 	o.registered = true

olm/exitnode.go

@@ -0,0 +1,287 @@
+package olm
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+	"strings"
+
+	"github.com/fosrl/newt/logger"
+	"github.com/fosrl/newt/network"
+	"github.com/fosrl/newt/util"
+	"github.com/fosrl/olm/peers"
+	"github.com/fosrl/olm/websocket"
+)
+
+// exitNodeAliasSiteId is the sentinel siteId used when registering exit node
+// aliases with the DNS proxy. It is not a real site, and the JIT handler
+// treats siteId 0 as "no JIT lookup", which is correct here since the exit
+// node is connected directly rather than on demand.
+const exitNodeAliasSiteId = 0
+
+// connectExitNode configures a WireGuard peer connection to an exit node, on the
+// same interface and WireGuard device already used for site peers. The exit node
+// lives in a different address space than the site tunnel, so a secondary address
+// (ExitNodeConfig.TunnelIP) is added to the interface for it - the exit node's own
+// WireGuard peer entry only accepts traffic sourced from that address. Nothing here
+// is persisted; it's purely in-memory WireGuard/routing state, same as site peers.
+func (o *Olm) connectExitNode(cfg ExitNodeConfig) error {
+	if !o.tunnelRunning {
+		return fmt.Errorf("tunnel not running")
+	}
+	if cfg.PublicKey == "" || cfg.Endpoint == "" || cfg.ServerIP == "" || cfg.TunnelIP == "" {
+		return fmt.Errorf("incomplete exit node configuration")
+	}
+
+	o.exitNodeMu.Lock()
+	defer o.exitNodeMu.Unlock()
+
+	dev := o.dev
+	if dev == nil {
+		return fmt.Errorf("wireguard device not initialized")
+	}
+
+	if o.exitNode != nil && o.exitNode.PublicKey != cfg.PublicKey {
+		logger.Info("Switching exit nodes, removing previous exit node peer")
+		if err := o.removeExitNodePeerLocked(); err != nil {
+			logger.Warn("Failed to remove previous exit node peer: %v", err)
+		}
+	}
+
+	endpoint := cfg.Endpoint
+	if !strings.Contains(endpoint, ":") {
+		relayPort := cfg.RelayPort
+		if relayPort == 0 {
+			relayPort = 21820
+		}
+		endpoint = fmt.Sprintf("%s:%d", endpoint, relayPort)
+	}
+
+	resolvedEndpoint, err := util.ResolveDomain(endpoint)
+	if err != nil {
+		return fmt.Errorf("failed to resolve exit node endpoint: %w", err)
+	}
+
+	persistentKeepalive := 0
+	if pm := o.getPeerManager(); pm != nil {
+		persistentKeepalive = pm.PersistentKeepalive
+	}
+
+	allowedIP := strings.Split(cfg.ServerIP, "/")[0] + "/32"
+
+	wgConfig := fmt.Sprintf(`public_key=%s
+allowed_ip=%s
+endpoint=%s
+persistent_keepalive_interval=%d`, util.FixKey(cfg.PublicKey), allowedIP, resolvedEndpoint, persistentKeepalive)
+
+	if err := dev.IpcSet(wgConfig); err != nil {
+		return fmt.Errorf("failed to configure exit node peer: %w", err)
+	}
+
+	interfaceName := o.tunnelConfig.InterfaceName
+	tunnelIP := cfg.TunnelIP
+	if !strings.Contains(tunnelIP, "/") {
+		tunnelIP += "/32"
+	}
+	if err := network.AddSecondaryAddress(interfaceName, tunnelIP); err != nil {
+		logger.Warn("Failed to add secondary address %s for exit node: %v", tunnelIP, err)
+	}
+
+	if err := network.AddRouteForServerIP(cfg.ServerIP, interfaceName); err != nil {
+		logger.Warn("Failed to add route for exit node server IP: %v", err)
+	}
+
+	cfgCopy := cfg
+	o.exitNode = &cfgCopy
+
+	if o.dnsProxy != nil {
+		serverIP := net.ParseIP(cfg.ServerIP)
+		if serverIP != nil {
+			for _, alias := range cfg.Aliases {
+				logger.Debug("Adding alias %s to the edit node", alias)
+				if err := o.dnsProxy.AddDNSRecord(alias, serverIP, exitNodeAliasSiteId); err != nil {
+					logger.Warn("Failed to add DNS record for exit node alias %s: %v", alias, err)
+				}
+			}
+		}
+	}
+
+	logger.Info("Connected to exit node at %s", resolvedEndpoint)
+	return nil
+}
+
+// disconnectExitNode tears down the current exit node peer connection, if any.
+func (o *Olm) disconnectExitNode() error {
+	o.exitNodeMu.Lock()
+	defer o.exitNodeMu.Unlock()
+
+	return o.removeExitNodePeerLocked()
+}
+
+// removeExitNodePeerLocked removes the current exit node peer, its secondary
+// interface address, and its server IP route. Must be called with exitNodeMu held.
+func (o *Olm) removeExitNodePeerLocked() error {
+	if o.exitNode == nil {
+		return nil
+	}
+	cfg := o.exitNode
+	o.exitNode = nil
+
+	if o.dnsProxy != nil {
+		serverIP := net.ParseIP(cfg.ServerIP)
+		if serverIP != nil {
+			for _, alias := range cfg.Aliases {
+				o.dnsProxy.RemoveDNSRecordForSite(alias, serverIP, exitNodeAliasSiteId)
+			}
+		}
+	}
+
+	if o.dev != nil {
+		if err := peers.RemovePeer(o.dev, 0, cfg.PublicKey); err != nil {
+			logger.Warn("Failed to remove exit node peer: %v", err)
+		}
+	}
+
+	interfaceName := o.tunnelConfig.InterfaceName
+	if err := network.RemoveRouteForServerIP(cfg.ServerIP, interfaceName); err != nil {
+		logger.Warn("Failed to remove route for exit node server IP: %v", err)
+	}
+
+	tunnelIP := cfg.TunnelIP
+	if !strings.Contains(tunnelIP, "/") {
+		tunnelIP += "/32"
+	}
+	if err := network.RemoveSecondaryAddress(interfaceName, tunnelIP); err != nil {
+		logger.Warn("Failed to remove secondary address %s for exit node: %v", tunnelIP, err)
+	}
+
+	logger.Info("Disconnected from exit node")
+	return nil
+}
+
+// handleExitNodeConnect handles a server-initiated request to connect to (or switch to)
+// an exit node, delivered as a full ExitNodeConfig payload.
+func (o *Olm) handleExitNodeConnect(msg websocket.WSMessage) {
+	logger.Debug("Received exit node connect message: %v", msg.Data)
+
+	if !o.tunnelRunning {
+		logger.Debug("Tunnel stopped, ignoring exit node connect message")
+		return
+	}
+
+	jsonData, err := json.Marshal(msg.Data)
+	if err != nil {
+		logger.Error("Error marshaling exit node connect data: %v", err)
+		return
+	}
+
+	var cfg ExitNodeConfig
+	if err := json.Unmarshal(jsonData, &cfg); err != nil {
+		logger.Error("Error unmarshaling exit node connect data: %v", err)
+		return
+	}
+
+	if !cfg.Connect {
+		logger.Debug("Exit node connect message has connect=false, disconnecting instead")
+		if err := o.disconnectExitNode(); err != nil {
+			logger.Error("Failed to disconnect from exit node: %v", err)
+		}
+		return
+	}
+
+	if err := o.connectExitNode(cfg); err != nil {
+		logger.Error("Failed to connect to exit node: %v", err)
+	}
+}
+
+// handleExitNodeDisconnect handles a server-initiated request to disconnect from the
+// currently connected exit node.
+func (o *Olm) handleExitNodeDisconnect(msg websocket.WSMessage) {
+	logger.Debug("Received exit node disconnect message: %v", msg.Data)
+
+	if !o.tunnelRunning {
+		logger.Debug("Tunnel stopped, ignoring exit node disconnect message")
+		return
+	}
+
+	if err := o.disconnectExitNode(); err != nil {
+		logger.Error("Failed to disconnect from exit node: %v", err)
+	}
+}
+
+// handleExitNodeUpdateData handles a server-initiated request to change data
+// associated with the currently connected exit node, such as its aliases (e.g. a
+// resource was renamed). Unlike site aliases, there is no per-alias address to
+// track since every exit node alias resolves to the exit node's own ServerIP.
+func (o *Olm) handleExitNodeUpdateData(msg websocket.WSMessage) {
+	logger.Debug("Received exit node update data message: %v", msg.Data)
+
+	if !o.tunnelRunning {
+		logger.Debug("Tunnel stopped, ignoring exit node update data message")
+		return
+	}
+
+	jsonData, err := json.Marshal(msg.Data)
+	if err != nil {
+		logger.Error("Error marshaling exit node update data: %v", err)
+		return
+	}
+
+	var update ExitNodeUpdateData
+	if err := json.Unmarshal(jsonData, &update); err != nil {
+		logger.Error("Error unmarshaling exit node update data: %v", err)
+		return
+	}
+
+	o.exitNodeMu.Lock()
+	defer o.exitNodeMu.Unlock()
+
+	if o.exitNode == nil {
+		logger.Debug("Ignoring exit node update data message: no exit node connected")
+		return
+	}
+
+	serverIP := net.ParseIP(o.exitNode.ServerIP)
+
+	// Add new aliases BEFORE removing old ones, same as site aliases, so a rename
+	// that keeps the same underlying address never has a gap in resolution.
+	if o.dnsProxy != nil && serverIP != nil {
+		for _, alias := range update.NewAliases {
+			if err := o.dnsProxy.AddDNSRecord(alias, serverIP, exitNodeAliasSiteId); err != nil {
+				logger.Warn("Failed to add DNS record for exit node alias %s: %v", alias, err)
+			}
+		}
+	}
+
+	if o.dnsProxy != nil && serverIP != nil {
+		for _, alias := range update.OldAliases {
+			o.dnsProxy.RemoveDNSRecordForSite(alias, serverIP, exitNodeAliasSiteId)
+		}
+	}
+
+	o.exitNode.Aliases = applyStringListUpdate(o.exitNode.Aliases, update.OldAliases, update.NewAliases)
+
+	logger.Info("Successfully updated exit node data")
+}
+
+// applyStringListUpdate returns list with every entry in removed dropped and every
+// entry in added appended, preserving the add-before-remove semantics of the caller.
+func applyStringListUpdate(list, removed, added []string) []string {
+	next := make([]string, 0, len(list)+len(added))
+	next = append(next, list...)
+	next = append(next, added...)
+
+	removedSet := make(map[string]struct{}, len(removed))
+	for _, alias := range removed {
+		removedSet[alias] = struct{}{}
+	}
+
+	filtered := next[:0]
+	for _, alias := range next {
+		if _, ok := removedSet[alias]; ok {
+			continue
+		}
+		filtered = append(filtered, alias)
+	}
+	return filtered
+}

olm/olm.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"net"
 	"net/http"
@@ -15,6 +16,7 @@ import (
 
 	"github.com/fosrl/newt/bind"
 	"github.com/fosrl/newt/clients/permissions"
+	"github.com/fosrl/newt/exitnode"
 	"github.com/fosrl/newt/holepunch"
 	"github.com/fosrl/newt/logger"
 	"github.com/fosrl/newt/network"
@@ -55,6 +57,11 @@ type Olm struct {
 	holePunchManager   *holepunch.Manager
 	peerManager        *peers.PeerManager
 	peerManagerMu      sync.RWMutex
+
+	// exitNode tracks the currently connected exit node peer, if any. It lives on a
+	// secondary address on the same interface/WireGuard device as the site peers.
+	exitNode   *ExitNodeConfig
+	exitNodeMu sync.Mutex
 	// Power mode management
 	currentPowerMode string
 	powerModeMu      sync.Mutex
@@ -75,6 +82,11 @@ type Olm struct {
 	stopRegister   func()
 	updateRegister func(newData any)
 
+	// Exit node ping dance, run before registration so the server can pick
+	// the best exit node (mirrors newt's newt/ping/request flow).
+	stopPingRequest    func()
+	pendingPingChainId string
+
 	stopPeerSends   map[string]func()
 	stopPeerInits   map[string]func()
 	jitPendingSites map[int]string // siteId -> chainId for in-flight JIT requests
@@ -550,6 +562,66 @@ func (o *Olm) StartTunnel(config TunnelConfig) {
 	o.websocket.RegisterHandler("olm/wg/peer/chain/cancel", o.handleCancelChain)
 	o.websocket.RegisterHandler("olm/sync", o.handleSync)
 
+	// Handlers for the server to direct connecting/disconnecting an exit node after registration
+	o.websocket.RegisterHandler("olm/wg/exitnode/connect", o.handleExitNodeConnect)
+	o.websocket.RegisterHandler("olm/wg/exitnode/disconnect", o.handleExitNodeDisconnect)
+	o.websocket.RegisterHandler("olm/wg/exitnode/data/update", o.handleExitNodeUpdateData)
+
+	o.websocket.RegisterHandler("olm/ping/exitNodes", func(msg websocket.WSMessage) {
+		logger.Debug("Received exit node ping request")
+
+		if o.stopPingRequest != nil {
+			o.stopPingRequest()
+			o.stopPingRequest = nil
+		}
+
+		if !o.tunnelRunning {
+			logger.Debug("Tunnel is no longer running, skipping exit node ping")
+			return
+		}
+
+		var exitNodeData exitnode.ExitNodeData
+		jsonData, err := json.Marshal(msg.Data)
+		if err != nil {
+			logger.Error("Error marshaling exit node data: %v", err)
+			return
+		}
+		if err := json.Unmarshal(jsonData, &exitNodeData); err != nil {
+			logger.Error("Error unmarshaling exit node data: %v", err)
+			return
+		}
+
+		if exitNodeData.ChainId != "" {
+			if exitNodeData.ChainId != o.pendingPingChainId {
+				logger.Debug("Discarding duplicate/stale olm/ping/exitNodes (chainId=%s, expected=%s)", exitNodeData.ChainId, o.pendingPingChainId)
+				return
+			}
+			o.pendingPingChainId = ""
+		}
+
+		if len(exitNodeData.ExitNodes) == 0 {
+			logger.Info("No exit nodes provided")
+			return
+		}
+
+		pingResults := exitnode.PingExitNodes(exitNodeData.ExitNodes, "", false)
+
+		publicKey := o.privateKey.PublicKey()
+		logger.Debug("Sending registration message to server with public key: %s, relay: %v, pingResults: %+v", publicKey, !config.Holepunch, pingResults)
+		o.stopRegister, o.updateRegister = o.websocket.SendMessageInterval("olm/wg/register", map[string]any{
+			"publicKey":   publicKey.String(),
+			"relay":       !config.Holepunch,
+			"olmVersion":  o.olmConfig.Version,
+			"olmAgent":    o.olmConfig.Agent,
+			"orgId":       config.OrgID,
+			"userToken":   userToken,
+			"fingerprint": o.fingerprint,
+			"postures":    o.postures,
+			"pingResults": pingResults,
+			"chainId":     generateChainId(), // use a random chainId for registration updates - it won't be used for cancellation since registration is a one-time message but for tracking the session
+		}, 2*time.Second, 20) // after 18 tries on the server side we send the error so dont change this without changing that
+	})
+
 	o.websocket.OnConnect(func() error {
 		logger.Info("Websocket Connected")
 
@@ -568,8 +640,6 @@ func (o *Olm) StartTunnel(config TunnelConfig) {
 			return nil
 		}
 
-		publicKey := o.privateKey.PublicKey()
-
 		// delay for 500ms to allow for time for the hp to get processed
 		time.Sleep(500 * time.Millisecond)
 
@@ -579,19 +649,36 @@ func (o *Olm) StartTunnel(config TunnelConfig) {
 			return nil
 		}
 
-		if o.stopRegister == nil {
-			logger.Debug("Sending registration message to server with public key: %s and relay: %v", publicKey, !config.Holepunch)
-			o.stopRegister, o.updateRegister = o.websocket.SendMessageInterval("olm/wg/register", map[string]any{
-				"publicKey":   publicKey.String(),
-				"relay":       !config.Holepunch,
-				"olmVersion":  o.olmConfig.Version,
-				"olmAgent":    o.olmConfig.Agent,
-				"orgId":       config.OrgID,
-				"userToken":   userToken,
-				"fingerprint": o.fingerprint,
-				"postures":    o.postures,
-				"chainId":     generateChainId(), // use a random chainId for registration updates - it won't be used for cancellation since registration is a one-time message but for tracking the session
-			}, 2*time.Second, 20) // after 18 tries on the server side we send the error so dont change this without changing that
+		if o.stopRegister == nil && o.stopPingRequest == nil {
+			publicKey := o.privateKey.PublicKey()
+
+			pingChainId := generateChainId()
+			o.pendingPingChainId = pingChainId
+			logger.Debug("Requesting exit nodes from server for ping selection")
+			o.stopPingRequest, _ = o.websocket.SendMessageInterval("olm/ping/request", map[string]any{
+				"chainId": pingChainId,
+			}, 3*time.Second, 10)
+
+			// Backwards-compatible one-shot registration, with no pingResults,
+			// for servers that predate the exit node ping dance. Servers that
+			// support it ignore backwardsCompatible register messages (see
+			// handleOlmRegisterMessage server-side) and wait for the real
+			// registration sent from the olm/ping/exitNodes handler above.
+			bcChainId := generateChainId()
+			if err := o.websocket.SendMessage("olm/wg/register", map[string]any{
+				"publicKey":           publicKey.String(),
+				"relay":               !config.Holepunch,
+				"olmVersion":          o.olmConfig.Version,
+				"olmAgent":            o.olmConfig.Agent,
+				"orgId":               config.OrgID,
+				"userToken":           userToken,
+				"fingerprint":         o.fingerprint,
+				"postures":            o.postures,
+				"backwardsCompatible": true,
+				"chainId":             bcChainId,
+			}); err != nil {
+				logger.Error("Failed to send registration message: %v", err)
+			}
 
 			// Invoke onRegistered callback if configured
 			if o.olmConfig.OnRegistered != nil {
@@ -688,6 +775,12 @@ func (o *Olm) Close() {
 		o.stopRegister = nil
 	}
 
+	if o.stopPingRequest != nil {
+		logger.Debug("Stopping exit node ping request interval")
+		o.stopPingRequest()
+		o.stopPingRequest = nil
+	}
+
 	// Stop all pending peer init and send senders before closing websocket
 	o.peerSendMu.Lock()
 	for _, stop := range o.stopPeerInits {
@@ -744,6 +837,13 @@ func (o *Olm) Close() {
 	}
 	o.peerManagerMu.Unlock()
 
+	// The WireGuard device and TUN interface are being torn down below, which takes
+	// the exit node peer and its secondary address with them - just clear the
+	// in-memory record so a stale config isn't reused on the next connect.
+	o.exitNodeMu.Lock()
+	o.exitNode = nil
+	o.exitNodeMu.Unlock()
+
 	if o.uapiListener != nil {
 		_ = o.uapiListener.Close()
 		o.uapiListener = nil

olm/types.go

@@ -10,6 +10,34 @@ type WgData struct {
 	Sites         []peers.SiteConfig `json:"sites"`
 	TunnelIP      string             `json:"tunnelIP"`
 	UtilitySubnet string             `json:"utilitySubnet"` // this is for things like the DNS server, and alias addresses
+	ExitNode      *ExitNodeConfig    `json:"exitNode,omitempty"`
+}
+
+// ExitNodeConfig describes an exit node the olm client can connect to for
+// resources (e.g. inference) hosted on that node, separate from the site
+// peers. It lives in a different address space than the site tunnel - the
+// client is assigned TunnelIP (within the exit node's subnet) to reach the
+// node at ServerIP. It arrives on the initial "olm/wg/connect" message and can
+// also be sent later via "olm/wg/exitnode/connect" / "olm/wg/exitnode/disconnect"
+// so the server can direct a client to connect/disconnect after registration.
+type ExitNodeConfig struct {
+	Connect   bool     `json:"connect"`
+	Endpoint  string   `json:"endpoint"`
+	RelayPort uint16   `json:"relayPort"`
+	PublicKey string   `json:"publicKey"`
+	ServerIP  string   `json:"serverIP"`
+	TunnelIP  string   `json:"tunnelIP"`
+	Aliases   []string `json:"aliases,omitempty"`
+}
+
+// ExitNodeUpdateData describes a change to data associated with the currently
+// connected exit node, e.g. when a resource's alias is renamed on the server.
+// Aliases have no per-alias address here since every exit node alias resolves
+// to the exit node's own ServerIP. More fields can be added here in the
+// future as other exit node data becomes updatable.
+type ExitNodeUpdateData struct {
+	OldAliases []string `json:"oldAliases,omitempty"`
+	NewAliases []string `json:"newAliases,omitempty"`
 }
 
 type SyncData struct {

websocket/client.go

@@ -22,6 +22,14 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+// writeDeadline bounds how long a websocket write may block before it is
+// treated as a failure. Without this, a write to a TCP connection whose
+// underlying network interface has disappeared (e.g. laptop sleep/resume,
+// Wi-Fi roam) can sit buffered in the kernel for minutes without erroring,
+// which prevents the ping monitor from ever detecting the dead connection
+// and reconnecting.
+const writeDeadline = 10 * time.Second
+
 // AuthError represents an authentication/authorization error (401/403)
 type AuthError struct {
 	StatusCode int
@@ -83,6 +91,7 @@ type Client struct {
 	isDisconnected    bool // Flag to track if client is intentionally disconnected
 	reconnectMux      sync.RWMutex
 	pingInterval      time.Duration
+	pongWait          time.Duration // read deadline window; if no pong/message arrives within it, the connection is considered dead
 	onConnect         func() error
 	onTokenUpdate     func(token string, exitNodes []ExitNode)
 	onAuthError       func(statusCode int, message string) // Callback for auth errors
@@ -167,6 +176,16 @@ func NewClient(ID, secret, userToken, orgId, endpoint string, pingInterval time.
 		OrgID:     orgId,
 	}
 
+	// Read deadline window: must exceed pingInterval so a healthy connection
+	// (which gets a pong/message at least every pingInterval) is never torn
+	// down, but a dead/half-open one — including one where writes keep
+	// "succeeding" because small pings fit in the kernel send buffer even
+	// under total packet loss — is detected within ~2 ping cycles.
+	pongWait := pingInterval * 2
+	if pongWait < 20*time.Second {
+		pongWait = 20 * time.Second
+	}
+
 	client := &Client{
 		config:            config,
 		baseURL:           endpoint, // default value
@@ -175,6 +194,7 @@ func NewClient(ID, secret, userToken, orgId, endpoint string, pingInterval time.
 		reconnectInterval: 3 * time.Second,
 		isConnected:       false,
 		pingInterval:      pingInterval,
+		pongWait:          pongWait,
 		clientType:        "olm",
 		pingDone:          make(chan struct{}),
 	}
@@ -268,6 +288,9 @@ func (c *Client) SendMessage(messageType string, data interface{}) error {
 
 	c.writeMux.Lock()
 	defer c.writeMux.Unlock()
+	if err := c.conn.SetWriteDeadline(time.Now().Add(writeDeadline)); err != nil {
+		return err
+	}
 	return c.conn.WriteJSON(msg)
 }
 
@@ -582,6 +605,18 @@ func (c *Client) establishConnection() error {
 	c.conn = conn
 	c.setConnected(true)
 
+	// Arm a read deadline and refresh it whenever a pong arrives. Combined with
+	// the protocol-level ping sent alongside the app-level one in sendPing,
+	// this detects a dead or half-open connection (e.g. the route disappearing
+	// on sleep/resume, or total packet loss) that a write-side check alone
+	// misses: small periodic pings fit in the kernel send buffer and keep
+	// "succeeding" even when nothing is actually reaching the peer.
+	_ = c.conn.SetReadDeadline(time.Now().Add(c.pongWait))
+	c.conn.SetPongHandler(func(appData string) error {
+		_ = c.conn.SetReadDeadline(time.Now().Add(c.pongWait))
+		return nil
+	})
+
 	// Note: ping monitor is NOT started here - it will be started when
 	// StartPingMonitor() is called after registration completes
 
@@ -697,7 +732,17 @@ func (c *Client) sendPing() {
 	logger.Debug("websocket: Sending ping: %+v", pingMsg)
 
 	c.writeMux.Lock()
-	err := c.conn.WriteJSON(pingMsg)
+	err := c.conn.SetWriteDeadline(time.Now().Add(writeDeadline))
+	if err == nil {
+		err = c.conn.WriteJSON(pingMsg)
+	}
+	if err == nil {
+		// Protocol-level ping: a standards-compliant server replies with a
+		// PONG, which refreshes the read deadline via SetPongHandler. This is
+		// what actually detects a half-open connection where writes still
+		// "succeed" (buffered by the kernel) but nothing is reaching the peer.
+		_ = c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeDeadline))
+	}
 	c.writeMux.Unlock()
 	if err != nil {
 		// Check if we're shutting down before logging error and reconnecting
@@ -803,6 +848,13 @@ func (c *Client) readPumpWithDisconnectDetection() {
 			return
 		default:
 			messageType, p, err := c.conn.ReadMessage()
+			if err == nil {
+				// Any inbound traffic means the peer is alive — extend the
+				// read deadline (also covers servers that answer the
+				// app-level "olm/ping" with a message rather than a
+				// protocol pong).
+				_ = c.conn.SetReadDeadline(time.Now().Add(c.pongWait))
+			}
 			if err != nil {
 				// Check if we're shutting down or explicitly disconnected before logging error
 				select {