Merge pull request #133 from fosrl/dev

1.18.2

device/middle_device.go

@@ -1,6 +1,7 @@
 package device
 
 import (
+	"bytes"
 	"io"
 	"net/netip"
 	"os"
@@ -8,6 +9,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/fosrl/newt/bind"
 	"github.com/fosrl/newt/logger"
 	"golang.zx2c4.com/wireguard/tun"
 )
@@ -24,7 +26,7 @@ type FilterRule struct {
 // closeAwareDevice wraps a tun.Device along with a flag
 // indicating whether its Close method was called.
 type closeAwareDevice struct {
-	isClosed     atomic.Bool
+	isClosed atomic.Bool
 	tun.Device
 	closeEventCh chan struct{}
 	wg           sync.WaitGroup
@@ -423,6 +425,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 +526,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 +570,74 @@ 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 && isMagicPacket(payload)
+}
+
+// IsMagicPacket reports whether payload is one of our connectivity-test magic
+// packets (a MagicTestRequest or MagicTestResponse). These packets are meant to
+// travel directly between physical UDP sockets and must never be encapsulated by
+// WireGuard - e.g. if OS routing mistakenly sends one into a WireGuard TUN
+// interface (because the destination falls inside a routed tunnel subnet), it
+// should be dropped there rather than tunneled, which would otherwise make a
+// LAN-local endpoint test falsely appear to succeed over the tunnel.
+func isMagicPacket(payload []byte) bool {
+	if len(payload) >= bind.MagicTestRequestLen && bytes.HasPrefix(payload, bind.MagicTestRequest) {
+		return true
+	}
+	if len(payload) >= bind.MagicTestResponseLen && bytes.HasPrefix(payload, bind.MagicTestResponse) {
+		return true
+	}
+	return false
+}
+
+// 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 +657,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
@@ -660,4 +728,4 @@ func (d *MiddleDevice) WriteToTun(bufs [][]byte, offset int) (int, error) {
 
 		return n, err
 	}
-}
+}

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)
+	}
+}

peers/manager.go

@@ -58,8 +58,29 @@ type PeerManager struct {
 
 	routeOptimizerStop chan struct{}
 	optimizerTrigger   chan struct{}
+
+	// lastOwnerChange tracks, per allowed-IP CIDR, when ownership was last transferred.
+	// Used to enforce a cooldown so routes don't flap between two similarly-performing sites.
+	lastOwnerChange map[string]time.Time
 }
 
+const (
+	// routeSwitchRTTMargin requires a candidate site's RTT to be at least this much
+	// better (as a fraction) than the current owner's before we consider it worth
+	// switching, so two similarly-performing sites don't flap back and forth.
+	routeSwitchRTTMargin = 0.20 // candidate must be >=20% faster
+
+	// routeSwitchMinAbsMargin is a floor on the RTT improvement required, so the
+	// percentage margin above doesn't become meaningless at very low RTTs (e.g. a
+	// 1ms vs 0.8ms "20% improvement" shouldn't trigger a switch).
+	routeSwitchMinAbsMargin = 5 * time.Millisecond
+
+	// routeSwitchCooldown is the minimum time to wait after transferring ownership
+	// of a route before it can be transferred again, unless the current owner's
+	// connection quality degrades (disconnects or falls back to relay).
+	routeSwitchCooldown = 30 * time.Second
+)
+
 // NewPeerManager creates a new PeerManager with an internal PeerMonitor
 func NewPeerManager(config PeerManagerConfig) *PeerManager {
 	pm := &PeerManager{
@@ -72,6 +93,7 @@ func NewPeerManager(config PeerManagerConfig) *PeerManager {
 		allowedIPClaims: make(map[string]map[int]bool),
 		APIServer:       config.APIServer,
 		publicDNS:       config.PublicDNS,
+		lastOwnerChange: make(map[string]time.Time),
 	}
 
 	// Create the peer monitor
@@ -515,6 +537,7 @@ func (pm *PeerManager) releaseAllowedIP(siteId int, cidr string) (newOwner int,
 		delete(claims, siteId)
 		if len(claims) == 0 {
 			delete(pm.allowedIPClaims, cidr)
+			delete(pm.lastOwnerChange, cidr)
 		}
 	}
 
@@ -1114,6 +1137,49 @@ func (pm *PeerManager) selectBestOwner(claims map[int]bool) int {
 	return bestSiteId
 }
 
+// shouldSwitchOwner decides whether ownership of cidr should move from the current
+// owner to the candidate. It applies hysteresis so two sites with roughly equal
+// performance don't flap back and forth:
+//   - A switch driven by connectivity class (connected vs not, direct vs relayed) is
+//     always allowed immediately - these are correctness issues, not noise.
+//   - A switch driven purely by RTT requires both a minimum improvement margin and
+//     that the cooldown since the last switch of this route has elapsed.
+//
+// Must be called with pm.mu held.
+func (pm *PeerManager) shouldSwitchOwner(cidr string, currentSiteId, candidateSiteId int) bool {
+	curConn, curRelay, curRTT := pm.peerMonitor.GetConnectionQuality(currentSiteId)
+	candConn, candRelay, candRTT := pm.peerMonitor.GetConnectionQuality(candidateSiteId)
+
+	// Connectivity-class differences (up/down, direct/relayed) are not subject to
+	// hysteresis - always act on them so we don't stay stuck on a broken route.
+	if curConn != candConn || curRelay != candRelay {
+		return true
+	}
+	if !curConn {
+		return false // both down, nothing to do
+	}
+
+	// Same connectivity class: only switch on a meaningful, sustained RTT win.
+	if candRTT == 0 || curRTT == 0 {
+		return false
+	}
+	minImprovement := time.Duration(float64(curRTT) * routeSwitchRTTMargin)
+	if minImprovement < routeSwitchMinAbsMargin {
+		minImprovement = routeSwitchMinAbsMargin
+	}
+	if candRTT > curRTT-minImprovement {
+		return false // not enough of an improvement to be worth switching
+	}
+
+	if lastChange, ok := pm.lastOwnerChange[cidr]; ok {
+		if time.Since(lastChange) < routeSwitchCooldown {
+			return false // switched too recently, avoid flapping
+		}
+	}
+
+	return true
+}
+
 // getWireGuardAllowedIPs returns the full set of IPs that should be in WireGuard
 // for a peer: server IP /32 plus all shared IPs it currently owns.
 // Must be called with pm.mu held.
@@ -1181,6 +1247,7 @@ func (pm *PeerManager) optimizeRoutes() {
 		if !hasOwner {
 			// No current owner, just assign
 			pm.allowedIPOwners[cidr] = bestOwner
+			pm.lastOwnerChange[cidr] = time.Now()
 			if toPeer, exists := pm.peers[bestOwner]; exists {
 				if err := AddAllowedIP(pm.device, toPeer.PublicKey, cidr); err != nil {
 					logger.Error("Failed to assign IP %s to site %d: %v", cidr, bestOwner, err)
@@ -1189,10 +1256,16 @@ func (pm *PeerManager) optimizeRoutes() {
 			continue
 		}
 
+		if !pm.shouldSwitchOwner(cidr, currentOwner, bestOwner) {
+			continue // Not a big enough or sustained enough improvement, avoid flapping
+		}
+
 		logger.Info("Route optimizer: moving %s from site %d to site %d", cidr, currentOwner, bestOwner)
 		if err := pm.transferOwnership(cidr, currentOwner, bestOwner); err != nil {
 			logger.Error("Failed to transfer ownership of %s from site %d to site %d: %v",
 				cidr, currentOwner, bestOwner, err)
+		} else {
+			pm.lastOwnerChange[cidr] = time.Now()
 		}
 	}
 }

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 {