[client] make AllowedIPs swap self-healing on WireGuard errors

Address CodeRabbit review on #6799.

Decrement now decides whether to reprogram based on whether the currently
installed peer still holds references (e.peers[e.active] > 0) instead of
whether the released peer was the active one. If a prior swap's remove/add
failed, e.active was left pointing at a peer with zero references and every
later Decrement returned early, permanently stranding the prefix and leaking
the entry. The active peer is now detached only when e.active != "", and a
failed hand-off is retried on the next Decrement/Increment.

Also clear currentPeerKey unconditionally in the static handler's
RemoveAllowedIPs (matching dynamic/dnsinterceptor) and assert Increment
errors in the tests. Added self-heal tests for the failed remove/add paths.

client/android/client.go

@@ -247,9 +247,6 @@ func (c *Client) DebugBundle(platformFiles PlatformFiles, anonymize bool) (strin
 		deps.SyncResponse = resp
 
 		if e := cc.Engine(); e != nil {
-			deps.RefreshStatus = func() {
-				e.RunHealthProbes(context.Background(), true)
-			}
 			if cm := e.GetClientMetrics(); cm != nil {
 				deps.ClientMetrics = cm
 			}

client/firewall/uspfilter/filter.go

@@ -121,7 +121,6 @@ type Manager struct {
 	udpTracker     *conntrack.UDPTracker
 	icmpTracker    *conntrack.ICMPTracker
 	tcpTracker     *conntrack.TCPTracker
-	fragments      *fragmentTracker
 	forwarder      atomic.Pointer[forwarder.Forwarder]
 	pendingCapture atomic.Pointer[forwarder.PacketCapture]
 	logger         *nblog.Logger
@@ -184,41 +183,6 @@ func (d *decoder) decodePacket(data []byte) error {
 	}
 }
 
-// decodeTransport decodes the transport header of a first fragment (which
-// gopacket leaves undecoded) into the decoder and appends its layer type to
-// decoded, so the ACL pipeline can evaluate it like a normal packet. It returns
-// false if the protocol is unsupported or the header is truncated.
-func (d *decoder) decodeTransport(proto layers.IPProtocol, payload []byte) bool {
-	var l4 gopacket.DecodingLayer
-	var layerType gopacket.LayerType
-	var minLen int
-	switch proto {
-	case layers.IPProtocolTCP:
-		l4, layerType, minLen = &d.tcp, layers.LayerTypeTCP, 20
-	case layers.IPProtocolUDP:
-		l4, layerType, minLen = &d.udp, layers.LayerTypeUDP, 8
-	case layers.IPProtocolICMPv4:
-		l4, layerType, minLen = &d.icmp4, layers.LayerTypeICMPv4, 8
-	case layers.IPProtocolICMPv6:
-		l4, layerType, minLen = &d.icmp6, layers.LayerTypeICMPv6, 8
-	default:
-		return false
-	}
-
-	// Reject a fragment too small to hold the full transport header before
-	// decoding: it can't be ACL-evaluated (tiny-fragment attack), and skipping
-	// the decode avoids gopacket allocating an error on the drop path.
-	if len(payload) < minLen {
-		return false
-	}
-
-	if err := l4.DecodeFromBytes(payload, gopacket.NilDecodeFeedback); err != nil {
-		return false
-	}
-	d.decoded = append(d.decoded, layerType)
-	return true
-}
-
 // Create userspace firewall manager constructor
 func Create(iface common.IFaceMapper, disableServerRoutes bool, flowLogger nftypes.FlowLogger, mtu uint16) (*Manager, error) {
 	return create(iface, nil, disableServerRoutes, flowLogger, mtu)
@@ -322,8 +286,6 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 	if err := m.localipmanager.UpdateLocalIPs(iface); err != nil {
 		return nil, fmt.Errorf("update local IPs: %w", err)
 	}
-	m.fragments = newFragmentTracker(m.logger)
-
 	if disableConntrack {
 		log.Info("conntrack is disabled")
 	} else {
@@ -337,7 +299,6 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 		}
 	}
 	if err := iface.SetFilter(m); err != nil {
-		m.fragments.Close()
 		return nil, fmt.Errorf("set filter: %w", err)
 	}
 	return m, nil
@@ -733,10 +694,6 @@ func (m *Manager) resetState() {
 		m.tcpTracker.Close()
 	}
 
-	if m.fragments != nil {
-		m.fragments.Close()
-	}
-
 	if fwder := m.forwarder.Load(); fwder != nil {
 		fwder.SetCapture(nil)
 		fwder.Stop()
@@ -1089,20 +1046,19 @@ func (m *Manager) filterInbound(packetData []byte, size int) bool {
 		return true
 	}
 
-	// gopacket does not decode the transport header of any IP fragment, so
-	// fragments take a dedicated path: the first fragment's header is decoded
-	// and ACL-evaluated here, and the remaining fragments inherit its verdict.
+	// TODO: pass fragments of routed packets to forwarder
 	if fragment {
-		return m.filterInboundFragment(d, srcIP, dstIP, size)
+		if m.logger.Enabled(nblog.LevelTrace) {
+			if d.decoded[0] == layers.LayerTypeIPv4 {
+				m.logger.Trace4("packet is a fragment: src=%v dst=%v id=%v flags=%v",
+					srcIP, dstIP, d.ip4.Id, d.ip4.Flags)
+			} else {
+				m.logger.Trace2("packet is an IPv6 fragment: src=%v dst=%v", srcIP, dstIP)
+			}
+		}
+		return false
 	}
 
-	return m.filterInboundDecoded(d, srcIP, dstIP, packetData, size)
-}
-
-// filterInboundDecoded runs the ACL, DNAT and conntrack pipeline on a fully
-// decoded (non-fragment) inbound packet. It returns true if the packet should
-// be dropped.
-func (m *Manager) filterInboundDecoded(d *decoder, srcIP, dstIP netip.Addr, packetData []byte, size int) bool {
 	// TODO: optimize port DNAT by caching matched rules in conntrack
 	if translated := m.translateInboundPortDNAT(packetData, d, srcIP, dstIP); translated {
 		// Re-decode after port DNAT translation to update port information
@@ -1133,226 +1089,33 @@ func (m *Manager) filterInboundDecoded(d *decoder, srcIP, dstIP netip.Addr, pack
 	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData, size)
 }
 
-// fragmentMeta holds the reassembly identity and layout of an IP fragment,
-// extracted uniformly for IPv4 and IPv6.
-type fragmentMeta struct {
-	key fragmentKey
-	// offset is the fragment offset in 8-byte units (zero for the first
-	// fragment).
-	offset uint16
-	// moreFragments is the More Fragments bit. A first fragment with it unset is
-	// an IPv6 atomic fragment (a complete datagram, RFC 6946): it has no trailing
-	// fragments to inherit a verdict, so it must not be recorded.
-	moreFragments bool
-	proto         layers.IPProtocol
-	// l4payload is the fragmentable payload of this fragment. For the first
-	// fragment it starts with the transport header.
-	l4payload []byte
-	// headerEndOctets is the first fragment's payload length in 8-byte units:
-	// the smallest offset a trailing fragment may start at without overlapping
-	// the inspected transport header.
-	headerEndOctets uint16
-}
-
-// fragmentMetadata extracts the fragment identity and layout from a decoded IP
-// fragment. It returns false for fragments it can't interpret (e.g. an IPv6
-// fragment header shorter than 8 bytes), which are then dropped.
-func fragmentMetadata(d *decoder, srcIP, dstIP netip.Addr) (fragmentMeta, bool) {
-	switch d.decoded[0] {
-	case layers.LayerTypeIPv4:
-		payload := d.ip4.Payload
-		return fragmentMeta{
-			key:             fragmentKey{srcIP: srcIP, dstIP: dstIP, id: uint32(d.ip4.Id), proto: uint8(d.ip4.Protocol)},
-			offset:          d.ip4.FragOffset,
-			moreFragments:   d.ip4.Flags&layers.IPv4MoreFragments != 0,
-			proto:           d.ip4.Protocol,
-			l4payload:       payload,
-			headerEndOctets: octets(len(payload)),
-		}, true
-
-	case layers.LayerTypeIPv6:
-		// IPv6 fragment extension header: 8 bytes, followed by the fragmentable
-		// payload. Layout: next header (1), reserved (1), offset+flags (2), id (4).
-		payload := d.ip6.Payload
-		if len(payload) < 8 {
-			return fragmentMeta{}, false
-		}
-		nextHeader := layers.IPProtocol(payload[0])
-		offsetFlags := binary.BigEndian.Uint16(payload[2:4])
-		id := binary.BigEndian.Uint32(payload[4:8])
-		l4 := payload[8:]
-		return fragmentMeta{
-			key:             fragmentKey{srcIP: srcIP, dstIP: dstIP, id: id, proto: uint8(nextHeader)},
-			offset:          offsetFlags >> 3,
-			moreFragments:   offsetFlags&1 != 0,
-			proto:           nextHeader,
-			l4payload:       l4,
-			headerEndOctets: octets(len(l4)),
-		}, true
-
-	default:
-		return fragmentMeta{}, false
-	}
-}
-
-// octets rounds a byte length up to whole 8-byte units, the granularity of the
-// IP fragment offset field.
-func octets(nbytes int) uint16 {
-	return uint16((nbytes + 7) / 8)
-}
-
-// filterInboundFragment decides the fate of an IP fragment. gopacket stops
-// decoding at the network layer for every fragment, so the first fragment's
-// transport header is decoded and ACL-evaluated here and its verdict recorded;
-// the remaining (headerless) fragments inherit that verdict. Anything that
-// cannot be tied to an allowed, non-overlapping first fragment is dropped.
-func (m *Manager) filterInboundFragment(d *decoder, srcIP, dstIP netip.Addr, size int) bool {
-	meta, ok := fragmentMetadata(d, srcIP, dstIP)
-	if !ok {
-		if m.logger.Enabled(nblog.LevelTrace) {
-			m.logger.Trace2("dropping unsupported fragment: src=%v dst=%v", srcIP, dstIP)
-		}
-		return true
-	}
-
-	if meta.offset != 0 {
-		return m.filterTrailingFragment(meta, srcIP, dstIP)
-	}
-
-	// A new first fragment supersedes any recorded verdict for this datagram, so
-	// a re-sent or overlapping offset-zero fragment can't inherit the old one.
-	m.fragments.poison(meta.key)
-
-	// First fragment: decode its transport header so the ACL can evaluate it. A
-	// decode failure means the fragment is too small to hold the full transport
-	// header (RFC 1858 §3 tiny-fragment attack); it can't be evaluated, so drop it.
-	if !d.decodeTransport(meta.proto, meta.l4payload) {
-		if m.logger.Enabled(nblog.LevelTrace) {
-			m.logger.Trace3("dropping first fragment without full L4 header: src=%v dst=%v id=%v",
-				srcIP, dstIP, meta.key.id)
-		}
-		return true
-	}
-
-	return m.filterFirstFragment(d, meta, srcIP, dstIP, size)
-}
-
-// filterTrailingFragment applies a recorded first-fragment verdict to a
-// non-first fragment.
-func (m *Manager) filterTrailingFragment(meta fragmentMeta, srcIP, dstIP netip.Addr) bool {
-	switch m.fragments.verdict(meta.key, meta.offset) {
-	case fragmentAllow:
-		return false
-	case fragmentOverlap:
-		if m.logger.Enabled(nblog.LevelTrace) {
-			m.logger.Trace3("dropping overlapping fragment rewriting inspected header: src=%v dst=%v id=%v",
-				srcIP, dstIP, meta.key.id)
-		}
-		return true
-	default:
-		if m.logger.Enabled(nblog.LevelTrace) {
-			m.logger.Trace3("dropping fragment with no allowed first fragment: src=%v dst=%v id=%v",
-				srcIP, dstIP, meta.key.id)
-		}
-		return true
-	}
-}
-
-// filterFirstFragment runs the verdict part of the inbound pipeline on a first
-// fragment with its transport header decoded. It mirrors filterInboundDecoded
-// but skips DNAT (port rewriting on fragments is unsupported) and forwarder
-// injection (fragments are left to the stack to reassemble, not forwarded).
-// Allowed fragments have their verdict recorded so the datagram's trailing
-// fragments inherit it.
-func (m *Manager) filterFirstFragment(d *decoder, meta fragmentMeta, srcIP, dstIP netip.Addr, size int) bool {
-	if m.stateful && m.isValidTrackedConnection(d, srcIP, dstIP, size) {
-		m.recordFirstFragment(meta)
-		return false
-	}
-
-	if m.localipmanager.IsLocalIP(dstIP) {
-		ruleID, blocked := m.peerACLsBlock(srcIP, d, nil)
-		if blocked {
-			m.storeDropFlow("Dropping local first fragment (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
-				d, srcIP, dstIP, ruleID, size)
-			return true
-		}
-		m.trackInbound(d, srcIP, dstIP, ruleID, size)
-		m.recordFirstFragment(meta)
-		return false
-	}
-
-	if !m.routingEnabled.Load() {
-		if m.logger.Enabled(nblog.LevelTrace) {
-			m.logger.Trace2("Dropping routed fragment (routing disabled): src=%s dst=%s", srcIP, dstIP)
-		}
-		return true
-	}
-	if m.nativeRouter.Load() {
-		m.trackInbound(d, srcIP, dstIP, nil, size)
-		m.recordFirstFragment(meta)
-		return false
-	}
-
-	// TODO: pass fragments of routed packets to the forwarder; until then
-	// allowed routed fragments go to the native stack.
-	srcPort, dstPort := getPortsFromPacket(d)
-	ruleID, pass := m.routeACLsPass(srcIP, dstIP, d.decoded[1], srcPort, dstPort)
-	if !pass {
-		m.storeDropFlow("Dropping routed first fragment (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
-			d, srcIP, dstIP, ruleID, size)
-		return true
-	}
-
-	m.recordFirstFragment(meta)
-	return false
-}
-
-// recordFirstFragment caches an allowed first fragment's verdict for its
-// trailing fragments to inherit. Atomic fragments (no More Fragments bit) are
-// complete datagrams with no trailing fragments, so they are not cached and
-// cannot exhaust the verdict table.
-func (m *Manager) recordFirstFragment(meta fragmentMeta) {
-	if !meta.moreFragments {
-		return
-	}
-	m.fragments.recordAllowed(meta.key, meta.headerEndOctets)
-}
-
-// storeDropFlow logs and records a netflow drop event for an inbound packet
-// denied by the ACLs. msg is the trace format taking rule id, protocol, source
-// and destination.
-func (m *Manager) storeDropFlow(msg string, d *decoder, srcIP, dstIP netip.Addr, ruleID []byte, size int) {
-	pnum := getProtocolFromPacket(d)
-	srcPort, dstPort := getPortsFromPacket(d)
-
-	if m.logger.Enabled(nblog.LevelTrace) {
-		m.logger.Trace6(msg, ruleID, pnum, srcIP, srcPort, dstIP, dstPort)
-	}
-
-	m.flowLogger.StoreEvent(nftypes.EventFields{
-		FlowID:     uuid.New(),
-		Type:       nftypes.TypeDrop,
-		RuleID:     ruleID,
-		Direction:  nftypes.Ingress,
-		Protocol:   pnum,
-		SourceIP:   srcIP,
-		DestIP:     dstIP,
-		SourcePort: srcPort,
-		DestPort:   dstPort,
-		// TODO: icmp type/code
-		RxPackets: 1,
-		RxBytes:   uint64(size),
-	})
-}
-
 // handleLocalTraffic handles local traffic.
 // If it returns true, the packet should be dropped.
 func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP netip.Addr, packetData []byte, size int) bool {
 	ruleID, blocked := m.peerACLsBlock(srcIP, d, packetData)
 	if blocked {
-		m.storeDropFlow("Dropping local packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
-			d, srcIP, dstIP, ruleID, size)
+		pnum := getProtocolFromPacket(d)
+		srcPort, dstPort := getPortsFromPacket(d)
+
+		if m.logger.Enabled(nblog.LevelTrace) {
+			m.logger.Trace6("Dropping local packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
+				ruleID, pnum, srcIP, srcPort, dstIP, dstPort)
+		}
+
+		m.flowLogger.StoreEvent(nftypes.EventFields{
+			FlowID:     uuid.New(),
+			Type:       nftypes.TypeDrop,
+			RuleID:     ruleID,
+			Direction:  nftypes.Ingress,
+			Protocol:   pnum,
+			SourceIP:   srcIP,
+			DestIP:     dstIP,
+			SourcePort: srcPort,
+			DestPort:   dstPort,
+			// TODO: icmp type/code
+			RxPackets: 1,
+			RxBytes:   uint64(size),
+		})
 		return true
 	}
 
@@ -1405,8 +1168,27 @@ func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packe
 
 	ruleID, pass := m.routeACLsPass(srcIP, dstIP, protoLayer, srcPort, dstPort)
 	if !pass {
-		m.storeDropFlow("Dropping routed packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
-			d, srcIP, dstIP, ruleID, size)
+		proto := getProtocolFromPacket(d)
+
+		if m.logger.Enabled(nblog.LevelTrace) {
+			m.logger.Trace6("Dropping routed packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
+				ruleID, proto, srcIP, srcPort, dstIP, dstPort)
+		}
+
+		m.flowLogger.StoreEvent(nftypes.EventFields{
+			FlowID:     uuid.New(),
+			Type:       nftypes.TypeDrop,
+			RuleID:     ruleID,
+			Direction:  nftypes.Ingress,
+			Protocol:   proto,
+			SourceIP:   srcIP,
+			DestIP:     dstIP,
+			SourcePort: srcPort,
+			DestPort:   dstPort,
+			// TODO: icmp type/code
+			RxPackets: 1,
+			RxBytes:   uint64(size),
+		})
 		return true
 	}
 

client/firewall/uspfilter/forwarder/forwarder.go

@@ -5,9 +5,7 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
-	"os"
 	"runtime"
-	"strconv"
 	"sync"
 	"time"
 
@@ -33,11 +31,6 @@ const (
 	defaultMaxInFlight   = 1024
 	iosReceiveWindow     = 16384
 	iosMaxInFlight       = 256
-
-	// envForceTCPRACK overrides the platform default for gVisor's RACK loss
-	// detection. Set to a truthy value to force RACK on, or a falsy value to
-	// force it off, on any platform.
-	envForceTCPRACK = "NB_FORCE_TCP_RACK"
 )
 
 type Forwarder struct {
@@ -159,8 +152,6 @@ func New(iface common.IFaceMapper, logger *nblog.Logger, flowLogger nftypes.Flow
 		maxInFlight = iosMaxInFlight
 	}
 
-	configureTCPRecovery(s)
-
 	tcpForwarder := tcp.NewForwarder(s, receiveWindow, maxInFlight, f.handleTCP)
 	s.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
 
@@ -475,31 +466,3 @@ func probeRawICMP(network, addr string, logger *nblog.Logger) bool {
 	logger.Debug1("forwarder: raw %s socket access available", network)
 	return true
 }
-
-// configureTCPRecovery disables gVisor's RACK loss detection on Windows, where
-// it interacts poorly with the host and collapses throughput on routed TCP
-// connections (gVisor issue #9778). Other platforms keep the default. The
-// EnvForceTCPRACK environment variable overrides the platform default.
-func configureTCPRecovery(s *stack.Stack) {
-	disableRACK := runtime.GOOS == "windows"
-
-	if val := os.Getenv(envForceTCPRACK); val != "" {
-		force, err := strconv.ParseBool(val)
-		if err != nil {
-			log.Warnf("parse %s: %v", envForceTCPRACK, err)
-		} else {
-			disableRACK = !force
-		}
-	}
-
-	if !disableRACK {
-		return
-	}
-
-	opt := tcpip.TCPRecovery(0)
-	if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &opt); err != nil {
-		log.Warnf("disable TCP RACK loss detection: %v", err)
-		return
-	}
-	log.Info("forwarder: TCP RACK loss detection disabled")
-}

client/firewall/uspfilter/fragment.go

@@ -1,204 +0,0 @@
-package uspfilter
-
-import (
-	"context"
-	"net/netip"
-	"os"
-	"strconv"
-	"sync"
-	"time"
-
-	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
-)
-
-const (
-	// defaultFragmentTimeout bounds how long a first-fragment verdict is kept
-	// while the remaining fragments arrive. It mirrors the Linux IP reassembly
-	// timeout (net.ipv4.ipfrag_time).
-	defaultFragmentTimeout = 30 * time.Second
-	// fragmentCleanupInterval is how often expired verdicts are purged.
-	fragmentCleanupInterval = 10 * time.Second
-	// defaultMaxFragmentEntries caps the number of concurrently tracked
-	// fragmented datagrams. The table stays bounded because each datagram is a
-	// single small entry regardless of how many fragments it is split into, and
-	// the 13-bit IPv4 fragment-offset field limits any datagram to 64 KiB.
-	defaultMaxFragmentEntries = 16384
-
-	// EnvFragmentMaxEntries overrides defaultMaxFragmentEntries.
-	EnvFragmentMaxEntries = "NB_FRAGMENT_MAX_ENTRIES"
-)
-
-// fragmentVerdict is the decision for a trailing (headerless) fragment.
-type fragmentVerdict int
-
-const (
-	// fragmentDeny drops the fragment: no allowed first fragment is on record.
-	fragmentDeny fragmentVerdict = iota
-	// fragmentAllow passes the fragment: it belongs to an allowed datagram and
-	// does not overlap the already-inspected transport header.
-	fragmentAllow
-	// fragmentOverlap drops the fragment and poisons its datagram: it overlaps
-	// the transport header the ACL inspected (RFC 1858 §4, RFC 3128; RFC 5722
-	// requires discarding the whole datagram on overlap for IPv6).
-	fragmentOverlap
-)
-
-// fragmentKey identifies a fragmented datagram. It matches the RFC 791 / RFC
-// 8200 reassembly key: source, destination, protocol and identification. The id
-// is 32-bit to hold both the IPv4 (16-bit) and IPv6 (32-bit) identification.
-type fragmentKey struct {
-	srcIP netip.Addr
-	dstIP netip.Addr
-	id    uint32
-	proto uint8
-}
-
-// fragmentEntry records the verdict of an allowed first fragment.
-type fragmentEntry struct {
-	// headerEndOctets is the offset, in 8-byte units, at which the first
-	// fragment's payload ended. A trailing fragment starting before this
-	// overlaps bytes the ACL already inspected and is rejected.
-	headerEndOctets uint16
-	// recordedAt is when the first fragment was accepted. The verdict expires a
-	// fixed timeout later and is not refreshed, mirroring the kernel reassembly
-	// timer so a trailing-fragment flood can't keep a datagram alive.
-	recordedAt time.Time
-}
-
-// fragmentTracker records the ACL verdict of a datagram's first fragment so the
-// remaining fragments, which carry no L4 header, can inherit the decision
-// without reassembling the datagram. Only allowed first fragments are stored;
-// anything that cannot be tied to an allowed, non-overlapping first fragment is
-// dropped (fail closed).
-type fragmentTracker struct {
-	logger  *nblog.Logger
-	mutex   sync.Mutex
-	entries map[fragmentKey]fragmentEntry
-	timeout time.Duration
-	// maxEntries caps the table; atCapacity dedups the capacity warning until
-	// the table drains below the cap again.
-	maxEntries    int
-	atCapacity    bool
-	cleanupTicker *time.Ticker
-	cancel        context.CancelFunc
-}
-
-func newFragmentTracker(logger *nblog.Logger) *fragmentTracker {
-	ctx, cancel := context.WithCancel(context.Background())
-	t := &fragmentTracker{
-		logger:        logger,
-		entries:       make(map[fragmentKey]fragmentEntry),
-		timeout:       defaultFragmentTimeout,
-		maxEntries:    fragmentMaxEntries(logger),
-		cleanupTicker: time.NewTicker(fragmentCleanupInterval),
-		cancel:        cancel,
-	}
-	go t.cleanupRoutine(ctx)
-	return t
-}
-
-func fragmentMaxEntries(logger *nblog.Logger) int {
-	v := os.Getenv(EnvFragmentMaxEntries)
-	if v == "" {
-		return defaultMaxFragmentEntries
-	}
-	n, err := strconv.Atoi(v)
-	if err != nil || n <= 0 {
-		logger.Warn2("invalid %s=%q, using default", EnvFragmentMaxEntries, v)
-		return defaultMaxFragmentEntries
-	}
-	return n
-}
-
-// recordAllowed stores the verdict of an allowed first fragment. headerEndOctets
-// is the first fragment's payload length in 8-byte units. When the table is full
-// the record is dropped, which fails closed: the datagram's trailing fragments
-// will be denied.
-func (t *fragmentTracker) recordAllowed(key fragmentKey, headerEndOctets uint16) {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-
-	if t.entries == nil {
-		return
-	}
-	if _, ok := t.entries[key]; !ok && len(t.entries) >= t.maxEntries {
-		if !t.atCapacity {
-			t.atCapacity = true
-			t.logger.Warn2("fragment verdict table at capacity (%d/%d): trailing fragments of new datagrams will be dropped",
-				len(t.entries), t.maxEntries)
-		}
-		return
-	}
-	t.entries[key] = fragmentEntry{
-		headerEndOctets: headerEndOctets,
-		recordedAt:      time.Now(),
-	}
-}
-
-// poison drops any recorded verdict for a datagram, so its later fragments are
-// denied until a new allowed first fragment is recorded. Called on every
-// offset-zero fragment to defeat offset-zero overlap rewrites (RFC 3128).
-func (t *fragmentTracker) poison(key fragmentKey) {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	delete(t.entries, key)
-}
-
-// verdict decides the fate of a trailing fragment at fragOffsetOctets (the IPv4
-// fragment offset, in 8-byte units). A fragment overlapping the inspected
-// header poisons the datagram: the entry is removed so all further fragments of
-// that datagram are denied too.
-func (t *fragmentTracker) verdict(key fragmentKey, fragOffsetOctets uint16) fragmentVerdict {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-
-	entry, ok := t.entries[key]
-	if !ok {
-		return fragmentDeny
-	}
-	if time.Since(entry.recordedAt) > t.timeout {
-		delete(t.entries, key)
-		return fragmentDeny
-	}
-	if fragOffsetOctets < entry.headerEndOctets {
-		delete(t.entries, key)
-		return fragmentOverlap
-	}
-	return fragmentAllow
-}
-
-func (t *fragmentTracker) cleanupRoutine(ctx context.Context) {
-	defer t.cleanupTicker.Stop()
-	for {
-		select {
-		case <-t.cleanupTicker.C:
-			t.cleanup()
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
-func (t *fragmentTracker) cleanup() {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-
-	for key, entry := range t.entries {
-		if time.Since(entry.recordedAt) > t.timeout {
-			delete(t.entries, key)
-		}
-	}
-
-	if len(t.entries) < t.maxEntries {
-		t.atCapacity = false
-	}
-}
-
-// Close stops the cleanup routine and releases resources.
-func (t *fragmentTracker) Close() {
-	t.cancel()
-
-	t.mutex.Lock()
-	t.entries = nil
-	t.mutex.Unlock()
-}

client/firewall/uspfilter/fragment_bench_test.go

@@ -1,115 +0,0 @@
-package uspfilter
-
-import (
-	"encoding/binary"
-	"testing"
-)
-
-// benchFilterInbound drives filterInbound over a fixed packet in a tight loop.
-// Packets are built once, outside the timed region, so the benchmark measures
-// only pipeline cost, which is what an attacker can amplify.
-func benchFilterInbound(b *testing.B, pkt []byte) {
-	b.Helper()
-	b.ReportAllocs()
-	b.SetBytes(int64(len(pkt)))
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		m := benchManager
-		m.filterInbound(pkt, len(pkt))
-	}
-}
-
-// benchManager is a package-level manager reused across fragment benchmarks so
-// setup cost stays out of the timed region.
-var benchManager *Manager
-
-func setupBenchManager(b *testing.B) *Manager {
-	b.Helper()
-	m := newFragmentTestManager(b)
-	allowUDP(b, m, 8080)
-	// Disable conntrack so the allowed-first-fragment path measures transport
-	// decode + ACL every iteration instead of matching the connection tracked
-	// on the first iteration.
-	m.stateful = false
-	benchManager = m
-	return m
-}
-
-// BenchmarkInbound_NormalPacket is the baseline: a full, non-fragmented UDP
-// packet that passes the ACL. Fragment paths should stay comparable to this.
-func BenchmarkInbound_NormalPacket(b *testing.B) {
-	setupBenchManager(b)
-	pkt := normalUDPPacket(b, 8080, 32)
-	benchFilterInbound(b, pkt)
-}
-
-// BenchmarkInbound_FirstFragmentAllowed measures the first-fragment path:
-// transport decode + ACL evaluation + verdict record.
-func BenchmarkInbound_FirstFragmentAllowed(b *testing.B) {
-	setupBenchManager(b)
-	pkt := firstFragmentUDP(b, 0x2000, 8080, 32)
-	benchFilterInbound(b, pkt)
-}
-
-// BenchmarkInbound_TrailingFragmentAllowed measures the common trailing-fragment
-// path: a single map lookup after the first fragment is on record.
-func BenchmarkInbound_TrailingFragmentAllowed(b *testing.B) {
-	m := setupBenchManager(b)
-	first := firstFragmentUDP(b, 0x3000, 8080, 32)
-	m.filterInbound(first, len(first))
-	pkt := trailingFragment(b, 0x3000, 5, false, 24)
-	benchFilterInbound(b, pkt)
-}
-
-// BenchmarkInbound_TrailingFragmentNoFirst is the primary DoS vector: an
-// attacker floods trailing fragments with no first fragment on record. Each is
-// a map miss and must be cheap.
-func BenchmarkInbound_TrailingFragmentNoFirst(b *testing.B) {
-	setupBenchManager(b)
-	pkt := trailingFragment(b, 0x4000, 185, false, 40)
-	benchFilterInbound(b, pkt)
-}
-
-// BenchmarkInbound_TinyFirstFragment measures the tiny-fragment drop path: a
-// first fragment too small to decode a transport header.
-func BenchmarkInbound_TinyFirstFragment(b *testing.B) {
-	setupBenchManager(b)
-	pkt := trailingFragment(b, 0x5000, 0, true, 4)
-	benchFilterInbound(b, pkt)
-}
-
-// BenchmarkInbound_TrailingFragmentDistinctIDs is the worst case for the
-// verdict table: an attacker varies the datagram id on every packet so no first
-// fragment ever matches. Verdict lookups always miss and nothing is recorded,
-// so the table cannot grow. Each iteration rewrites the id field in place.
-func BenchmarkInbound_TrailingFragmentDistinctIDs(b *testing.B) {
-	setupBenchManager(b)
-	pkt := trailingFragment(b, 0x6000, 185, false, 40)
-	m := benchManager
-
-	b.ReportAllocs()
-	b.SetBytes(int64(len(pkt)))
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		// IPv4 identification field is at bytes 4:6.
-		binary.BigEndian.PutUint16(pkt[4:6], uint16(i))
-		m.filterInbound(pkt, len(pkt))
-	}
-}
-
-// BenchmarkInbound_FirstFragmentDistinctIDs measures sustained first-fragment
-// pressure with distinct ids: transport decode + ACL + verdict insert until the
-// table caps, exercising the map growth and capacity guard.
-func BenchmarkInbound_FirstFragmentDistinctIDs(b *testing.B) {
-	setupBenchManager(b)
-	pkt := firstFragmentUDP(b, 0x7000, 8080, 32)
-	m := benchManager
-
-	b.ReportAllocs()
-	b.SetBytes(int64(len(pkt)))
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		binary.BigEndian.PutUint16(pkt[4:6], uint16(i))
-		m.filterInbound(pkt, len(pkt))
-	}
-}

client/firewall/uspfilter/fragment_test.go

@@ -1,554 +0,0 @@
-package uspfilter
-
-import (
-	"encoding/binary"
-	"net"
-	"net/netip"
-	"testing"
-	"time"
-
-	"github.com/google/gopacket"
-	"github.com/google/gopacket/layers"
-	"github.com/stretchr/testify/require"
-
-	fw "github.com/netbirdio/netbird/client/firewall/manager"
-	nbiface "github.com/netbirdio/netbird/client/iface"
-	"github.com/netbirdio/netbird/client/iface/device"
-	"github.com/netbirdio/netbird/client/iface/wgaddr"
-)
-
-const (
-	fragTestSrc   = "100.10.0.1"
-	fragTestDst   = "100.10.0.100"
-	fragTestSrcV6 = "fd00::1"
-	fragTestDstV6 = "fd00::100"
-)
-
-func newFragmentTestManager(tb testing.TB) *Manager {
-	tb.Helper()
-
-	ifaceMock := &IFaceMock{
-		SetFilterFunc: func(device.PacketFilter) error { return nil },
-		AddressFunc: func() wgaddr.Address {
-			return wgaddr.Address{
-				IP:      netip.MustParseAddr(fragTestDst),
-				Network: netip.MustParsePrefix("100.10.0.0/16"),
-				IPv6:    netip.MustParseAddr(fragTestDstV6),
-				IPv6Net: netip.MustParsePrefix("fd00::/64"),
-			}
-		},
-	}
-
-	m, err := Create(ifaceMock, false, flowLogger, nbiface.DefaultMTU)
-	require.NoError(tb, err)
-	require.NoError(tb, m.UpdateLocalIPs())
-	tb.Cleanup(func() { require.NoError(tb, m.Close(nil)) })
-	return m
-}
-
-// firstFragmentUDPTo builds the first fragment of a fragmented UDP datagram to
-// the given destination: it carries the full UDP header plus payloadLen bytes
-// of data, with the More Fragments flag set and offset zero.
-func firstFragmentUDPTo(tb testing.TB, dst string, id uint16, dstPort uint16, payloadLen int) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv4{
-		Version:  4,
-		TTL:      64,
-		Id:       id,
-		Protocol: layers.IPProtocolUDP,
-		SrcIP:    net.ParseIP(fragTestSrc),
-		DstIP:    net.ParseIP(dst),
-		Flags:    layers.IPv4MoreFragments,
-	}
-	udp := &layers.UDP{SrcPort: 40000, DstPort: layers.UDPPort(dstPort)}
-	require.NoError(tb, udp.SetNetworkLayerForChecksum(ip))
-
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
-	require.NoError(tb, gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(make([]byte, payloadLen))))
-	return buf.Bytes()
-}
-
-func firstFragmentUDP(tb testing.TB, id uint16, dstPort uint16, payloadLen int) []byte {
-	tb.Helper()
-	return firstFragmentUDPTo(tb, fragTestDst, id, dstPort, payloadLen)
-}
-
-// firstFragmentTCP builds the first fragment of a fragmented TCP datagram: the
-// full 20-byte TCP header plus 12 bytes of data, with the More Fragments flag
-// set and offset zero.
-func firstFragmentTCP(tb testing.TB, id uint16, dstPort uint16) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv4{
-		Version:  4,
-		TTL:      64,
-		Id:       id,
-		Protocol: layers.IPProtocolTCP,
-		SrcIP:    net.ParseIP(fragTestSrc),
-		DstIP:    net.ParseIP(fragTestDst),
-		Flags:    layers.IPv4MoreFragments,
-	}
-	tcp := &layers.TCP{SrcPort: 40000, DstPort: layers.TCPPort(dstPort), SYN: true, Window: 64240}
-	require.NoError(tb, tcp.SetNetworkLayerForChecksum(ip))
-
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
-	require.NoError(tb, gopacket.SerializeLayers(buf, opts, ip, tcp, gopacket.Payload(make([]byte, 12))))
-	return buf.Bytes()
-}
-
-// trailingFragmentTo builds a non-first fragment to the given destination: an
-// IPv4 header at the given fragment offset (in 8-byte units) carrying raw
-// payload and no L4 header.
-func trailingFragmentTo(tb testing.TB, dst string, proto layers.IPProtocol, id uint16, fragOffsetOctets uint16, moreFragments bool, payloadLen int) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv4{
-		Version:    4,
-		TTL:        64,
-		Id:         id,
-		Protocol:   proto,
-		SrcIP:      net.ParseIP(fragTestSrc),
-		DstIP:      net.ParseIP(dst),
-		FragOffset: fragOffsetOctets,
-	}
-	if moreFragments {
-		ip.Flags = layers.IPv4MoreFragments
-	}
-
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{FixLengths: true}
-	require.NoError(tb, gopacket.SerializeLayers(buf, opts, ip, gopacket.Payload(make([]byte, payloadLen))))
-	return buf.Bytes()
-}
-
-func trailingFragment(tb testing.TB, id uint16, fragOffsetOctets uint16, moreFragments bool, payloadLen int) []byte {
-	tb.Helper()
-	return trailingFragmentTo(tb, fragTestDst, layers.IPProtocolUDP, id, fragOffsetOctets, moreFragments, payloadLen)
-}
-
-// outboundUDPPacket builds a complete outbound UDP packet from the local
-// address, used to establish conntrack state for reply-direction tests.
-func outboundUDPPacket(tb testing.TB, srcPort, dstPort uint16) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv4{
-		Version:  4,
-		TTL:      64,
-		Id:       1,
-		Protocol: layers.IPProtocolUDP,
-		SrcIP:    net.ParseIP(fragTestDst),
-		DstIP:    net.ParseIP(fragTestSrc),
-	}
-	udp := &layers.UDP{SrcPort: layers.UDPPort(srcPort), DstPort: layers.UDPPort(dstPort)}
-	require.NoError(tb, udp.SetNetworkLayerForChecksum(ip))
-
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
-	require.NoError(tb, gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(make([]byte, 16))))
-	return buf.Bytes()
-}
-
-// normalUDPPacket builds a complete, non-fragmented UDP packet for baseline
-// comparisons against the fragment paths.
-func normalUDPPacket(tb testing.TB, dstPort uint16, payloadLen int) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv4{
-		Version:  4,
-		TTL:      64,
-		Id:       1,
-		Protocol: layers.IPProtocolUDP,
-		SrcIP:    net.ParseIP(fragTestSrc),
-		DstIP:    net.ParseIP(fragTestDst),
-	}
-	udp := &layers.UDP{SrcPort: 40000, DstPort: layers.UDPPort(dstPort)}
-	require.NoError(tb, udp.SetNetworkLayerForChecksum(ip))
-
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
-	require.NoError(tb, gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(make([]byte, payloadLen))))
-	return buf.Bytes()
-}
-
-func allowUDP(tb testing.TB, m *Manager, dstPort uint16) {
-	tb.Helper()
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrc), fw.ProtocolUDP, nil,
-		&fw.Port{Values: []uint16{dstPort}}, fw.ActionAccept, "")
-	require.NoError(tb, err)
-}
-
-// TestFragment_TrailingWithoutFirstDropped is the core bypass repro: a trailing
-// fragment with no allowed first fragment on record must be dropped. Before the
-// fix, filterInbound returned false (allow) for any fragment.
-func TestFragment_TrailingWithoutFirstDropped(t *testing.T) {
-	m := newFragmentTestManager(t)
-
-	frag := trailingFragment(t, 0x1234, 185, false, 40)
-	require.True(t, m.filterInbound(frag, len(frag)),
-		"trailing fragment without an allowed first fragment must be dropped")
-}
-
-// TestFragment_AllowedFirstPassesTrailing verifies that once a first fragment
-// passes the ACL, its trailing fragments inherit the allow verdict.
-func TestFragment_AllowedFirstPassesTrailing(t *testing.T) {
-	m := newFragmentTestManager(t)
-	allowUDP(t, m, 8080)
-
-	// First fragment: UDP header (8) + 32 payload = 40 octets -> headerEnd = 5.
-	first := firstFragmentUDP(t, 0x2222, 8080, 32)
-	require.False(t, m.filterInbound(first, len(first)),
-		"allowed first fragment should pass and be recorded")
-
-	trailing := trailingFragment(t, 0x2222, 5, false, 24)
-	require.False(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of an allowed datagram should pass")
-}
-
-// TestFragment_DeniedFirstDropsTrailing verifies that a first fragment blocked
-// by the ACL leaves no verdict, so its trailing fragments are dropped.
-func TestFragment_DeniedFirstDropsTrailing(t *testing.T) {
-	m := newFragmentTestManager(t)
-	// No accept rule: local traffic defaults to deny.
-
-	first := firstFragmentUDP(t, 0x3333, 9999, 32)
-	require.True(t, m.filterInbound(first, len(first)),
-		"first fragment to a blocked port should be dropped by the ACL")
-
-	trailing := trailingFragment(t, 0x3333, 5, false, 24)
-	require.True(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of a denied datagram must be dropped")
-}
-
-// TestFragment_OverlappingHeaderDropped covers the RFC 1858 §4 / RFC 3128
-// overlapping-fragment rewrite: a trailing fragment starting inside the range
-// the ACL already inspected is dropped and poisons the datagram. TCP is used so
-// the overlap lands on real header bytes (the flags at byte 13).
-func TestFragment_OverlappingHeaderDropped(t *testing.T) {
-	m := newFragmentTestManager(t)
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrc), fw.ProtocolTCP, nil,
-		&fw.Port{Values: []uint16{8080}}, fw.ActionAccept, "")
-	require.NoError(t, err)
-
-	// First fragment: TCP header (20) + 12 data = 32 bytes -> headerEnd = 4 octets.
-	first := firstFragmentTCP(t, 0x4444, 8080)
-	require.False(t, m.filterInbound(first, len(first)))
-
-	// Overlapping fragment at offset 1 (byte 8) falls inside the inspected TCP
-	// header, so it could rewrite the flags or port on reassembly.
-	overlap := trailingFragmentTo(t, fragTestDst, layers.IPProtocolTCP, 0x4444, 1, true, 32)
-	require.True(t, m.filterInbound(overlap, len(overlap)),
-		"fragment overlapping the inspected header must be dropped")
-
-	// The datagram is now poisoned: a later, non-overlapping fragment is also
-	// dropped because the verdict was removed.
-	later := trailingFragmentTo(t, fragTestDst, layers.IPProtocolTCP, 0x4444, 4, false, 24)
-	require.True(t, m.filterInbound(later, len(later)),
-		"fragments after an overlap must be dropped (datagram poisoned)")
-}
-
-// TestFragment_OffsetZeroOverlapPoisons covers the RFC 3128 offset-zero rewrite:
-// an allowed first fragment followed by a denied offset-zero fragment for the
-// same datagram must not leave the earlier allow verdict in place.
-func TestFragment_OffsetZeroOverlapPoisons(t *testing.T) {
-	m := newFragmentTestManager(t)
-	allowUDP(t, m, 8080)
-
-	allowed := firstFragmentUDP(t, 0x5A5A, 8080, 32)
-	require.False(t, m.filterInbound(allowed, len(allowed)),
-		"allowed first fragment should pass and be recorded")
-
-	// A second offset-zero fragment to a denied port supersedes the datagram's
-	// verdict; it is dropped and must not leave the allow in place.
-	denied := firstFragmentUDP(t, 0x5A5A, 9999, 32)
-	require.True(t, m.filterInbound(denied, len(denied)),
-		"denied offset-zero fragment must be dropped")
-
-	trailing := trailingFragment(t, 0x5A5A, 5, false, 24)
-	require.True(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment must be denied after the datagram was poisoned")
-}
-
-// TestFragment_TinyFirstDropped covers the tiny-fragment attack: a first
-// fragment too small to contain the full transport header can't be
-// ACL-evaluated and must be dropped.
-func TestFragment_TinyFirstDropped(t *testing.T) {
-	m := newFragmentTestManager(t)
-	allowUDP(t, m, 8080)
-
-	// IPv4 header + 4 raw bytes, MF set, offset 0: too small for the 8-byte UDP
-	// header, so it decodes to L3 only.
-	tiny := trailingFragment(t, 0x5555, 0, true, 4)
-	require.True(t, m.filterInbound(tiny, len(tiny)),
-		"tiny first fragment without a full L4 header must be dropped")
-}
-
-// TestFragment_TCPFirstFragment verifies the TCP arm of the transport decode: a
-// first fragment carrying the full 20-byte TCP header is ACL-evaluated and its
-// trailing fragments inherit the verdict.
-func TestFragment_TCPFirstFragment(t *testing.T) {
-	m := newFragmentTestManager(t)
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrc), fw.ProtocolTCP, nil,
-		&fw.Port{Values: []uint16{8080}}, fw.ActionAccept, "")
-	require.NoError(t, err)
-
-	// TCP header (20) + 12 data = 32 bytes -> headerEnd = 4 octets.
-	first := firstFragmentTCP(t, 0x6666, 8080)
-	require.False(t, m.filterInbound(first, len(first)),
-		"allowed TCP first fragment should pass and be recorded")
-
-	trailing := trailingFragmentTo(t, fragTestDst, layers.IPProtocolTCP, 0x6666, 4, false, 24)
-	require.False(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of an allowed TCP datagram should pass")
-}
-
-// TestFragment_TCPTinyFirstDropped verifies the TCP minimum header length: 12
-// bytes would satisfy a UDP header but falls short of the 20-byte TCP header.
-func TestFragment_TCPTinyFirstDropped(t *testing.T) {
-	m := newFragmentTestManager(t)
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrc), fw.ProtocolTCP, nil,
-		&fw.Port{Values: []uint16{8080}}, fw.ActionAccept, "")
-	require.NoError(t, err)
-
-	tiny := trailingFragmentTo(t, fragTestDst, layers.IPProtocolTCP, 0x7777, 0, true, 12)
-	require.True(t, m.filterInbound(tiny, len(tiny)),
-		"first fragment shorter than the TCP header must be dropped")
-}
-
-// TestFragment_ConntrackAllowsFirstFragment verifies the conntrack branch: reply
-// fragments of an outbound-established UDP flow pass without any inbound rule.
-func TestFragment_ConntrackAllowsFirstFragment(t *testing.T) {
-	m := newFragmentTestManager(t)
-
-	out := outboundUDPPacket(t, 12345, 40000)
-	require.False(t, m.filterOutbound(out, len(out)))
-
-	first := firstFragmentUDP(t, 0x8888, 12345, 32)
-	require.False(t, m.filterInbound(first, len(first)),
-		"reply first fragment should pass via conntrack")
-
-	trailing := trailingFragment(t, 0x8888, 5, false, 24)
-	require.False(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of a tracked flow should pass")
-}
-
-// TestFragment_RoutingDisabledDropsFragment verifies routed first fragments are
-// dropped when routing is disabled.
-func TestFragment_RoutingDisabledDropsFragment(t *testing.T) {
-	m := newFragmentTestManager(t)
-	m.routingEnabled.Store(false)
-
-	first := firstFragmentUDPTo(t, "198.51.100.10", 0x9999, 8080, 32)
-	require.True(t, m.filterInbound(first, len(first)),
-		"routed first fragment must be dropped when routing is disabled")
-}
-
-// TestFragment_RouteACL verifies the route-ACL branch: fragments to a non-local
-// destination follow the route rules, allowed datagrams pass their trailing
-// fragments and denied ones don't.
-func TestFragment_RouteACL(t *testing.T) {
-	m := newFragmentTestManager(t)
-	m.routingEnabled.Store(true)
-	m.nativeRouter.Store(false)
-
-	_, err := m.AddRouteFiltering(
-		[]byte("rt-1"),
-		[]netip.Prefix{netip.MustParsePrefix("100.10.0.0/16")},
-		fw.Network{Prefix: netip.MustParsePrefix("198.51.100.0/24")},
-		fw.ProtocolUDP,
-		nil,
-		&fw.Port{Values: []uint16{8080}},
-		fw.ActionAccept,
-	)
-	require.NoError(t, err)
-
-	first := firstFragmentUDPTo(t, "198.51.100.10", 0xAAAA, 8080, 32)
-	require.False(t, m.filterInbound(first, len(first)),
-		"route-ACL-allowed first fragment should pass")
-	trailing := trailingFragmentTo(t, "198.51.100.10", layers.IPProtocolUDP, 0xAAAA, 5, false, 24)
-	require.False(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of an allowed routed datagram should pass")
-
-	denied := firstFragmentUDPTo(t, "198.51.100.10", 0xBBBB, 9999, 32)
-	require.True(t, m.filterInbound(denied, len(denied)),
-		"route-ACL-denied first fragment must be dropped")
-	deniedTrailing := trailingFragmentTo(t, "198.51.100.10", layers.IPProtocolUDP, 0xBBBB, 5, false, 24)
-	require.True(t, m.filterInbound(deniedTrailing, len(deniedTrailing)),
-		"trailing fragment of a denied routed datagram must be dropped")
-}
-
-// TestFragment_ExpiredVerdictDropsTrailing verifies a verdict older than the
-// tracker timeout no longer admits trailing fragments.
-func TestFragment_ExpiredVerdictDropsTrailing(t *testing.T) {
-	m := newFragmentTestManager(t)
-	allowUDP(t, m, 8080)
-
-	first := firstFragmentUDP(t, 0xCCCC, 8080, 32)
-	require.False(t, m.filterInbound(first, len(first)))
-
-	m.fragments.mutex.Lock()
-	for key, entry := range m.fragments.entries {
-		entry.recordedAt = time.Now().Add(-defaultFragmentTimeout - time.Second)
-		m.fragments.entries[key] = entry
-	}
-	m.fragments.mutex.Unlock()
-
-	trailing := trailingFragment(t, 0xCCCC, 5, false, 24)
-	require.True(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment after verdict expiry must be dropped")
-}
-
-// TestFragment_CapacityFailsClosed verifies the table cap: at capacity, new
-// datagram verdicts are not recorded (their trailing fragments are dropped)
-// while already-recorded datagrams keep working.
-func TestFragment_CapacityFailsClosed(t *testing.T) {
-	m := newFragmentTestManager(t)
-	allowUDP(t, m, 8080)
-
-	m.fragments.mutex.Lock()
-	m.fragments.maxEntries = 1
-	m.fragments.mutex.Unlock()
-
-	first1 := firstFragmentUDP(t, 0x0101, 8080, 32)
-	require.False(t, m.filterInbound(first1, len(first1)))
-
-	first2 := firstFragmentUDP(t, 0x0202, 8080, 32)
-	require.False(t, m.filterInbound(first2, len(first2)),
-		"first fragment itself still passes at capacity")
-
-	trailing2 := trailingFragment(t, 0x0202, 5, false, 24)
-	require.True(t, m.filterInbound(trailing2, len(trailing2)),
-		"trailing fragment of an unrecorded datagram must be dropped at capacity")
-
-	trailing1 := trailingFragment(t, 0x0101, 5, false, 24)
-	require.False(t, m.filterInbound(trailing1, len(trailing1)),
-		"already-recorded datagram should keep passing at capacity")
-}
-
-// v6FragmentHeader builds the 8-byte IPv6 fragment extension header for the
-// given inner protocol, offset (8-byte units), More Fragments bit and id.
-func v6FragmentHeader(proto layers.IPProtocol, offsetOctets uint16, moreFragments bool, id uint32) []byte {
-	offsetFlags := offsetOctets << 3
-	if moreFragments {
-		offsetFlags |= 1
-	}
-	hdr := make([]byte, 8)
-	hdr[0] = uint8(proto)
-	binary.BigEndian.PutUint16(hdr[2:4], offsetFlags)
-	binary.BigEndian.PutUint32(hdr[4:8], id)
-	return hdr
-}
-
-func v6UDPHeader(dstPort uint16, dataLen int) []byte {
-	hdr := make([]byte, 8)
-	binary.BigEndian.PutUint16(hdr[0:2], 40000)
-	binary.BigEndian.PutUint16(hdr[2:4], dstPort)
-	binary.BigEndian.PutUint16(hdr[4:6], uint16(8+dataLen))
-	return hdr
-}
-
-// firstFragmentUDPv6 builds the first fragment of a fragmented IPv6 UDP
-// datagram: fragment header (offset 0, More Fragments set) + full UDP header +
-// data.
-func firstFragmentUDPv6(tb testing.TB, id uint32, dstPort uint16, dataLen int) []byte {
-	tb.Helper()
-	return fragmentUDPv6(tb, id, dstPort, dataLen, true)
-}
-
-// fragmentUDPv6 builds an offset-zero IPv6 UDP fragment. With moreFragments
-// false it is an atomic fragment (a complete datagram, RFC 6946).
-func fragmentUDPv6(tb testing.TB, id uint32, dstPort uint16, dataLen int, moreFragments bool) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv6{
-		Version:    6,
-		NextHeader: layers.IPProtocolIPv6Fragment,
-		HopLimit:   64,
-		SrcIP:      net.ParseIP(fragTestSrcV6),
-		DstIP:      net.ParseIP(fragTestDstV6),
-	}
-	payload := append(v6FragmentHeader(layers.IPProtocolUDP, 0, moreFragments, id), v6UDPHeader(dstPort, dataLen)...)
-	payload = append(payload, make([]byte, dataLen)...)
-
-	buf := gopacket.NewSerializeBuffer()
-	require.NoError(tb, gopacket.SerializeLayers(buf, gopacket.SerializeOptions{FixLengths: true}, ip, gopacket.Payload(payload)))
-	return buf.Bytes()
-}
-
-// trailingFragmentV6 builds a non-first IPv6 fragment: fragment header at the
-// given offset carrying raw data and no transport header.
-func trailingFragmentV6(tb testing.TB, id uint32, offsetOctets uint16, moreFragments bool, dataLen int) []byte {
-	tb.Helper()
-
-	ip := &layers.IPv6{
-		Version:    6,
-		NextHeader: layers.IPProtocolIPv6Fragment,
-		HopLimit:   64,
-		SrcIP:      net.ParseIP(fragTestSrcV6),
-		DstIP:      net.ParseIP(fragTestDstV6),
-	}
-	payload := append(v6FragmentHeader(layers.IPProtocolUDP, offsetOctets, moreFragments, id), make([]byte, dataLen)...)
-
-	buf := gopacket.NewSerializeBuffer()
-	require.NoError(tb, gopacket.SerializeLayers(buf, gopacket.SerializeOptions{FixLengths: true}, ip, gopacket.Payload(payload)))
-	return buf.Bytes()
-}
-
-// TestFragmentV6_TrailingWithoutFirstDropped verifies the IPv6 bypass is closed:
-// a trailing fragment with no allowed first fragment is dropped.
-func TestFragmentV6_TrailingWithoutFirstDropped(t *testing.T) {
-	m := newFragmentTestManager(t)
-
-	frag := trailingFragmentV6(t, 0xAABBCCDD, 100, false, 40)
-	require.True(t, m.filterInbound(frag, len(frag)),
-		"IPv6 trailing fragment without an allowed first fragment must be dropped")
-}
-
-// TestFragmentV6_AllowedFirstPassesTrailing verifies IPv6 fragments are
-// evaluated like IPv4: an allowed first fragment lets its trailing fragments
-// through.
-func TestFragmentV6_AllowedFirstPassesTrailing(t *testing.T) {
-	m := newFragmentTestManager(t)
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrcV6), fw.ProtocolUDP, nil,
-		&fw.Port{Values: []uint16{8080}}, fw.ActionAccept, "")
-	require.NoError(t, err)
-
-	// First fragment: UDP header (8) + 32 data = 40 octets -> headerEnd = 5.
-	first := firstFragmentUDPv6(t, 0xAABBCCDD, 8080, 32)
-	require.False(t, m.filterInbound(first, len(first)),
-		"allowed IPv6 first fragment should pass and be recorded")
-
-	trailing := trailingFragmentV6(t, 0xAABBCCDD, 5, false, 24)
-	require.False(t, m.filterInbound(trailing, len(trailing)),
-		"trailing fragment of an allowed IPv6 datagram should pass")
-}
-
-// TestFragmentV6_AtomicNotCached verifies an IPv6 atomic fragment (fragment
-// header with offset 0 and no More Fragments, a complete datagram per RFC 6946)
-// is evaluated but not recorded, so a flood of allowed atomic fragments can't
-// exhaust the verdict table.
-func TestFragmentV6_AtomicNotCached(t *testing.T) {
-	m := newFragmentTestManager(t)
-	_, err := m.AddPeerFiltering(nil, net.ParseIP(fragTestSrcV6), fw.ProtocolUDP, nil,
-		&fw.Port{Values: []uint16{8080}}, fw.ActionAccept, "")
-	require.NoError(t, err)
-
-	atomic := fragmentUDPv6(t, 0xA70301C, 8080, 16, false)
-	require.False(t, m.filterInbound(atomic, len(atomic)),
-		"allowed IPv6 atomic fragment should pass")
-
-	m.fragments.mutex.Lock()
-	n := len(m.fragments.entries)
-	m.fragments.mutex.Unlock()
-	require.Zero(t, n, "atomic fragment must not create a verdict entry")
-
-	// A genuine fragmented datagram (More Fragments set) is still recorded.
-	first := fragmentUDPv6(t, 0xBEEF, 8080, 32, true)
-	require.False(t, m.filterInbound(first, len(first)))
-	m.fragments.mutex.Lock()
-	n = len(m.fragments.entries)
-	m.fragments.mutex.Unlock()
-	require.Equal(t, 1, n, "genuine first fragment must record a verdict")
-}

client/iface/netstack/env.go

@@ -3,31 +3,14 @@
 package netstack
 
 import (
-	"net"
+	"fmt"
 	"os"
 	"strconv"
 
 	log "github.com/sirupsen/logrus"
 )
 
-const (
-	EnvUseNetstackMode = "NB_USE_NETSTACK_MODE"
-
-	// EnvSocks5ListenerPort overrides the port the SOCKS5 proxy listens on.
-	EnvSocks5ListenerPort = "NB_SOCKS5_LISTENER_PORT"
-
-	// EnvSocks5ListenerAddress overrides the host/IP the SOCKS5 proxy binds to.
-	// The proxy is a bridge for local host applications into the userspace
-	// WireGuard netstack, so it binds to loopback by default. Override this only
-	// when the proxy must be reachable from other hosts (e.g. a container
-	// gateway); doing so exposes an unauthenticated SOCKS5 proxy on that
-	// address.
-	EnvSocks5ListenerAddress = "NB_SOCKS5_LISTENER_ADDRESS"
-
-	// defaultSocks5Host is the loopback address the SOCKS5 proxy binds to unless
-	// overridden via EnvSocks5ListenerAddress.
-	defaultSocks5Host = "127.0.0.1"
-)
+const EnvUseNetstackMode = "NB_USE_NETSTACK_MODE"
 
 // IsEnabled todo: move these function to cmd layer
 func IsEnabled() bool {
@@ -35,40 +18,24 @@ func IsEnabled() bool {
 }
 
 func ListenAddr() string {
-	return net.JoinHostPort(listenHost(), strconv.Itoa(listenPort()))
-}
-
-// listenHost returns the host/IP the SOCKS5 proxy binds to. It defaults to
-// loopback and only honors EnvSocks5ListenerAddress when it holds a valid IP.
-func listenHost() string {
-	addr := os.Getenv(EnvSocks5ListenerAddress)
-	if addr == "" {
-		return defaultSocks5Host
-	}
-	if net.ParseIP(addr) == nil {
-		log.Warnf("invalid socks5 listener address %q, falling back to default: %s", addr, defaultSocks5Host)
-		return defaultSocks5Host
-	}
-	return addr
-}
-
-// listenPort returns the port the SOCKS5 proxy binds to, defaulting to
-// DefaultSocks5Port when EnvSocks5ListenerPort is unset or invalid.
-func listenPort() int {
-	sPort := os.Getenv(EnvSocks5ListenerPort)
+	sPort := os.Getenv("NB_SOCKS5_LISTENER_PORT")
 	if sPort == "" {
-		return DefaultSocks5Port
+		return listenAddr(DefaultSocks5Port)
 	}
 
 	port, err := strconv.Atoi(sPort)
 	if err != nil {
 		log.Warnf("invalid socks5 listener port, unable to convert it to int, falling back to default: %d", DefaultSocks5Port)
-		return DefaultSocks5Port
+		return listenAddr(DefaultSocks5Port)
 	}
 	if port < 1 || port > 65535 {
 		log.Warnf("invalid socks5 listener port, it should be in the range 1-65535, falling back to default: %d", DefaultSocks5Port)
-		return DefaultSocks5Port
+		return listenAddr(DefaultSocks5Port)
 	}
 
-	return port
+	return listenAddr(port)
+}
+
+func listenAddr(port int) string {
+	return fmt.Sprintf("0.0.0.0:%d", port)
 }

client/iface/netstack/env_test.go

@@ -1,63 +0,0 @@
-//go:build !js
-
-package netstack
-
-import (
-	"net"
-	"strconv"
-	"testing"
-)
-
-func TestListenAddr_DefaultsToLoopback(t *testing.T) {
-	// No env overrides: must bind loopback, never all interfaces.
-	got := ListenAddr()
-	want := net.JoinHostPort("127.0.0.1", strconv.Itoa(DefaultSocks5Port))
-	if got != want {
-		t.Fatalf("ListenAddr() = %q, want %q", got, want)
-	}
-}
-
-func TestListenAddr_AddressOverride(t *testing.T) {
-	tests := []struct {
-		name string
-		env  string
-		want string
-	}{
-		{name: "valid override honored", env: "0.0.0.0", want: "0.0.0.0"},
-		{name: "valid specific ip honored", env: "10.0.0.5", want: "10.0.0.5"},
-		{name: "ipv6 loopback bracketed", env: "::1", want: "::1"},
-		{name: "invalid falls back to loopback", env: "not-an-ip", want: "127.0.0.1"},
-		{name: "empty falls back to loopback", env: "", want: "127.0.0.1"},
-	}
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			t.Setenv(EnvSocks5ListenerAddress, tc.env)
-			want := net.JoinHostPort(tc.want, strconv.Itoa(DefaultSocks5Port))
-			if got := ListenAddr(); got != want {
-				t.Fatalf("ListenAddr() = %q, want %q", got, want)
-			}
-		})
-	}
-}
-
-func TestListenAddr_PortOverride(t *testing.T) {
-	tests := []struct {
-		name string
-		env  string
-		want int
-	}{
-		{name: "valid port honored", env: "1081", want: 1081},
-		{name: "non-numeric falls back", env: "abc", want: DefaultSocks5Port},
-		{name: "out of range falls back", env: "70000", want: DefaultSocks5Port},
-		{name: "zero falls back", env: "0", want: DefaultSocks5Port},
-	}
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			t.Setenv(EnvSocks5ListenerPort, tc.env)
-			want := net.JoinHostPort("127.0.0.1", strconv.Itoa(tc.want))
-			if got := ListenAddr(); got != want {
-				t.Fatalf("ListenAddr() = %q, want %q", got, want)
-			}
-		})
-	}
-}

client/internal/auth/device_flow.go

@@ -299,7 +299,7 @@ func (d *DeviceAuthorizationFlow) WaitToken(ctx context.Context, info AuthFlowIn
 				UseIDToken:   d.providerConfig.UseIDToken,
 			}
 
-			err = validateTokenAudience(tokenInfo.GetTokenToUse(), d.providerConfig.Audience)
+			err = isValidAccessToken(tokenInfo.GetTokenToUse(), d.providerConfig.Audience)
 			if err != nil {
 				return TokenInfo{}, fmt.Errorf("validate access token failed with error: %v", err)
 			}

client/internal/auth/pkce_flow.go

@@ -306,7 +306,7 @@ func (p *PKCEAuthorizationFlow) parseOAuthToken(token *oauth2.Token) (TokenInfo,
 		audience = p.providerConfig.ClientID
 	}
 
-	if err := validateTokenAudience(tokenInfo.GetTokenToUse(), audience); err != nil {
+	if err := isValidAccessToken(tokenInfo.GetTokenToUse(), audience); err != nil {
 		return TokenInfo{}, fmt.Errorf("authentication failed: invalid access token - %w", err)
 	}
 
@@ -320,11 +320,6 @@ func (p *PKCEAuthorizationFlow) parseOAuthToken(token *oauth2.Token) (TokenInfo,
 	return tokenInfo, nil
 }
 
-// parseEmailFromIDToken extracts the email (or name) claim from an ID token
-// without verifying its signature. The value is best-effort and used only as a
-// UX convenience (login hint prefill and display); it never drives an
-// authorization decision. The authoritative identity is established server-side
-// from the signature-verified token.
 func parseEmailFromIDToken(token string) (string, error) {
 	parts := strings.Split(token, ".")
 	if len(parts) < 2 {

client/internal/auth/util.go

@@ -20,26 +20,14 @@ func randomBytesInHex(count int) (string, error) {
 	return hex.EncodeToString(buf), nil
 }
 
-// validateTokenAudience checks that the token is a well-formed JWT whose
-// audience claim matches the expected audience.
-//
-// It does NOT verify the token's cryptographic signature and therefore must not
-// be treated as an authenticity check. The token is obtained by the client
-// directly from the IdP token endpoint over TLS, and its signature is verified
-// server-side by the management server against the IdP's JWKS
-// (see shared/auth/jwt/validator.go). This function is only a client-side
-// sanity check that the returned token targets the expected audience.
-func validateTokenAudience(token string, audience string) error {
+// isValidAccessToken is a simple validation of the access token
+func isValidAccessToken(token string, audience string) error {
 	if token == "" {
 		return fmt.Errorf("token received is empty")
 	}
 
-	parts := strings.Split(token, ".")
-	if len(parts) != 3 {
-		return fmt.Errorf("token is not a well-formed JWT")
-	}
-
-	claimsString, err := base64.RawURLEncoding.DecodeString(parts[1])
+	encodedClaims := strings.Split(token, ".")[1]
+	claimsString, err := base64.RawURLEncoding.DecodeString(encodedClaims)
 	if err != nil {
 		return err
 	}

client/internal/auth/util_test.go

@@ -1,108 +0,0 @@
-package auth
-
-import (
-	"encoding/base64"
-	"encoding/json"
-	"testing"
-)
-
-// makeJWT builds an unsigned JWT-shaped string (header.payload.signature) with
-// the given claims payload. The signature part is arbitrary because
-// validateTokenAudience intentionally does not verify it.
-func makeJWT(t *testing.T, claims map[string]interface{}) string {
-	t.Helper()
-	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
-	payloadBytes, err := json.Marshal(claims)
-	if err != nil {
-		t.Fatalf("marshal claims: %v", err)
-	}
-	payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
-	return header + "." + payload + ".unverified-signature"
-}
-
-func TestValidateTokenAudience(t *testing.T) {
-	tests := []struct {
-		name     string
-		token    string
-		audience string
-		wantErr  bool
-	}{
-		{
-			name:     "empty token",
-			token:    "",
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "not a JWT - no dots",
-			token:    "notajwt",
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "not a JWT - two parts only",
-			token:    "header.payload",
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "matching string audience",
-			token:    makeJWT(t, map[string]interface{}{"aud": "netbird"}),
-			audience: "netbird",
-			wantErr:  false,
-		},
-		{
-			name:     "mismatching string audience",
-			token:    makeJWT(t, map[string]interface{}{"aud": "other"}),
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "matching audience in array",
-			token:    makeJWT(t, map[string]interface{}{"aud": []interface{}{"other", "netbird"}}),
-			audience: "netbird",
-			wantErr:  false,
-		},
-		{
-			name:     "mismatching audience array",
-			token:    makeJWT(t, map[string]interface{}{"aud": []interface{}{"a", "b"}}),
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "missing audience claim",
-			token:    makeJWT(t, map[string]interface{}{"sub": "user"}),
-			audience: "netbird",
-			wantErr:  true,
-		},
-		{
-			name:     "invalid base64 payload",
-			token:    "header.!!!not-base64!!!.sig",
-			audience: "netbird",
-			wantErr:  true,
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			err := validateTokenAudience(tc.token, tc.audience)
-			if tc.wantErr && err == nil {
-				t.Fatalf("expected error, got nil")
-			}
-			if !tc.wantErr && err != nil {
-				t.Fatalf("expected no error, got %v", err)
-			}
-		})
-	}
-}
-
-// TestValidateTokenAudienceNoPanic guards the regression where a non-empty
-// token without the JWT dot structure caused an index-out-of-range panic.
-func TestValidateTokenAudienceNoPanic(t *testing.T) {
-	inputs := []string{"a", ".", "a.", "aaaa", "no-dots-here"}
-	for _, in := range inputs {
-		if err := validateTokenAudience(in, "netbird"); err == nil {
-			t.Fatalf("expected error for malformed token %q, got nil", in)
-		}
-	}
-}

client/internal/dns/service_listener.go

@@ -292,16 +292,18 @@ func (s *serviceViaListener) generateFreePort() (uint16, error) {
 		return customPort, nil
 	}
 
-	probeListener, err := net.ListenUDP("udp4", &net.UDPAddr{})
+	udpAddr := net.UDPAddrFromAddrPort(netip.MustParseAddrPort("0.0.0.0:0"))
+	probeListener, err := net.ListenUDP("udp", udpAddr)
 	if err != nil {
 		log.Debugf("failed to bind random port for DNS: %s", err)
 		return 0, err
 	}
 
-	port := uint16(probeListener.LocalAddr().(*net.UDPAddr).Port)
-	if err = probeListener.Close(); err != nil {
+	addrPort := netip.MustParseAddrPort(probeListener.LocalAddr().String()) // might panic if address is incorrect
+	err = probeListener.Close()
+	if err != nil {
 		log.Debugf("failed to free up DNS port: %s", err)
 		return 0, err
 	}
-	return port, nil
+	return addrPort.Port(), nil
 }

client/internal/peer/worker_relay.go

@@ -54,15 +54,19 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relaySupportedOnRemotePeer.Store(true)
 
 	// the relayManager will return with error in case if the connection has lost with relay server
-	_, _, err := w.relayManager.RelayInstanceAddress()
+	currentRelayAddress, _, err := w.relayManager.RelayInstanceAddress()
 	if err != nil {
 		w.log.Errorf("failed to handle new offer: %s", err)
 		return
 	}
 
-	preferForeign := !w.isController
-	remoteRelayServer := relayClient.RelayServer{Addr: remoteOfferAnswer.RelaySrvAddress, IP: remoteOfferAnswer.RelaySrvIP}
-	relayedConn, err := w.relayManager.OpenConn(w.peerCtx, remoteRelayServer, w.config.Key, preferForeign)
+	srv := w.preferredRelayServer(currentRelayAddress, remoteOfferAnswer.RelaySrvAddress)
+	var serverIP netip.Addr
+	if srv == remoteOfferAnswer.RelaySrvAddress {
+		serverIP = remoteOfferAnswer.RelaySrvIP
+	}
+
+	relayedConn, err := w.relayManager.OpenConn(w.peerCtx, srv, w.config.Key, serverIP)
 	if err != nil {
 		if errors.Is(err, relayClient.ErrConnAlreadyExists) {
 			w.log.Debugf("handled offer by reusing existing relay connection")
@@ -76,13 +80,14 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relayedConn = relayedConn
 	w.relayLock.Unlock()
 
-	if err := w.relayManager.AddCloseListener(relayedConn.RemoteAddr().String(), w.onRelayClientDisconnected); err != nil {
-		w.log.Errorf("failed to add close listener: %s", err)
+	err = w.relayManager.AddCloseListener(srv, w.onRelayClientDisconnected)
+	if err != nil {
+		log.Errorf("failed to add close listener: %s", err)
 		_ = relayedConn.Close()
 		return
 	}
 
-	w.log.Debugf("peer conn opened via Relay: %s", relayedConn.RemoteAddr())
+	w.log.Debugf("peer conn opened via Relay: %s", srv)
 	go w.conn.onRelayConnectionIsReady(RelayConnInfo{
 		relayedConn:     relayedConn,
 		rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
@@ -121,6 +126,13 @@ func (w *WorkerRelay) isRelaySupported(answer *OfferAnswer) bool {
 	return answer.RelaySrvAddress != ""
 }
 
+func (w *WorkerRelay) preferredRelayServer(myRelayAddress, remoteRelayAddress string) string {
+	if w.isController {
+		return myRelayAddress
+	}
+	return remoteRelayAddress
+}
+
 func (w *WorkerRelay) onRelayClientDisconnected() {
 	go w.conn.onRelayDisconnected()
 }

client/internal/routemanager/dnsinterceptor/handler.go

@@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
 
 			// AllowedIPs should use real IPs
 			if d.currentPeerKey != "" {
-				if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+				if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
 					merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
 				}
 			}
@@ -172,7 +172,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error {
 	}
 
 	// AllowedIPs use real IPs
-	if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil {
+	if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil {
 		return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err)
 	}
 
@@ -205,7 +205,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error {
 	for _, prefixes := range d.interceptedDomains {
 		for _, prefix := range prefixes {
 			// AllowedIPs use real IPs
-			if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
 			}
 		}

client/internal/routemanager/dynamic/route.go

@@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error {
 	var merr *multierror.Error
 	for _, domainPrefixes := range r.dynamicDomains {
 		for _, prefix := range domainPrefixes {
-			if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
 			}
 		}
@@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) {
 			merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err))
 		}
 		if r.currentPeerKey != "" {
-			if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
+			if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
 				merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
 			}
 		}

client/internal/routemanager/manager.go

@@ -215,7 +215,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
 		)
 	}
 
-	m.allowedIPsRefCounter = refcounter.New(
+	m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
 		func(prefix netip.Prefix, peerKey string) (string, error) {
 			// save peerKey to use it in the remove function
 			return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix)

client/internal/routemanager/refcounter/allowedips.go

@@ -0,0 +1,185 @@
+package refcounter
+
+import (
+	"errors"
+	"fmt"
+	"net/netip"
+	"sort"
+	"sync"
+
+	"github.com/hashicorp/go-multierror"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+)
+
+// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is
+// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most
+// one peer is active at a time even when several peers reference the prefix.
+type allowedIPsEntry struct {
+	// peers maps a peerKey to the number of references holding the prefix for that peer.
+	peers map[string]int
+	// active is the peerKey currently installed in WireGuard for this prefix ("" if none).
+	active string
+	// total is the sum of all per-peer reference counts (kept in sync with peers).
+	total int
+}
+
+// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs.
+//
+// The generic Counter keys only by prefix and remembers a single Out value set by the first
+// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or
+// multiple resolved domains) can reference the same prefix through different peers, and when the
+// peer currently installed in WireGuard releases its last reference the prefix must be handed over
+// to a surviving peer instead of being left pointing at the released one.
+//
+// It calls add/remove (which program WireGuard) only on the transitions that matter:
+//   - add on the first reference for a prefix, or when swapping the active peer;
+//   - remove on the last reference for a prefix, or on the old peer during a swap.
+type AllowedIPsRefCounter struct {
+	mu      sync.Mutex
+	entries map[netip.Prefix]*allowedIPsEntry
+	add     AddFunc[netip.Prefix, string, string]
+	remove  RemoveFunc[netip.Prefix, string]
+}
+
+// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter.
+// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer.
+// remove unprograms the prefix from the given peer.
+func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter {
+	return &AllowedIPsRefCounter{
+		entries: map[netip.Prefix]*allowedIPsEntry{},
+		add:     add,
+		remove:  remove,
+	}
+}
+
+// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first
+// reference to a prefix; while a different peer is already installed the prefix is left with it
+// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept.
+func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	e, ok := rm.entries[prefix]
+	if !ok {
+		e = &allowedIPsEntry{peers: map[string]int{}}
+		rm.entries[prefix] = e
+	}
+
+	logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
+		prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active)
+
+	// Program WireGuard only when nothing is installed yet for this prefix.
+	if e.active == "" {
+		out, err := rm.add(prefix, peerKey)
+		if errors.Is(err, ErrIgnore) {
+			if e.total == 0 {
+				delete(rm.entries, prefix)
+			}
+			return Ref[string]{Count: e.total, Out: e.active}, nil
+		}
+		if err != nil {
+			if e.total == 0 {
+				delete(rm.entries, prefix)
+			}
+			return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err)
+		}
+		e.active = out
+	}
+
+	e.peers[peerKey]++
+	e.total++
+
+	return Ref[string]{Count: e.total, Out: e.active}, nil
+}
+
+// Decrement removes a reference to prefix for peerKey. When the peer currently installed in
+// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists,
+// otherwise it is removed from WireGuard.
+func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	e, ok := rm.entries[prefix]
+	if !ok {
+		logCallerF("No allowed IP reference found for prefix %v", prefix)
+		return Ref[string]{}, nil
+	}
+
+	if e.peers[peerKey] > 0 {
+		logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
+			prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active)
+		e.peers[peerKey]--
+		e.total--
+		if e.peers[peerKey] == 0 {
+			delete(e.peers, peerKey)
+		}
+	} else {
+		logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey)
+	}
+
+	// If the peer currently installed in WireGuard still holds references, nothing to reprogram.
+	// Keying the check on the active peer (not the one just released) makes this self-healing:
+	// a prior swap whose remove/add failed leaves e.active pointing at a peer with no references,
+	// and this retries the hand-off on the next Decrement instead of getting stuck.
+	if e.active != "" && e.peers[e.active] > 0 {
+		return Ref[string]{Count: e.total, Out: e.active}, nil
+	}
+
+	// Detach the stale/gone active peer from WireGuard before reprogramming.
+	if e.active != "" {
+		if err := rm.remove(prefix, e.active); err != nil {
+			return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)
+		}
+		e.active = ""
+	}
+
+	// Hand the prefix over to a surviving peer, or drop the entry when none remain.
+	if survivor, ok := pickSurvivor(e.peers); ok {
+		out, err := rm.add(prefix, survivor)
+		if err != nil {
+			return Ref[string]{Count: e.total, Out: ""}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err)
+		}
+		e.active = out
+		return Ref[string]{Count: e.total, Out: e.active}, nil
+	}
+
+	delete(rm.entries, prefix)
+	return Ref[string]{Count: 0, Out: ""}, nil
+}
+
+// Flush removes all prefixes from WireGuard and clears the counter.
+func (rm *AllowedIPsRefCounter) Flush() error {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	var merr *multierror.Error
+	for prefix, e := range rm.entries {
+		if e.active == "" {
+			continue
+		}
+		logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active)
+		if err := rm.remove(prefix, e.active); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err))
+		}
+	}
+
+	clear(rm.entries)
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do
+// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable
+// (lowest peerKey) for predictable behavior and testability.
+func pickSurvivor(peers map[string]int) (string, bool) {
+	if len(peers) == 0 {
+		return "", false
+	}
+	keys := make([]string, 0, len(peers))
+	for k := range peers {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	return keys[0], true
+}

client/internal/routemanager/refcounter/allowedips_test.go

@@ -0,0 +1,241 @@
+package refcounter
+
+import (
+	"errors"
+	"net/netip"
+	"testing"
+)
+
+// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer.
+// failAdd/failRemove make the next add/remove fail once, to exercise the self-healing error paths.
+type fakeWG struct {
+	installed  map[netip.Prefix]string
+	adds       int
+	removes    int
+	failAdd    bool
+	failRemove bool
+}
+
+func newFakeWG() *fakeWG {
+	return &fakeWG{installed: map[netip.Prefix]string{}}
+}
+
+func (f *fakeWG) counter() *AllowedIPsRefCounter {
+	return NewAllowedIPs(
+		func(prefix netip.Prefix, peerKey string) (string, error) {
+			if f.failAdd {
+				f.failAdd = false
+				return "", errors.New("add failed")
+			}
+			f.adds++
+			f.installed[prefix] = peerKey
+			return peerKey, nil
+		},
+		func(prefix netip.Prefix, peerKey string) error {
+			if f.failRemove {
+				f.failRemove = false
+				return errors.New("remove failed")
+			}
+			f.removes++
+			// only clear if this peer is the one installed, mirroring wg semantics
+			if f.installed[prefix] == peerKey {
+				delete(f.installed, prefix)
+			}
+			return nil
+		},
+	)
+}
+
+func mustPrefix(t *testing.T, s string) netip.Prefix {
+	t.Helper()
+	p, err := netip.ParsePrefix(s)
+	if err != nil {
+		t.Fatalf("parse prefix %q: %v", s, err)
+	}
+	return p
+}
+
+func mustIncrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
+	t.Helper()
+	ref, err := c.Increment(p, peer)
+	if err != nil {
+		t.Fatalf("Increment(%v, %s): %v", p, peer, err)
+	}
+	return ref
+}
+
+func mustDecrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
+	t.Helper()
+	ref, err := c.Decrement(p, peer)
+	if err != nil {
+		t.Fatalf("Decrement(%v, %s): %v", p, peer, err)
+	}
+	return ref
+}
+
+// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same
+// prefix routed by different peers. Removing the network whose peer is installed must hand the
+// prefix over to the surviving peer instead of leaving it on the removed one.
+func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	// First peer wins while both are present.
+	if got := f.installed[p]; got != "peerA" {
+		t.Fatalf("expected peerA installed, got %q", got)
+	}
+
+	// Remove the active peer's network -> must swap to peerB.
+	mustDecrement(t, c, p, "peerA")
+	if got := f.installed[p]; got != "peerB" {
+		t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got)
+	}
+
+	// Remove the last one -> prefix gone.
+	mustDecrement(t, c, p, "peerB")
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("expected prefix removed, still installed on %q", f.installed[p])
+	}
+}
+
+// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard.
+func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	removesBefore := f.removes
+
+	mustDecrement(t, c, p, "peerB")
+	if f.installed[p] != "peerA" {
+		t.Fatalf("active peer must stay peerA, got %q", f.installed[p])
+	}
+	if f.removes != removesBefore {
+		t.Fatalf("removing a non-active peer must not call wg remove")
+	}
+}
+
+// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until
+// the last reference is released (the reason the per-peer count must be an int, not a set).
+func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerA")
+	if f.adds != 1 {
+		t.Fatalf("expected a single wg add for the same peer, got %d", f.adds)
+	}
+
+	mustDecrement(t, c, p, "peerA")
+	if f.installed[p] != "peerA" {
+		t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p])
+	}
+	if f.removes != 0 {
+		t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes)
+	}
+
+	mustDecrement(t, c, p, "peerA")
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("prefix must be removed after last reference")
+	}
+}
+
+// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log).
+func TestAllowedIPs_RefCountAndActive(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	ref := mustIncrement(t, c, p, "peerA")
+	if ref.Count != 1 || ref.Out != "peerA" {
+		t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out)
+	}
+	ref = mustIncrement(t, c, p, "peerB")
+	if ref.Count != 2 || ref.Out != "peerA" {
+		t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out)
+	}
+}
+
+// TestAllowedIPs_Flush removes everything installed and clears the counter.
+func TestAllowedIPs_Flush(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p1 := mustPrefix(t, "10.44.8.0/24")
+	p2 := mustPrefix(t, "10.44.9.0/24")
+
+	mustIncrement(t, c, p1, "peerA")
+	mustIncrement(t, c, p2, "peerB")
+
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+	if len(f.installed) != 0 {
+		t.Fatalf("expected all prefixes removed, got %v", f.installed)
+	}
+	// After flush, a fresh increment must add again.
+	mustIncrement(t, c, p1, "peerC")
+	if f.installed[p1] != "peerC" {
+		t.Fatalf("counter not reset after flush")
+	}
+}
+
+// TestAllowedIPs_SelfHealAfterSwapAddError ensures a failed add during a swap does not permanently
+// strand the prefix: the next Decrement (or Increment) must retry and install a surviving peer.
+func TestAllowedIPs_SelfHealAfterSwapAddError(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+	mustIncrement(t, c, p, "peerC")
+
+	// Removing the active peerA triggers a swap to a survivor; make the add fail once.
+	f.failAdd = true
+	if _, err := c.Decrement(p, "peerA"); err == nil {
+		t.Fatalf("expected error from failed swap add")
+	}
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("nothing should be installed after a failed swap add, got %q", f.installed[p])
+	}
+
+	// A later Decrement of a non-active survivor must retry the hand-off (self-heal), not stay stuck.
+	ref := mustDecrement(t, c, p, "peerC")
+	if got := f.installed[p]; got == "" {
+		t.Fatalf("self-heal failed: prefix left unrouted after add recovered")
+	}
+	if ref.Out == "" {
+		t.Fatalf("expected an active peer after self-heal, got empty")
+	}
+}
+
+// TestAllowedIPs_SelfHealAfterRemoveError ensures a failed remove during a swap is retried instead
+// of leaving e.active stuck on a peer that no longer holds references.
+func TestAllowedIPs_SelfHealAfterRemoveError(t *testing.T) {
+	f := newFakeWG()
+	c := f.counter()
+	p := mustPrefix(t, "10.44.8.0/24")
+
+	mustIncrement(t, c, p, "peerA")
+	mustIncrement(t, c, p, "peerB")
+
+	// Releasing active peerA must detach it (remove) then add peerB; fail the remove once.
+	f.failRemove = true
+	if _, err := c.Decrement(p, "peerA"); err == nil {
+		t.Fatalf("expected error from failed remove")
+	}
+
+	// Next Decrement of the non-active survivor retries: removes stale peerA, installs peerB.
+	mustDecrement(t, c, p, "peerB")
+	// peerB had only one ref, so after retry the prefix is fully released.
+	if _, ok := f.installed[p]; ok {
+		t.Fatalf("expected prefix released after self-heal, still on %q", f.installed[p])
+	}
+}

client/internal/routemanager/refcounter/types.go

@@ -5,5 +5,7 @@ import "net/netip"
 // RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement
 type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}]
 
-// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement
-type AllowedIPsRefCounter = Counter[netip.Prefix, string, string]
+// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware:
+// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer,
+// so the counter records the per-peer reference count and swaps the installed peer when the active one is released.
+// See allowedips.go.

client/internal/routemanager/static/route.go

@@ -15,6 +15,11 @@ type Route struct {
 	route                *route.Route
 	routeRefCounter      *refcounter.RouteRefCounter
 	allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
+	// currentPeerKey is the routing peer this watcher currently has the prefix installed on
+	// (the HA winner elected by the watcher). It can differ from route.Peer and change on
+	// failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement
+	// the exact peer that was incremented.
+	currentPeerKey string
 }
 
 func NewRoute(params common.HandlerParams) *Route {
@@ -52,12 +57,15 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
 			ref.Out,
 		)
 	}
+	r.currentPeerKey = peerKey
 	return nil
 }
 
 func (r *Route) RemoveAllowedIPs() error {
-	if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
-		return err
+	var err error
+	if _, decErr := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); decErr != nil {
+		err = fmt.Errorf("remove allowed IP %s: %w", r.route.Network, decErr)
 	}
-	return nil
+	r.currentPeerKey = ""
+	return err
 }

client/internal/statemanager/manager.go

@@ -1,7 +1,6 @@
 package statemanager
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -306,11 +305,6 @@ func (m *Manager) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage,
 
 	var rawStates map[string]json.RawMessage
 	if err := json.Unmarshal(data, &rawStates); err != nil {
-		if len(bytes.TrimSpace(data)) == 0 {
-			log.Warnf("state file %s is empty (%d bytes)", m.filePath, len(data))
-		} else {
-			log.Warnf("state file %s has malformed content (%d bytes)", m.filePath, len(data))
-		}
 		m.handleCorruptedState(deleteCorrupt)
 		return nil, fmt.Errorf("unmarshal states: %w", err)
 	}

client/ios/NetBirdSDK/client.go

@@ -233,9 +233,6 @@ func (c *Client) DebugBundle(anonymize bool) (string, error) {
 		deps.SyncResponse = resp
 
 		if e := cc.Engine(); e != nil {
-			deps.RefreshStatus = func() {
-				e.RunHealthProbes(context.Background(), true)
-			}
 			if cm := e.GetClientMetrics(); cm != nil {
 				deps.ClientMetrics = cm
 			}

client/ios/NetBirdSDK/login.go

@@ -44,25 +44,10 @@ type Auth struct {
 // NewAuth instantiate Auth struct and validate the management URL
 func NewAuth(cfgPath string, mgmURL string) (*Auth, error) {
 	inputCfg := profilemanager.ConfigInput{
-		ConfigPath:    cfgPath,
 		ManagementURL: mgmURL,
 	}
 
-	// Load the existing config when a config file is already present so an
-	// interactive re-login reuses the peer's persisted WireGuard private key
-	// (and thus its identity) instead of generating a fresh one. Generating a
-	// new key registers a brand-new peer on the management server on every
-	// re-auth (named after the fallback hostname). Only fall back to a fresh
-	// in-memory config for the first-time login when no config file exists yet.
-	// DirectUpdateOrCreateConfig uses non-atomic writes so it also works inside
-	// the tvOS App Group sandbox where atomic temp-file+rename is blocked.
-	var cfg *profilemanager.Config
-	var err error
-	if cfgPath != "" {
-		cfg, err = profilemanager.DirectUpdateOrCreateConfig(inputCfg)
-	} else {
-		cfg, err = profilemanager.CreateInMemoryConfig(inputCfg)
-	}
+	cfg, err := profilemanager.CreateInMemoryConfig(inputCfg)
 	if err != nil {
 		return nil, err
 	}

client/ssh/config/manager.go

@@ -14,7 +14,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	nbssh "github.com/netbirdio/netbird/client/ssh"
-	"github.com/netbirdio/netbird/shared/management/domain"
 )
 
 const (
@@ -219,20 +218,11 @@ func (m *Manager) buildHostPatterns(peer PeerSSHInfo) []string {
 	if peer.IPv6.IsValid() {
 		hostPatterns = append(hostPatterns, peer.IPv6.String())
 	}
-	// Peer FQDNs and hostnames originate from remote peers, so they must be
-	// validated as plain DNS names before being embedded in the ssh_config
-	// "Match host" pattern list. This prevents injection of arbitrary
-	// ssh_config directives via embedded quotes, whitespace, newlines, the
-	// comma pattern separator, or the "*"/"?" pattern metacharacters.
-	if domain.IsValidDomainNoWildcard(peer.FQDN) {
+	if peer.FQDN != "" {
 		hostPatterns = append(hostPatterns, peer.FQDN)
-	} else if peer.FQDN != "" {
-		log.Warnf("skipping peer FQDN with invalid characters in SSH config: %q", peer.FQDN)
 	}
-	if peer.Hostname != peer.FQDN && domain.IsValidDomainNoWildcard(peer.Hostname) {
+	if peer.Hostname != "" && peer.Hostname != peer.FQDN {
 		hostPatterns = append(hostPatterns, peer.Hostname)
-	} else if peer.Hostname != "" && peer.Hostname != peer.FQDN {
-		log.Warnf("skipping peer hostname with invalid characters in SSH config: %q", peer.Hostname)
 	}
 	return hostPatterns
 }

client/ssh/config/manager_test.go

@@ -148,45 +148,6 @@ func TestManager_MatchHostFormat(t *testing.T) {
 		"should use Match host with comma-separated patterns")
 }
 
-func TestManager_HostPatternInjection(t *testing.T) {
-	tempDir, err := os.MkdirTemp("", "netbird-ssh-config-test")
-	require.NoError(t, err)
-	defer func() { assert.NoError(t, os.RemoveAll(tempDir)) }()
-
-	manager := &Manager{
-		sshConfigDir:  filepath.Join(tempDir, "ssh_config.d"),
-		sshConfigFile: "99-netbird.conf",
-	}
-
-	// A malicious peer FQDN/hostname attempts to break out of the Match host
-	// directive and inject arbitrary ssh_config (a ProxyCommand executing a
-	// command). It must be rejected, not written to the config.
-	peers := []PeerSSHInfo{
-		{
-			Hostname: "evil\"\n    ProxyCommand touch /tmp/pwned\nHost x",
-			IP:       netip.MustParseAddr("100.125.1.1"),
-			FQDN:     "evil\"\n    ProxyCommand touch /tmp/pwned\nHost x.nb.internal",
-		},
-		{Hostname: "peer2", IP: netip.MustParseAddr("100.125.1.2"), FQDN: "peer2.nb.internal"},
-	}
-
-	err = manager.SetupSSHClientConfig(peers)
-	require.NoError(t, err)
-
-	configPath := filepath.Join(manager.sshConfigDir, manager.sshConfigFile)
-	content, err := os.ReadFile(configPath)
-	require.NoError(t, err)
-	configStr := string(content)
-
-	assert.NotContains(t, configStr, "ProxyCommand touch /tmp/pwned",
-		"injected directive must not appear in generated config")
-	assert.NotContains(t, configStr, "evil",
-		"malicious pattern must be dropped entirely")
-	// The valid peer must still be present, on a single Match host line.
-	assert.Contains(t, configStr, "Match host \"100.125.1.1,100.125.1.2,peer2.nb.internal,peer2\"",
-		"valid peers must survive, injected patterns dropped")
-}
-
 func TestManager_ForcedSSHConfig(t *testing.T) {
 	// Set force environment variable
 	t.Setenv(EnvForceSSHConfig, "true")

client/ssh/server/getent_unix.go

@@ -69,8 +69,7 @@ func parseGetentPasswd(output string) (*user.User, string, error) {
 
 // validateGetentInput checks that the input is safe to pass to getent or id.
 // Allows POSIX usernames, numeric UIDs, and common NSS extensions
-// (@ for Kerberos, $ for Samba, + for NIS compat). A leading hyphen is
-// rejected so the input can never be parsed as a command-line flag.
+// (@ for Kerberos, $ for Samba, + for NIS compat).
 func validateGetentInput(input string) bool {
 	maxLen := 32
 	if runtime.GOOS == "linux" {
@@ -81,10 +80,6 @@ func validateGetentInput(input string) bool {
 		return false
 	}
 
-	if input[0] == '-' {
-		return false
-	}
-
 	for _, r := range input {
 		if isAllowedGetentChar(r) {
 			continue

client/ssh/server/getent_unix_test.go

@@ -157,9 +157,6 @@ func TestValidateGetentInput(t *testing.T) {
 		{"numeric UID", "1001", true},
 		{"dots and underscores", "alice.bob_test", true},
 		{"hyphen", "alice-bob", true},
-		{"leading hyphen rejected", "-i", false},
-		{"leading double hyphen rejected", "--no-idn", false},
-		{"lone hyphen rejected", "-", false},
 		{"kerberos principal", "user@REALM", true},
 		{"samba machine account", "MACHINE$", true},
 		{"NIS compat", "+user", true},

client/ui/i18n/locales/_index.json

@@ -8,7 +8,6 @@
         {"code": "fr", "displayName": "Français", "englishName": "French"},
         {"code": "it", "displayName": "Italiano", "englishName": "Italian"},
         {"code": "pt", "displayName": "Português", "englishName": "Portuguese"},
-        {"code": "zh-CN", "displayName": "简体中文", "englishName": "Simplified Chinese"},
-        {"code": "ja", "displayName": "日本語", "englishName": "Japanese"}
+        {"code": "zh-CN", "displayName": "简体中文", "englishName": "Simplified Chinese"}
     ]
 }

combined/cmd/config.go

@@ -145,7 +145,6 @@ type AuthConfig struct {
 	CLIRedirectURIs                 []string          `yaml:"cliRedirectURIs"`
 	Owner                           *AuthOwnerConfig  `yaml:"owner,omitempty"`
 	DashboardPostLogoutRedirectURIs []string          `yaml:"dashboardPostLogoutRedirectURIs"`
-	GrantTypes                      []string          `yaml:"grantTypes"`
 }
 
 // AuthStorageConfig contains auth storage settings
@@ -605,7 +604,6 @@ func (c *CombinedConfig) buildEmbeddedIdPConfig(mgmt ManagementConfig) (*idp.Emb
 		DashboardRedirectURIs:           mgmt.Auth.DashboardRedirectURIs,
 		CLIRedirectURIs:                 mgmt.Auth.CLIRedirectURIs,
 		DashboardPostLogoutRedirectURIs: mgmt.Auth.DashboardPostLogoutRedirectURIs,
-		GrantTypes:                      mgmt.Auth.GrantTypes,
 	}
 
 	if mgmt.Auth.Owner != nil && mgmt.Auth.Owner.Email != "" {

combined/cmd/root.go

@@ -226,7 +226,7 @@ func (s *serverInstances) createRelayServer(cfg *CombinedConfig, tlsSupport bool
 	}
 
 	hashedSecret := sha256.Sum256([]byte(cfg.Relay.AuthSecret))
-	authenticator := auth.NewTimedHMACValidator(hashedSecret[:])
+	authenticator := auth.NewTimedHMACValidator(hashedSecret[:], 24*time.Hour)
 
 	relayCfg := relayServer.Config{
 		Meter:          s.metricsServer.Meter,

go.mod

@@ -113,7 +113,7 @@ require (
 	github.com/ti-mo/conntrack v0.5.1
 	github.com/ti-mo/netfilter v0.5.2
 	github.com/vmihailenco/msgpack/v5 v5.4.1
-	github.com/wailsapp/wails/v3 v3.0.0-alpha2.117
+	github.com/wailsapp/wails/v3 v3.0.0-alpha2.111
 	github.com/yusufpapurcu/wmi v1.2.4
 	github.com/zcalusic/sysinfo v1.1.3
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0
@@ -303,6 +303,7 @@ require (
 	github.com/tklauser/numcpus v0.10.0 // indirect
 	github.com/vishvananda/netns v0.0.5 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
+	github.com/wailsapp/wails/webview2 v1.0.27 // indirect
 	github.com/wlynxg/anet v0.0.5 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
 	github.com/zeebo/blake3 v0.2.3 // indirect
@@ -334,7 +335,7 @@ replace github.com/cloudflare/circl => codeberg.org/cunicu/circl v0.0.0-20230801
 
 replace github.com/pion/ice/v4 => github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51
 
-replace github.com/dexidp/dex => github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5
+replace github.com/dexidp/dex => github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1
 
 replace github.com/dexidp/dex/api/v2 => github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1
 

go.sum

@@ -476,8 +476,8 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
 github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
-github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5 h1:3PwQv8aR46qN2u16+Dv6udnH3sbVKX5KrGwF35CKSI0=
-github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5/go.mod h1:IHH+H8vK2GfqtIt5u/5OdPh18yk0oDHuj2vz5+Goetg=
+github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1 h1:4TaYr9O4xX0D2kszeOLclTiCbA3eHq3xWV+9ILJbIYs=
+github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1/go.mod h1:IHH+H8vK2GfqtIt5u/5OdPh18yk0oDHuj2vz5+Goetg=
 github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1 h1:neE7z+FPUkldl3faK/Jt+hJK2L+1XfQ1W33TQhU9m88=
 github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1/go.mod h1:awuTyT29CYALpEyET0S307EgNlPWrc7fFKRAyhsO45M=
 github.com/netbirdio/easyjson v0.9.0 h1:6Nw2lghSVuy8RSkAYDhDv1thBVEmfVbKZnV7T7Z6Aus=
@@ -660,8 +660,10 @@ github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IU
 github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
 github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
 github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
-github.com/wailsapp/wails/v3 v3.0.0-alpha2.117 h1:udyjqPG3AIgkod5QDR/WblCkpV8R86BFPSrsWxSyt5Y=
-github.com/wailsapp/wails/v3 v3.0.0-alpha2.117/go.mod h1:74WH2FScMsgucZvHHvv7eOefDXCm/CjuIxqhhZgPhKg=
+github.com/wailsapp/wails/v3 v3.0.0-alpha2.111 h1:MKx1nOnhnDuEGrRBmtxLOJq1NERwailu2cI4BvzWhi4=
+github.com/wailsapp/wails/v3 v3.0.0-alpha2.111/go.mod h1:wrdvmyeCsB/K3YqJDoH8E3MwcN8NXAMnEFaDTW46w60=
+github.com/wailsapp/wails/webview2 v1.0.27 h1:wjgAi/I8BBZ7kUGU8um3XF3ILEfzr96Q2Q1G4GPjMns=
+github.com/wailsapp/wails/webview2 v1.0.27/go.mod h1:zdM4jcO1IaC61RiJL5F1BzgoqBHFIdacz8gPr5exr0o=
 github.com/wlynxg/anet v0.0.5 h1:J3VJGi1gvo0JwZ/P1/Yc/8p63SoW98B5dHkYDmpgvvU=
 github.com/wlynxg/anet v0.0.5/go.mod h1:eay5PRQr7fIVAMbTbchTnO9gG65Hg/uYGdc7mguHxoA=
 github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=

idp/dex/config.go

@@ -613,10 +613,6 @@ func (c *YAMLConfig) ToServerConfig(stor storage.Storage, logger *slog.Logger) s
 		cfg.SupportedResponseTypes = c.OAuth2.ResponseTypes
 	}
 
-	if len(c.OAuth2.GrantTypes) > 0 {
-		cfg.AllowedGrantTypes = c.OAuth2.GrantTypes
-	}
-
 	// Apply expiry settings
 	if c.Expiry.IDTokens != "" {
 		if d, err := parseDuration(c.Expiry.IDTokens); err == nil {

idp/dex/provider.go

@@ -21,7 +21,7 @@ import (
 	"github.com/dexidp/dex/server/signer"
 	"github.com/dexidp/dex/storage"
 	"github.com/dexidp/dex/storage/sql"
-	"github.com/go-jose/go-jose/v4"
+	jose "github.com/go-jose/go-jose/v4"
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/crypto/bcrypt"

idp/dex/provider_test.go

@@ -595,90 +595,3 @@ enablePasswordDB: true
 	assert.True(t, cfg.ContinueOnConnectorFailure,
 		"buildDexConfig must set ContinueOnConnectorFailure to true so management starts even if an external IdP is down")
 }
-
-func TestToServerConfig_WiresGrantTypes(t *testing.T) {
-	tmpDir, err := os.MkdirTemp("", "dex-grants-*")
-	require.NoError(t, err)
-	defer os.RemoveAll(tmpDir)
-
-	stor := openTestStorage(t, tmpDir)
-	defer stor.Close()
-
-	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
-
-	grants := []string{"authorization_code", "refresh_token"}
-	cfg := &YAMLConfig{Issuer: "http://localhost:5599/oauth2", OAuth2: OAuth2{GrantTypes: grants}}
-	assert.Equal(t, grants, cfg.ToServerConfig(stor, logger).AllowedGrantTypes)
-
-	empty := &YAMLConfig{Issuer: "http://localhost:5599/oauth2"}
-	assert.Empty(t, empty.ToServerConfig(stor, logger).AllowedGrantTypes)
-}
-
-func newDeviceGuardProvider(t *testing.T, grantTypesYAML string) *Provider {
-	t.Helper()
-
-	tmpDir, err := os.MkdirTemp("", "dex-devguard-*")
-	require.NoError(t, err)
-	t.Cleanup(func() { _ = os.RemoveAll(tmpDir) })
-
-	yamlContent := `
-issuer: http://localhost:5599/oauth2
-storage:
-  type: sqlite3
-  config:
-    file: ` + filepath.Join(tmpDir, "dex.db") + `
-web:
-  http: 127.0.0.1:5599
-enablePasswordDB: true
-` + grantTypesYAML
-
-	configPath := filepath.Join(tmpDir, "config.yaml")
-	require.NoError(t, os.WriteFile(configPath, []byte(yamlContent), 0644))
-
-	yamlConfig, err := LoadConfig(configPath)
-	require.NoError(t, err)
-
-	provider, err := NewProviderFromYAML(context.Background(), yamlConfig)
-	require.NoError(t, err)
-	t.Cleanup(func() { _ = provider.Stop(context.Background()) })
-	return provider
-}
-
-func TestHandler_BlocksDeviceEndpointsWhenDeviceGrantDisabled(t *testing.T) {
-	provider := newDeviceGuardProvider(t, `
-oauth2:
-  grantTypes:
-    - authorization_code
-    - refresh_token
-`)
-
-	devicePaths := []string{
-		"/oauth2/device",
-		"/oauth2/device/code",
-		"/oauth2/device/token",
-		"/oauth2/device/auth/verify_code",
-		"/oauth2/device/callback",
-	}
-	for _, path := range devicePaths {
-		for _, method := range []string{http.MethodGet, http.MethodPost} {
-			req := httptest.NewRequest(method, path, nil)
-			rec := httptest.NewRecorder()
-			provider.Handler().ServeHTTP(rec, req)
-			assert.Equal(t, http.StatusNotFound, rec.Code, "%s %s must be blocked", method, path)
-		}
-	}
-
-	req := httptest.NewRequest(http.MethodGet, "/oauth2/.well-known/openid-configuration", nil)
-	rec := httptest.NewRecorder()
-	provider.Handler().ServeHTTP(rec, req)
-	assert.Equal(t, http.StatusOK, rec.Code)
-}
-
-func TestHandler_AllowsDeviceEndpointsWhenGrantsDefault(t *testing.T) {
-	provider := newDeviceGuardProvider(t, "")
-
-	req := httptest.NewRequest(http.MethodPost, "/oauth2/device/code", nil)
-	rec := httptest.NewRecorder()
-	provider.Handler().ServeHTTP(rec, req)
-	assert.NotEqual(t, http.StatusNotFound, rec.Code)
-}

management/internals/modules/agentnetwork/catalog/catalog.go

@@ -197,12 +197,6 @@ type JSONMetadataInjection struct {
 	// enforces a 128-char limit per value; oversized values are
 	// truncated rather than failing the request. 0 disables the cap.
 	MaxValueLength int
-	// Sanitize, when true, replaces characters outside the destination's
-	// accepted set with '_' before emitting each value. AWS Bedrock's
-	// X-Amzn-Bedrock-Request-Metadata restricts values to a limited character
-	// class, so unsanitized group display names (e.g. containing spaces) would
-	// make Bedrock reject the request with 400.
-	Sanitize bool
 }
 
 // providers is the canonical list of supported Agent Network providers.
@@ -335,18 +329,6 @@ var providers = []Provider{
 			{ID: "amazon.nova-lite", Label: "Amazon Nova Lite (Bedrock)", InputPer1k: 0.00006, OutputPer1k: 0.00024, ContextWindow: 300000},
 			{ID: "amazon.nova-micro", Label: "Amazon Nova Micro (Bedrock)", InputPer1k: 0.000035, OutputPer1k: 0.00014, ContextWindow: 128000},
 		},
-		// Bedrock accepts a cost-allocation metadata header; stamp the caller's
-		// user + authorizing group so spend can be attributed in AWS Cost
-		// Management. Sanitized because Bedrock restricts the value character set.
-		IdentityInjection: &IdentityInjection{
-			JSONMetadata: &JSONMetadataInjection{
-				Header:         "X-Amzn-Bedrock-Request-Metadata",
-				UserKey:        "user",
-				GroupsKey:      "group",
-				MaxValueLength: 256,
-				Sanitize:       true,
-			},
-		},
 	},
 	{
 		ID:                 "vertex_ai_api",

management/internals/modules/agentnetwork/synthesizer.go

@@ -540,7 +540,6 @@ type identityInjectJSONMetadata struct {
 	UserKey        string `json:"user_key,omitempty"`
 	GroupsKey      string `json:"groups_key,omitempty"`
 	MaxValueLength int    `json:"max_value_length,omitempty"`
-	Sanitize       bool   `json:"sanitize,omitempty"`
 }
 
 // buildIdentityInjectConfigJSON walks the enabled providers and emits
@@ -584,11 +583,9 @@ func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[s
 func buildIdentityInjectRule(p *types.Provider, entry catalog.Provider) (identityInjectProvider, bool) {
 	rule := identityInjectProvider{ProviderID: p.ID}
 	// Identity-stamping shape (one of HeaderPair / JSONMetadata). Skip the
-	// shape silently when the catalog entry doesn't declare one, or when the
-	// operator disabled metadata for this provider — extras can still apply,
-	// see below. MetadataDisabled suppresses only the identity dimensions
-	// (user + authorizing group), not the catalog's routing ExtraHeaders.
-	if !p.MetadataDisabled && entry.IdentityInjection != nil {
+	// shape silently when the catalog entry doesn't declare one — extras
+	// can still apply, see below.
+	if entry.IdentityInjection != nil {
 		switch {
 		case entry.IdentityInjection.HeaderPair != nil:
 			rule.HeaderPair = buildIdentityHeaderPair(p, entry.IdentityInjection.HeaderPair)
@@ -654,7 +651,6 @@ func buildIdentityJSONMetadata(p *types.Provider, jm *catalog.JSONMetadataInject
 		UserKey:        userKey,
 		GroupsKey:      groupsKey,
 		MaxValueLength: jm.MaxValueLength,
-		Sanitize:       jm.Sanitize,
 	}
 }
 

management/internals/modules/agentnetwork/synthesizer_test.go

@@ -698,94 +698,6 @@ func TestSynthesizeServices_IdentityInject_Portkey_NotCustomizable(t *testing.T)
 		"same fixed-schema guarantee for the groups dimension")
 }
 
-// TestSynthesizeServices_IdentityInject_Bedrock pins Bedrock's cost-allocation
-// metadata: a JSONMetadata shape emitting X-Amzn-Bedrock-Request-Metadata with
-// the reserved user/group keys, sanitized to Bedrock's accepted charset.
-func TestSynthesizeServices_IdentityInject_Bedrock(t *testing.T) {
-	ctx := context.Background()
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-	mockStore := store.NewMockStore(ctrl)
-
-	br := newSynthTestProvider()
-	br.ID = "prov-bedrock"
-	br.ProviderID = "bedrock_api"
-	br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
-	br.APIKey = "bedrock-bearer"
-	br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
-
-	policy := newSynthTestPolicy(br.ID, "grp-eng", "")
-	policy.ID = "pol-bedrock"
-
-	expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
-		[]*types.Provider{br},
-		[]*types.Policy{policy},
-		[]*types.Guardrail{})
-
-	services, err := SynthesizeServices(ctx, mockStore, testAccountID)
-	require.NoError(t, err)
-	require.Len(t, services, 1)
-
-	var injectCfg identityInjectConfig
-	for _, m := range services[0].Targets[0].Options.Middlewares {
-		if m.ID == middlewareIDLLMIdentityInject {
-			require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
-			break
-		}
-	}
-	require.Len(t, injectCfg.Providers, 1)
-	entry := injectCfg.Providers[0]
-	require.NotNil(t, entry.JSONMetadata, "Bedrock uses the JSONMetadata shape for cost-allocation metadata")
-	assert.Nil(t, entry.HeaderPair, "shapes are mutually exclusive")
-	assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", entry.JSONMetadata.Header,
-		"the caller identity lands in Bedrock's cost-allocation metadata header")
-	assert.Equal(t, "user", entry.JSONMetadata.UserKey)
-	assert.Equal(t, "group", entry.JSONMetadata.GroupsKey)
-	assert.True(t, entry.JSONMetadata.Sanitize,
-		"Bedrock restricts the metadata value charset, so values must be sanitized")
-}
-
-// TestSynthesizeServices_MetadataDisabled_SuppressesInjection verifies the
-// per-provider opt-out: a provider with MetadataDisabled set emits no
-// identity-inject entry (Bedrock has no catalog ExtraHeaders, so the whole
-// entry is dropped).
-func TestSynthesizeServices_MetadataDisabled_SuppressesInjection(t *testing.T) {
-	ctx := context.Background()
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-	mockStore := store.NewMockStore(ctrl)
-
-	br := newSynthTestProvider()
-	br.ID = "prov-bedrock"
-	br.ProviderID = "bedrock_api"
-	br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
-	br.APIKey = "bedrock-bearer"
-	br.MetadataDisabled = true
-	br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
-
-	policy := newSynthTestPolicy(br.ID, "grp-eng", "")
-	policy.ID = "pol-bedrock"
-
-	expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
-		[]*types.Provider{br},
-		[]*types.Policy{policy},
-		[]*types.Guardrail{})
-
-	services, err := SynthesizeServices(ctx, mockStore, testAccountID)
-	require.NoError(t, err)
-	require.Len(t, services, 1)
-
-	var injectCfg identityInjectConfig
-	for _, m := range services[0].Targets[0].Options.Middlewares {
-		if m.ID == middlewareIDLLMIdentityInject {
-			require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
-			break
-		}
-	}
-	assert.Empty(t, injectCfg.Providers,
-		"metadata_disabled must drop the provider's identity-inject entry")
-}
-
 // TestSynthesizeServices_IdentityInject_Vercel pins Vercel AI
 // Gateway's wiring: HeaderPair shape with fixed wire names dictated
 // by Vercel's Custom Reporting API (ai-reporting-user /

management/internals/modules/agentnetwork/types/provider.go

@@ -51,12 +51,6 @@ type Provider struct {
 	// private or self-signed certificate. The synthesiser propagates it into
 	// the router route so the proxy dials that provider's upstream insecurely.
 	SkipTLSVerification bool `gorm:"column:skip_tls_verification"`
-	// MetadataDisabled suppresses identity metadata injection for this provider.
-	// Metadata (the caller's user + authorizing group) is injected by default;
-	// when true the synthesiser omits the provider's identity-inject shape, so no
-	// user/group headers (e.g. Bedrock's X-Amzn-Bedrock-Request-Metadata) are
-	// stamped. Catalog ExtraHeaders (routing config) are unaffected.
-	MetadataDisabled bool `gorm:"column:metadata_disabled"`
 	// SessionPrivateKey + SessionPublicKey are the ed25519 keypair the
 	// synthesised reverse-proxy service uses to sign / verify session
 	// JWTs after a successful OIDC handshake. Generated once on
@@ -143,9 +137,6 @@ func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) {
 	if req.SkipTlsVerification != nil {
 		p.SkipTLSVerification = *req.SkipTlsVerification
 	}
-	if req.MetadataDisabled != nil {
-		p.MetadataDisabled = *req.MetadataDisabled
-	}
 	// Identity-header overrides for catalogs flagged Customizable.
 	// nil pointer = "field omitted on the wire" → leave the stored
 	// value untouched (per the openapi description). Empty string is
@@ -179,7 +170,6 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
 		Models:              models,
 		Enabled:             p.Enabled,
 		SkipTlsVerification: p.SkipTLSVerification,
-		MetadataDisabled:    p.MetadataDisabled,
 		CreatedAt:           &created,
 		UpdatedAt:           &updated,
 	}

management/internals/modules/agentnetwork/types/provider_test.go

@@ -42,38 +42,3 @@ func TestProvider_SkipTLSVerification_RoundTrip(t *testing.T) {
 	assert.False(t, p.SkipTLSVerification, "explicit false must clear skip_tls_verification")
 	assert.False(t, p.ToAPIResponse().SkipTlsVerification, "response must reflect the cleared value")
 }
-
-// TestProvider_MetadataDisabled_RoundTrip covers the request→provider→response
-// mapping of metadata_disabled, with the same update semantics: nil preserves,
-// explicit false clears.
-func TestProvider_MetadataDisabled_RoundTrip(t *testing.T) {
-	enable := true
-	disable := false
-
-	base := func() *api.AgentNetworkProviderRequest {
-		return &api.AgentNetworkProviderRequest{
-			ProviderId:  "bedrock_api",
-			Name:        "bedrock",
-			UpstreamUrl: "https://bedrock-runtime.us-east-1.amazonaws.com",
-		}
-	}
-
-	p := NewProvider("acc-1")
-
-	req := base()
-	req.MetadataDisabled = &enable
-	p.FromAPIRequest(req)
-	assert.True(t, p.MetadataDisabled, "create with metadata_disabled=true must set the field")
-	assert.True(t, p.ToAPIResponse().MetadataDisabled, "response must surface metadata_disabled")
-
-	// Omitting the field on update leaves the stored value untouched.
-	p.FromAPIRequest(base())
-	assert.True(t, p.MetadataDisabled, "omitting metadata_disabled on update must preserve it")
-
-	// Explicit false clears it (re-enables metadata).
-	req = base()
-	req.MetadataDisabled = &disable
-	p.FromAPIRequest(req)
-	assert.False(t, p.MetadataDisabled, "explicit false must clear metadata_disabled")
-	assert.False(t, p.ToAPIResponse().MetadataDisabled, "response must reflect the cleared value")
-}

management/server/account_test.go

@@ -4269,7 +4269,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_NetworkRangePreserved(t *te
 	}
 
 	// Sanity: an actually different range still triggers reallocation.
-	newRange := netip.MustParsePrefix("100.60.0.0/16")
+	newRange := netip.MustParsePrefix("100.99.0.0/16")
 	_, err = manager.UpdateAccountSettings(ctx, account.Id, userID, &types.Settings{
 		PeerLoginExpirationEnabled: true,
 		PeerLoginExpiration:        types.DefaultPeerLoginExpiration,

management/server/idp/embedded.go

@@ -76,9 +76,6 @@ type EmbeddedIdPConfig struct {
 	DashboardPostLogoutRedirectURIs []string
 	// StaticConnectors are additional connectors to seed during initialization
 	StaticConnectors []dex.Connector
-	// GrantTypes restricts allowed OAuth2 grants; empty means all (Dex default). Omit the
-	// device_code grant to disable the device flow; keep authorization_code and refresh_token.
-	GrantTypes []string
 }
 
 // EmbeddedStorageConfig holds storage configuration for the embedded IdP.
@@ -178,7 +175,6 @@ func (c *EmbeddedIdPConfig) ToYAMLConfig() (*dex.YAMLConfig, error) {
 		},
 		OAuth2: dex.OAuth2{
 			SkipApprovalScreen: true,
-			GrantTypes:         c.GrantTypes,
 		},
 		Frontend: dex.Frontend{
 			Issuer: "NetBird",

management/server/store/sql_store.go

@@ -1606,8 +1606,7 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
 			settings_routing_peer_dns_resolution_enabled, settings_dns_domain, settings_network_range,
 			settings_network_range_v6, settings_ipv6_enabled_groups, settings_lazy_connection_enabled,
 			settings_local_mfa_enabled, settings_metrics_push_enabled, settings_agent_network_only,
-			settings_dashboard_features, settings_auto_update_version, settings_auto_update_always,
-			settings_peer_expose_enabled, settings_peer_expose_groups,
+			settings_dashboard_features,
 			-- Embedded ExtraSettings
 			settings_extra_peer_approval_enabled, settings_extra_user_approval_required,
 			settings_extra_integrated_validator, settings_extra_integrated_validator_groups
@@ -1633,10 +1632,6 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
 		sMetricsPushEnabled              sql.NullBool
 		sAgentNetworkOnly                sql.NullBool
 		sDashboardFeatures               sql.NullString
-		autoUpdateVersion                sql.NullString
-		autoUpdateAlways                 sql.NullBool
-		peerExposeEnabled                sql.NullBool
-		peerExposeGroups                 sql.NullString
 		sExtraPeerApprovalEnabled        sql.NullBool
 		sExtraUserApprovalRequired       sql.NullBool
 		sExtraIntegratedValidator        sql.NullString
@@ -1660,8 +1655,7 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
 		&sRoutingPeerDNSResolutionEnabled, &sDNSDomain, &sNetworkRange,
 		&sNetworkRangeV6, &sIPv6EnabledGroups, &sLazyConnectionEnabled,
 		&sLocalMFAEnabled, &sMetricsPushEnabled, &sAgentNetworkOnly,
-		&sDashboardFeatures, &autoUpdateVersion, &autoUpdateAlways,
-		&peerExposeEnabled, &peerExposeGroups,
+		&sDashboardFeatures,
 		&sExtraPeerApprovalEnabled, &sExtraUserApprovalRequired,
 		&sExtraIntegratedValidator, &sExtraIntegratedValidatorGroups,
 	)
@@ -1753,18 +1747,6 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
 	if sIPv6EnabledGroups.Valid {
 		_ = json.Unmarshal([]byte(sIPv6EnabledGroups.String), &account.Settings.IPv6EnabledGroups)
 	}
-	if autoUpdateAlways.Valid {
-		account.Settings.AutoUpdateAlways = autoUpdateAlways.Bool
-	}
-	if autoUpdateVersion.Valid {
-		account.Settings.AutoUpdateVersion = autoUpdateVersion.String
-	}
-	if peerExposeEnabled.Valid {
-		account.Settings.PeerExposeEnabled = peerExposeEnabled.Bool
-	}
-	if peerExposeGroups.Valid {
-		_ = json.Unmarshal([]byte(peerExposeGroups.String), &account.Settings.PeerExposeGroups)
-	}
 
 	if sExtraPeerApprovalEnabled.Valid {
 		account.Settings.Extra.PeerApprovalEnabled = sExtraPeerApprovalEnabled.Bool

management/server/store/sql_store_test.go

@@ -9,7 +9,6 @@ import (
 	"net"
 	"net/netip"
 	"os"
-	"reflect"
 	"runtime"
 	"sort"
 	"sync"
@@ -35,7 +34,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/util"
 	nbroute "github.com/netbirdio/netbird/route"
 	"github.com/netbirdio/netbird/shared/management/status"
-	"github.com/netbirdio/netbird/shared/testing_helpers"
 	"github.com/netbirdio/netbird/util/crypt"
 )
 
@@ -298,53 +296,6 @@ func Test_SaveAccount(t *testing.T) {
 	})
 }
 
-func Test_AccountSettings_SaveAndRetrieve(t *testing.T) {
-	if runtime.GOOS == "windows" {
-		t.Skip("The SQLite store is not properly supported by Windows yet")
-	}
-
-	populateFields := testing_helpers.NewPopulateFields().WithCustomFieldSetter(
-		reflect.PointerTo(reflect.TypeOf(types.ExtraSettings{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
-			es := types.ExtraSettings{}
-			reflectedEs := reflect.ValueOf(&es).Elem()
-			n, err := this.PopulateAll(reflectedEs)
-			if err != nil {
-				return n, err
-			}
-			field.Set(reflectedEs.Addr())
-			return n, nil
-		}).WithCustomFieldSetter(
-		reflect.PointerTo(reflect.TypeOf(types.DashboardFeatures{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
-			t := true
-			df := types.DashboardFeatures{AgentNetwork: &t}
-			reflectedDf := reflect.ValueOf(&df).Elem()
-			field.Set(reflectedDf.Addr())
-			return 1, nil
-		}).WithSkippedTag("gorm", "-")
-
-	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
-		account := newAccountWithId(context.Background(), "account_id", "testuser", "")
-		setupKey, _ := types.GenerateDefaultSetupKey()
-		account.SetupKeys[setupKey.Key] = setupKey
-
-		settings := types.Settings{}
-		numOfExportedFields, err := populateFields.PopulateAll(reflect.ValueOf(&settings).Elem())
-		assert.NoError(t, err)
-		assert.Equal(t, 27, numOfExportedFields)
-		account.Settings = &settings
-
-		err = store.SaveAccount(context.Background(), account)
-		assert.NoError(t, err)
-
-		accountFromDb, err := store.GetAccount(context.Background(), account.Id)
-		assert.NoError(t, err)
-		assert.NotNil(t, accountFromDb)
-		assert.NotNil(t, accountFromDb.Settings)
-
-		assert.True(t, reflect.DeepEqual(&settings, accountFromDb.Settings), "created settings and settings retrieved from the db should match")
-	})
-}
-
 func TestSqlite_DeleteAccount(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("The SQLite store is not properly supported by Windows yet")

proxy/cmd/proxy/cmd/root.go

@@ -18,7 +18,6 @@ import (
 	"github.com/netbirdio/netbird/client/embed"
 	"github.com/netbirdio/netbird/proxy"
 	nbacme "github.com/netbirdio/netbird/proxy/internal/acme"
-	"github.com/netbirdio/netbird/trustedproxy"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -210,7 +209,7 @@ func runServer(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("invalid domain value %q: %w", proxyDomain, err)
 	}
 
-	parsedTrustedProxies, err := trustedproxy.Parse(trustedProxies)
+	parsedTrustedProxies, err := proxy.ParseTrustedProxies(trustedProxies)
 	if err != nil {
 		return fmt.Errorf("invalid --trusted-proxies: %w", err)
 	}

proxy/internal/accesslog/logger.go

@@ -16,7 +16,6 @@ import (
 	"github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/shared/management/proto"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 const (
@@ -67,7 +66,7 @@ type denyBucket struct {
 type Logger struct {
 	client         gRPCClient
 	logger         *log.Logger
-	trustedProxies *trustedproxy.List
+	trustedProxies []netip.Prefix
 
 	usageMux    sync.Mutex
 	domainUsage map[string]*domainUsage
@@ -83,7 +82,7 @@ type Logger struct {
 // NewLogger creates a new access log Logger. The trustedProxies parameter
 // configures which upstream proxy IP ranges are trusted for extracting
 // the real client IP from X-Forwarded-For headers.
-func NewLogger(client gRPCClient, logger *log.Logger, trustedProxies *trustedproxy.List) *Logger {
+func NewLogger(client gRPCClient, logger *log.Logger, trustedProxies []netip.Prefix) *Logger {
 	if logger == nil {
 		logger = log.StandardLogger()
 	}

proxy/internal/accesslog/requestip.go

@@ -4,13 +4,13 @@ import (
 	"net/http"
 	"net/netip"
 
-	"github.com/netbirdio/netbird/trustedproxy"
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
 )
 
 // extractSourceIP resolves the real client IP from the request using trusted
 // proxy configuration. When trustedProxies is non-empty and the direct
 // connection is from a trusted source, it walks X-Forwarded-For right-to-left
 // skipping trusted IPs. Otherwise it returns RemoteAddr directly.
-func extractSourceIP(r *http.Request, trustedProxies *trustedproxy.List) netip.Addr {
-	return trustedProxies.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For"))
+func extractSourceIP(r *http.Request, trustedProxies []netip.Prefix) netip.Addr {
+	return proxy.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For"), trustedProxies)
 }

proxy/internal/llm/bedrock_model.go

@@ -1,38 +0,0 @@
-package llm
-
-import (
-	"regexp"
-	"strings"
-)
-
-// bedrockRegionPrefixes are the cross-region inference-profile prefixes that
-// front a Bedrock model id (e.g. "eu.anthropic.claude-...").
-var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
-
-// bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]"
-// version/throughput suffix of a Bedrock model id.
-var bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
-
-// NormalizeBedrockModel strips an ARN wrapper, a cross-region inference-profile
-// prefix, and the version/throughput suffix from a Bedrock model id so it
-// matches the catalog/pricing key, e.g.
-// "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" -> "anthropic.claude-sonnet-4-5"
-// and the inference-profile ARN's last segment likewise. It is the single
-// source of truth shared by the request parser (which normalizes the request
-// model from the URL path) and the router (which normalizes the operator's
-// registered Bedrock model ids so both sides compare equal).
-func NormalizeBedrockModel(modelID string) string {
-	m := modelID
-	if strings.HasPrefix(m, "arn:") {
-		if i := strings.LastIndex(m, "/"); i >= 0 {
-			m = m[i+1:]
-		}
-	}
-	for _, p := range bedrockRegionPrefixes {
-		if strings.HasPrefix(m, p) {
-			m = m[len(p):]
-			break
-		}
-	}
-	return bedrockVersionSuffix.ReplaceAllString(m, "")
-}

proxy/internal/llm/bedrock_model_test.go

@@ -1,23 +0,0 @@
-package llm
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestNormalizeBedrockModel(t *testing.T) {
-	cases := map[string]string{
-		"eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
-		"us.anthropic.claude-haiku-4-5":                "anthropic.claude-haiku-4-5",
-		"us.anthropic.claude-opus-4-8-20250101-v1:0":   "anthropic.claude-opus-4-8",
-		"anthropic.claude-sonnet-4-5-20250929-v1:0":    "anthropic.claude-sonnet-4-5",
-		"meta.llama3-3-70b-instruct-v1:0":              "meta.llama3-3-70b-instruct",
-		"amazon.nova-pro-v1:0":                         "amazon.nova-pro",
-		// Inference-profile ARN — model id lives in the last path segment.
-		"arn:aws:bedrock:eu-central-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
-	}
-	for in, want := range cases {
-		require.Equal(t, want, NormalizeBedrockModel(in), "normalize %q", in)
-	}
-}

proxy/internal/middleware/builtin/llm_identity_inject/factory.go

@@ -64,11 +64,6 @@ type JSONMetadataRule struct {
 	UserKey        string `json:"user_key,omitempty"`
 	GroupsKey      string `json:"groups_key,omitempty"`
 	MaxValueLength int    `json:"max_value_length,omitempty"`
-	// Sanitize replaces characters outside the destination provider's accepted
-	// set with '_' before emitting each value. AWS Bedrock's
-	// X-Amzn-Bedrock-Request-Metadata restricts values to [A-Za-z0-9 +-=._:/@];
-	// group display names with other characters would otherwise 400.
-	Sanitize bool `json:"sanitize,omitempty"`
 }
 
 // Config is the on-wire configuration accepted by the factory. An

proxy/internal/middleware/builtin/llm_identity_inject/middleware.go

@@ -292,21 +292,15 @@ func applyJSONMetadata(rule *JSONMetadataRule, in *middleware.Input) *middleware
 	mutations := &middleware.Mutations{}
 	mutations.HeadersRemove = append(mutations.HeadersRemove, rule.Header)
 
-	emit := func(v string) string {
-		if rule.Sanitize {
-			v = sanitizeMetadataValue(v)
-		}
-		return truncate(v, rule.MaxValueLength)
-	}
 	payload := map[string]string{}
 	if rule.UserKey != "" {
 		if identity := identityFor(in); identity != "" {
-			payload[rule.UserKey] = emit(identity)
+			payload[rule.UserKey] = truncate(identity, rule.MaxValueLength)
 		}
 	}
 	if rule.GroupsKey != "" {
 		if csv := authorisingTagsCSV(in); csv != "" {
-			payload[rule.GroupsKey] = emit(csv)
+			payload[rule.GroupsKey] = truncate(csv, rule.MaxValueLength)
 		}
 	}
 	if len(payload) == 0 {
@@ -365,36 +359,6 @@ func truncate(s string, maxBytes int) string {
 	return s[:maxBytes]
 }
 
-// sanitizeMetadataValue replaces any character outside AWS Bedrock's accepted
-// request-metadata class — letters, digits, space, and + - = . _ : / @ — with
-// '_'. This keeps values (notably the groups CSV, whose commas are rejected, and
-// group display names with arbitrary characters) from making Bedrock reject the
-// request with 400. The result stays opaque to the gateway.
-func sanitizeMetadataValue(s string) string {
-	var b strings.Builder
-	b.Grow(len(s))
-	for _, r := range s {
-		if metadataCharAllowed(r) {
-			b.WriteRune(r)
-		} else {
-			b.WriteByte('_')
-		}
-	}
-	return b.String()
-}
-
-func metadataCharAllowed(r rune) bool {
-	switch {
-	case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
-		return true
-	}
-	switch r {
-	case ' ', '+', '-', '=', '.', '_', ':', '/', '@':
-		return true
-	}
-	return false
-}
-
 // tagsIDsFromAuthorising reads llm_router's authorising-groups metadata
 // (a CSV of group ids) and returns the parsed slice. Returns nil when
 // the key is absent or empty so the caller can fall back to the full

proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go

@@ -304,46 +304,6 @@ func TestInject_JSONMetadata_TruncatesValues(t *testing.T) {
 		"per-value byte length must be capped at MaxValueLength")
 }
 
-// TestInject_JSONMetadata_Sanitize pins the AWS-Bedrock sanitization path: when
-// Sanitize is set, characters outside Bedrock's accepted metadata class
-// (notably the groups CSV comma and arbitrary characters in group display
-// names) are replaced with '_' so Bedrock doesn't reject the request. Allowed
-// characters (letters, digits, spaces, and @ . _ : / + - =) pass through.
-func TestInject_JSONMetadata_Sanitize(t *testing.T) {
-	rule := ProviderInjection{
-		ProviderID: portkeyProvider,
-		JSONMetadata: &JSONMetadataRule{
-			Header:         "X-Amzn-Bedrock-Request-Metadata",
-			UserKey:        "user",
-			GroupsKey:      "group",
-			MaxValueLength: 256,
-			Sanitize:       true,
-		},
-	}
-	mw := New(Config{Providers: []ProviderInjection{rule}})
-	in := newInput(portkeyProvider, "alice", []string{"g1", "g2"})
-	in.UserEmail = "alice@example.com"
-	// Group display names carry characters Bedrock rejects (comma, '#'); the CSV
-	// join adds another comma between the two groups.
-	in.UserGroupNames = []string{"Eng,Team", "Ops#1"}
-
-	out, err := mw.Invoke(context.Background(), in)
-	require.NoError(t, err)
-	require.NotNil(t, out.Mutations)
-	require.Len(t, out.Mutations.HeadersAdd, 1)
-	added := out.Mutations.HeadersAdd[0]
-	assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", added.Key,
-		"the Bedrock cost-allocation header carries the metadata JSON")
-
-	var payload map[string]string
-	require.NoError(t, json.Unmarshal([]byte(added.Value), &payload))
-	assert.Equal(t, "alice@example.com", payload["user"],
-		"'@' and '.' are in Bedrock's accepted set and must be preserved")
-	assert.NotContains(t, payload["group"], ",", "commas must be sanitized — Bedrock rejects them")
-	assert.NotContains(t, payload["group"], "#", "disallowed characters must be sanitized")
-	assert.Contains(t, payload["group"], "Eng", "allowed characters must be preserved")
-}
-
 // TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd verifies the
 // anti-spoof Remove still fires when there's nothing to stamp.
 func TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) {

proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go

@@ -1,30 +0,0 @@
-package llm_router
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-// TestRouteClaimsModel_BedrockNormalizesCandidate guards the fix for the native
-// Bedrock routing gap: the request model reaches the router already normalized
-// (the parser strips the region/inference-profile prefix and version suffix),
-// so a provider registered with the raw inference-profile id must still match.
-func TestRouteClaimsModel_BedrockNormalizesCandidate(t *testing.T) {
-	route := ProviderRoute{Bedrock: true, Models: []string{"us.anthropic.claude-haiku-4-5"}}
-	assert.True(t, routeClaimsModel(route, "anthropic.claude-haiku-4-5"),
-		"raw region-prefixed Bedrock model must match the normalized request model")
-	assert.False(t, routeClaimsModel(route, "anthropic.claude-opus-4-8"),
-		"a model outside the provider's list must not match")
-
-	// A provider registered with the already-normalized id also matches.
-	normalized := ProviderRoute{Bedrock: true, Models: []string{"anthropic.claude-haiku-4-5"}}
-	assert.True(t, routeClaimsModel(normalized, "anthropic.claude-haiku-4-5"),
-		"normalized Bedrock model must match")
-
-	// Non-Bedrock routes keep exact matching (no prefix stripping).
-	openai := ProviderRoute{Models: []string{"gpt-4o"}}
-	assert.True(t, routeClaimsModel(openai, "gpt-4o"), "exact model must match")
-	assert.False(t, routeClaimsModel(openai, "us.gpt-4o"),
-		"non-Bedrock routes must not strip a us. prefix")
-}

proxy/internal/middleware/builtin/llm_router/middleware.go

@@ -23,7 +23,6 @@ import (
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
 
-	"github.com/netbirdio/netbird/proxy/internal/llm"
 	"github.com/netbirdio/netbird/proxy/internal/middleware"
 )
 
@@ -556,14 +555,6 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
 		if candidate == model {
 			return true
 		}
-		// Bedrock request models reach the router already normalized (the parser
-		// strips the region / inference-profile prefix and version suffix), but
-		// the operator may register the raw inference-profile id (e.g.
-		// "us.anthropic.claude-haiku-4-5"). Normalize the candidate so both sides
-		// compare equal; otherwise a native Bedrock request denies as not-routable.
-		if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
-			return true
-		}
 	}
 	return false
 }

proxy/internal/proxy/reverseproxy.go

@@ -22,7 +22,6 @@ import (
 	"github.com/netbirdio/netbird/proxy/internal/roundtrip"
 	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/proxy/web"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 type ReverseProxy struct {
@@ -30,10 +29,10 @@ type ReverseProxy struct {
 	// forwardedProto overrides the X-Forwarded-Proto header value.
 	// Valid values: "auto" (detect from TLS), "http", "https".
 	forwardedProto string
-	// trustedProxies is the set of trusted upstream proxies. When the direct
-	// connection comes from a trusted proxy, forwarding headers are preserved
-	// and appended to instead of being stripped.
-	trustedProxies *trustedproxy.List
+	// trustedProxies is a list of IP prefixes for trusted upstream proxies.
+	// When the direct connection comes from a trusted proxy, forwarding
+	// headers are preserved and appended to instead of being stripped.
+	trustedProxies []netip.Prefix
 	mappingsMux    sync.RWMutex
 	mappings       map[string]Mapping
 	logger         *log.Logger
@@ -64,7 +63,7 @@ func WithMiddlewareManager(m *middleware.Manager) Option {
 // between requested URLs and targets.
 // The internal mappings can be modified using the AddMapping
 // and RemoveMapping functions.
-func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trustedProxies *trustedproxy.List, logger *log.Logger, opts ...Option) *ReverseProxy {
+func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trustedProxies []netip.Prefix, logger *log.Logger, opts ...Option) *ReverseProxy {
 	if logger == nil {
 		logger = log.StandardLogger()
 	}
@@ -528,7 +527,7 @@ func (p *ReverseProxy) isSelfTargetLoop(r *http.Request, target *url.URL) bool {
 	if !types.IsOverlayOrigin(r.Context()) {
 		return false
 	}
-	srcIP := trustedproxy.ExtractHostIP(r.RemoteAddr)
+	srcIP := extractHostIP(r.RemoteAddr)
 	if !srcIP.IsValid() {
 		return false
 	}
@@ -579,9 +578,9 @@ func (p *ReverseProxy) rewriteFunc(target *url.URL, matchedPath string, passHost
 
 		stampNetBirdIdentity(r)
 
-		clientIP := trustedproxy.ExtractHostIP(r.In.RemoteAddr)
+		clientIP := extractHostIP(r.In.RemoteAddr)
 
-		if p.trustedProxies.Contains(clientIP) {
+		if isTrustedAddr(clientIP, p.trustedProxies) {
 			p.setTrustedForwardingHeaders(r, clientIP)
 		} else {
 			p.setUntrustedForwardingHeaders(r, clientIP)
@@ -665,7 +664,7 @@ func (p *ReverseProxy) setTrustedForwardingHeaders(r *httputil.ProxyRequest, cli
 	if realIP := r.In.Header.Get("X-Real-IP"); realIP != "" {
 		r.Out.Header.Set("X-Real-IP", realIP)
 	} else {
-		resolved := p.trustedProxies.ResolveClientIP(r.In.RemoteAddr, r.In.Header.Get("X-Forwarded-For"))
+		resolved := ResolveClientIP(r.In.RemoteAddr, r.In.Header.Get("X-Forwarded-For"), p.trustedProxies)
 		r.Out.Header.Set("X-Real-IP", resolved.String())
 	}
 

proxy/internal/proxy/reverseproxy_test.go

@@ -23,7 +23,6 @@ import (
 	"github.com/netbirdio/netbird/proxy/internal/roundtrip"
 	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/proxy/web"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 func TestRewriteFunc_HostRewriting(t *testing.T) {
@@ -303,7 +302,7 @@ func TestExtractHostIP(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			assert.Equal(t, tt.expected, trustedproxy.ExtractHostIP(tt.remoteAddr))
+			assert.Equal(t, tt.expected, extractHostIP(tt.remoteAddr))
 		})
 	}
 }
@@ -331,7 +330,7 @@ func TestExtractForwardedPort(t *testing.T) {
 
 func TestRewriteFunc_TrustedProxy(t *testing.T) {
 	target, _ := url.Parse("http://backend.internal:8080")
-	trusted := trustedproxy.FromPrefixes([]netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")})
+	trusted := []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}
 
 	t.Run("appends to X-Forwarded-For", func(t *testing.T) {
 		p := &ReverseProxy{forwardedProto: "auto", trustedProxies: trusted}

proxy/internal/proxy/trustedproxy.go

@@ -0,0 +1,81 @@
+package proxy
+
+import (
+	"net/netip"
+	"strings"
+)
+
+// IsTrustedProxy checks if the given IP string falls within any of the trusted prefixes.
+func IsTrustedProxy(ipStr string, trusted []netip.Prefix) bool {
+	addr, err := netip.ParseAddr(ipStr)
+	if err != nil || len(trusted) == 0 {
+		return false
+	}
+	return isTrustedAddr(addr.Unmap(), trusted)
+}
+
+// ResolveClientIP extracts the real client IP from X-Forwarded-For using the trusted proxy list.
+// It walks the XFF chain right-to-left, skipping IPs that match trusted prefixes.
+// The first untrusted IP is the real client.
+//
+// If the trusted list is empty or remoteAddr is not trusted, it returns the
+// remoteAddr IP directly (ignoring any forwarding headers).
+func ResolveClientIP(remoteAddr, xff string, trusted []netip.Prefix) netip.Addr {
+	remoteIP := extractHostIP(remoteAddr)
+
+	if len(trusted) == 0 || !isTrustedAddr(remoteIP, trusted) {
+		return remoteIP
+	}
+
+	if xff == "" {
+		return remoteIP
+	}
+
+	parts := strings.Split(xff, ",")
+	for i := len(parts) - 1; i >= 0; i-- {
+		ip := strings.TrimSpace(parts[i])
+		if ip == "" {
+			continue
+		}
+		addr, err := netip.ParseAddr(ip)
+		if err != nil {
+			continue
+		}
+		addr = addr.Unmap()
+		if !isTrustedAddr(addr, trusted) {
+			return addr
+		}
+	}
+
+	// All IPs in XFF are trusted; return the leftmost as best guess.
+	if first := strings.TrimSpace(parts[0]); first != "" {
+		if addr, err := netip.ParseAddr(first); err == nil {
+			return addr.Unmap()
+		}
+	}
+	return remoteIP
+}
+
+// extractHostIP parses the IP from a host:port string and returns it unmapped.
+func extractHostIP(hostPort string) netip.Addr {
+	if ap, err := netip.ParseAddrPort(hostPort); err == nil {
+		return ap.Addr().Unmap()
+	}
+	if addr, err := netip.ParseAddr(hostPort); err == nil {
+		return addr.Unmap()
+	}
+	return netip.Addr{}
+}
+
+// isTrustedAddr checks if the given address falls within any of the trusted prefixes.
+func isTrustedAddr(addr netip.Addr, trusted []netip.Prefix) bool {
+	if !addr.IsValid() {
+		return false
+	}
+	for _, prefix := range trusted {
+		if prefix.Contains(addr) {
+			return true
+		}
+	}
+	return false
+}

proxy/internal/proxy/trustedproxy_test.go

@@ -0,0 +1,129 @@
+package proxy
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsTrustedProxy(t *testing.T) {
+	trusted := []netip.Prefix{
+		netip.MustParsePrefix("10.0.0.0/8"),
+		netip.MustParsePrefix("192.168.1.0/24"),
+		netip.MustParsePrefix("fd00::/8"),
+	}
+
+	tests := []struct {
+		name    string
+		ip      string
+		trusted []netip.Prefix
+		want    bool
+	}{
+		{"empty trusted list", "10.0.0.1", nil, false},
+		{"IP within /8 prefix", "10.1.2.3", trusted, true},
+		{"IP within /24 prefix", "192.168.1.100", trusted, true},
+		{"IP outside all prefixes", "203.0.113.50", trusted, false},
+		{"boundary IP just outside prefix", "192.168.2.1", trusted, false},
+		{"unparsable IP", "not-an-ip", trusted, false},
+		{"IPv6 in trusted range", "fd00::1", trusted, true},
+		{"IPv6 outside range", "2001:db8::1", trusted, false},
+		{"empty string", "", trusted, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, IsTrustedProxy(tt.ip, tt.trusted))
+		})
+	}
+}
+
+func TestResolveClientIP(t *testing.T) {
+	trusted := []netip.Prefix{
+		netip.MustParsePrefix("10.0.0.0/8"),
+		netip.MustParsePrefix("172.16.0.0/12"),
+	}
+
+	tests := []struct {
+		name       string
+		remoteAddr string
+		xff        string
+		trusted    []netip.Prefix
+		want       netip.Addr
+	}{
+		{
+			name:       "empty trusted list returns RemoteAddr",
+			remoteAddr: "203.0.113.50:9999",
+			xff:        "1.2.3.4",
+			trusted:    nil,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "untrusted RemoteAddr ignores XFF",
+			remoteAddr: "203.0.113.50:9999",
+			xff:        "1.2.3.4, 10.0.0.1",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "trusted RemoteAddr with single client in XFF",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "203.0.113.50",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "trusted RemoteAddr walks past trusted entries in XFF",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "203.0.113.50, 10.0.0.2, 172.16.0.5",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "trusted RemoteAddr with empty XFF falls back to RemoteAddr",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("10.0.0.1"),
+		},
+		{
+			name:       "all XFF IPs trusted returns leftmost",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "10.0.0.2, 172.16.0.1, 10.0.0.3",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("10.0.0.2"),
+		},
+		{
+			name:       "XFF with whitespace",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        " 203.0.113.50 , 10.0.0.2 ",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "XFF with empty segments",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "203.0.113.50,,10.0.0.2",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "multi-hop with mixed trust",
+			remoteAddr: "10.0.0.1:5000",
+			xff:        "8.8.8.8, 203.0.113.50, 172.16.0.1",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+		{
+			name:       "RemoteAddr without port",
+			remoteAddr: "10.0.0.1",
+			xff:        "203.0.113.50",
+			trusted:    trusted,
+			want:       netip.MustParseAddr("203.0.113.50"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, ResolveClientIP(tt.remoteAddr, tt.xff, tt.trusted))
+		})
+	}
+}

proxy/lifecycle.go

@@ -2,13 +2,13 @@ package proxy
 
 import (
 	"context"
+	"net/netip"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/embed"
 	"github.com/netbirdio/netbird/proxy/internal/acme"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 // Config bundles every knob the proxy reads at construction time. It mirrors
@@ -83,9 +83,9 @@ type Config struct {
 	// ForwardedProto overrides the X-Forwarded-Proto value sent to
 	// backends. Valid values: "auto", "http", "https".
 	ForwardedProto string
-	// TrustedProxies is the set of trusted upstream proxies that may set
-	// forwarding headers.
-	TrustedProxies *trustedproxy.List
+	// TrustedProxies is a list of IP prefixes for trusted upstream
+	// proxies that may set forwarding headers.
+	TrustedProxies []netip.Prefix
 	// WireguardPort is the UDP port for the embedded NetBird tunnel.
 	// Zero asks the OS for a random port.
 	WireguardPort uint16

proxy/proxyprotocol_test.go

@@ -10,14 +10,12 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 func TestWrapProxyProtocol_OverridesRemoteAddr(t *testing.T) {
 	srv := &Server{
 		Logger:         log.StandardLogger(),
-		TrustedProxies: trustedproxy.FromPrefixes([]netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")}),
+		TrustedProxies: []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")},
 		ProxyProtocol:  true,
 	}
 
@@ -68,7 +66,7 @@ func TestWrapProxyProtocol_OverridesRemoteAddr(t *testing.T) {
 func TestProxyProtocolPolicy_TrustedRequires(t *testing.T) {
 	srv := &Server{
 		Logger:         log.StandardLogger(),
-		TrustedProxies: trustedproxy.FromPrefixes([]netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}),
+		TrustedProxies: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
 	}
 
 	opts := proxyproto.ConnPolicyOptions{
@@ -82,7 +80,7 @@ func TestProxyProtocolPolicy_TrustedRequires(t *testing.T) {
 func TestProxyProtocolPolicy_UntrustedIgnores(t *testing.T) {
 	srv := &Server{
 		Logger:         log.StandardLogger(),
-		TrustedProxies: trustedproxy.FromPrefixes([]netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}),
+		TrustedProxies: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
 	}
 
 	opts := proxyproto.ConnPolicyOptions{
@@ -96,7 +94,7 @@ func TestProxyProtocolPolicy_UntrustedIgnores(t *testing.T) {
 func TestProxyProtocolPolicy_InvalidIPRejects(t *testing.T) {
 	srv := &Server{
 		Logger:         log.StandardLogger(),
-		TrustedProxies: trustedproxy.FromPrefixes([]netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}),
+		TrustedProxies: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
 	}
 
 	opts := proxyproto.ConnPolicyOptions{

proxy/server.go

@@ -67,7 +67,6 @@ import (
 	"github.com/netbirdio/netbird/proxy/web"
 	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/proto"
-	"github.com/netbirdio/netbird/trustedproxy"
 	"github.com/netbirdio/netbird/util/embeddedroots"
 )
 
@@ -80,19 +79,19 @@ type portRouter struct {
 
 type Server struct {
 	ctx               context.Context
-	mgmtClient        proto.ProxyServiceClient
-	proxy             *proxy.ReverseProxy
-	netbird           *roundtrip.NetBird
-	acme              *acme.Manager
+	mgmtClient    proto.ProxyServiceClient
+	proxy         *proxy.ReverseProxy
+	netbird       *roundtrip.NetBird
+	acme          *acme.Manager
 	staticCertWatcher *certwatch.Watcher
-	auth              *auth.Middleware
-	http              *http.Server
-	https             *http.Server
-	debug             *http.Server
-	healthServer      *health.Server
-	healthChecker     *health.Checker
-	meter             *proxymetrics.Metrics
-	accessLog         *accesslog.Logger
+	auth          *auth.Middleware
+	http          *http.Server
+	https         *http.Server
+	debug         *http.Server
+	healthServer  *health.Server
+	healthChecker *health.Checker
+	meter         *proxymetrics.Metrics
+	accessLog     *accesslog.Logger
 	// middlewareManager drives per-target middleware dispatch. Always
 	// constructed during boot; an empty registry produces empty chains and
 	// the reverse-proxy stays on the no-capture fast path.
@@ -100,16 +99,16 @@ type Server struct {
 	// middlewareRegistry is the source of registered middleware factories.
 	// Concrete middlewares register themselves through init().
 	middlewareRegistry *middleware.Registry
-	mainRouter         *nbtcp.Router
-	mainPort           uint16
-	udpMu              sync.Mutex
-	udpRelays          map[types.ServiceID]*udprelay.Relay
-	udpRelayWg         sync.WaitGroup
-	portMu             sync.RWMutex
-	portRouters        map[uint16]*portRouter
-	svcPorts           map[types.ServiceID][]uint16
-	lastMappings       map[types.ServiceID]*proto.ProxyMapping
-	portRouterWg       sync.WaitGroup
+	mainRouter    *nbtcp.Router
+	mainPort      uint16
+	udpMu         sync.Mutex
+	udpRelays     map[types.ServiceID]*udprelay.Relay
+	udpRelayWg    sync.WaitGroup
+	portMu        sync.RWMutex
+	portRouters   map[uint16]*portRouter
+	svcPorts      map[types.ServiceID][]uint16
+	lastMappings  map[types.ServiceID]*proto.ProxyMapping
+	portRouterWg  sync.WaitGroup
 
 	// hijackTracker tracks hijacked connections (e.g. WebSocket upgrades)
 	// so they can be closed during graceful shutdown, since http.Server.Shutdown
@@ -193,10 +192,10 @@ type Server struct {
 	// ForwardedProto overrides the X-Forwarded-Proto value sent to backends.
 	// Valid values: "auto" (detect from TLS), "http", "https".
 	ForwardedProto string
-	// TrustedProxies is the set of trusted upstream proxies. When set,
-	// forwarding headers from these sources are preserved and appended to
-	// instead of being stripped.
-	TrustedProxies *trustedproxy.List
+	// TrustedProxies is a list of IP prefixes for trusted upstream proxies.
+	// When set, forwarding headers from these sources are preserved and
+	// appended to instead of being stripped.
+	TrustedProxies []netip.Prefix
 	// WireguardPort is the port for the NetBird tunnel interface. Use 0
 	// for a random OS-assigned port. A fixed port only works with
 	// single-account deployments; multiple accounts will fail to bind
@@ -719,7 +718,7 @@ func (s *Server) wrapProxyProtocol(ln net.Listener) net.Listener {
 		Listener:          ln,
 		ReadHeaderTimeout: proxyProtoHeaderTimeout,
 	}
-	if !s.TrustedProxies.Empty() {
+	if len(s.TrustedProxies) > 0 {
 		ppListener.ConnPolicy = s.proxyProtocolPolicy
 	} else {
 		s.Logger.Warn("PROXY protocol enabled without trusted proxies; any source may send PROXY headers")
@@ -743,8 +742,10 @@ func (s *Server) proxyProtocolPolicy(opts proxyproto.ConnPolicyOptions) (proxypr
 	addr = addr.Unmap()
 
 	// called per accept
-	if s.TrustedProxies.Contains(addr) {
-		return proxyproto.REQUIRE, nil
+	for _, prefix := range s.TrustedProxies {
+		if prefix.Contains(addr) {
+			return proxyproto.REQUIRE, nil
+		}
 	}
 	return proxyproto.IGNORE, nil
 }

proxy/trustedproxy.go

@@ -0,0 +1,43 @@
+package proxy
+
+import (
+	"fmt"
+	"net/netip"
+	"strings"
+)
+
+// ParseTrustedProxies parses a comma-separated list of CIDR prefixes or bare IPs
+// into a slice of netip.Prefix values suitable for trusted proxy configuration.
+// Bare IPs are converted to single-host prefixes (/32 or /128).
+func ParseTrustedProxies(raw string) ([]netip.Prefix, error) {
+	if raw == "" {
+		return nil, nil
+	}
+
+	parts := strings.Split(raw, ",")
+	prefixes := make([]netip.Prefix, 0, len(parts))
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+
+		prefix, err := netip.ParsePrefix(part)
+		if err == nil {
+			prefixes = append(prefixes, prefix)
+			continue
+		}
+
+		addr, addrErr := netip.ParseAddr(part)
+		if addrErr != nil {
+			return nil, fmt.Errorf("parse trusted proxy %q: not a valid CIDR or IP: %w", part, addrErr)
+		}
+
+		bits := 32
+		if addr.Is6() {
+			bits = 128
+		}
+		prefixes = append(prefixes, netip.PrefixFrom(addr, bits))
+	}
+	return prefixes, nil
+}

proxy/trustedproxy_test.go

@@ -0,0 +1,90 @@
+package proxy
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseTrustedProxies(t *testing.T) {
+	tests := []struct {
+		name    string
+		raw     string
+		want    []netip.Prefix
+		wantErr bool
+	}{
+		{
+			name: "empty string returns nil",
+			raw:  "",
+			want: nil,
+		},
+		{
+			name: "single CIDR",
+			raw:  "10.0.0.0/8",
+			want: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
+		},
+		{
+			name: "single bare IPv4",
+			raw:  "1.2.3.4",
+			want: []netip.Prefix{netip.MustParsePrefix("1.2.3.4/32")},
+		},
+		{
+			name: "single bare IPv6",
+			raw:  "::1",
+			want: []netip.Prefix{netip.MustParsePrefix("::1/128")},
+		},
+		{
+			name: "comma-separated CIDRs",
+			raw:  "10.0.0.0/8, 192.168.1.0/24",
+			want: []netip.Prefix{
+				netip.MustParsePrefix("10.0.0.0/8"),
+				netip.MustParsePrefix("192.168.1.0/24"),
+			},
+		},
+		{
+			name: "mixed CIDRs and bare IPs",
+			raw:  "10.0.0.0/8, 1.2.3.4, fd00::/8",
+			want: []netip.Prefix{
+				netip.MustParsePrefix("10.0.0.0/8"),
+				netip.MustParsePrefix("1.2.3.4/32"),
+				netip.MustParsePrefix("fd00::/8"),
+			},
+		},
+		{
+			name: "whitespace around entries",
+			raw:  "  10.0.0.0/8 , 192.168.0.0/16  ",
+			want: []netip.Prefix{
+				netip.MustParsePrefix("10.0.0.0/8"),
+				netip.MustParsePrefix("192.168.0.0/16"),
+			},
+		},
+		{
+			name: "trailing comma produces no extra entry",
+			raw:  "10.0.0.0/8,",
+			want: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
+		},
+		{
+			name:    "invalid entry",
+			raw:     "not-an-ip",
+			wantErr: true,
+		},
+		{
+			name:    "partially invalid",
+			raw:     "10.0.0.0/8, garbage",
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := ParseTrustedProxies(tt.raw)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

relay/cmd/root.go

@@ -24,7 +24,6 @@ import (
 	"github.com/netbirdio/netbird/shared/metrics"
 	"github.com/netbirdio/netbird/shared/relay/auth"
 	"github.com/netbirdio/netbird/stun"
-	"github.com/netbirdio/netbird/trustedproxy"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -46,9 +45,6 @@ type Config struct {
 	LogLevel                 string
 	LogFile                  string
 	HealthcheckListenAddress string
-	// TrustedProxies is a comma-separated list of upstream proxy CIDRs/IPs whose
-	// X-Real-Ip/X-Real-Port headers are trusted. Empty means never trust these headers.
-	TrustedProxies string
 	// STUN server configuration
 	EnableSTUN   bool
 	STUNPorts    []int
@@ -120,7 +116,6 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&cobraConfig.LogLevel, "log-level", "info", "log level")
 	rootCmd.PersistentFlags().StringVar(&cobraConfig.LogFile, "log-file", "console", "log file")
 	rootCmd.PersistentFlags().StringVarP(&cobraConfig.HealthcheckListenAddress, "health-listen-address", "H", ":9000", "listen address of healthcheck server")
-	rootCmd.PersistentFlags().StringVar(&cobraConfig.TrustedProxies, "trusted-proxies", "", "comma-separated list of upstream proxy CIDRs or IPs whose X-Real-Ip/X-Real-Port headers are trusted; leave empty to always use the direct connection address")
 	rootCmd.PersistentFlags().BoolVar(&cobraConfig.EnableSTUN, "enable-stun", false, "enable embedded STUN server")
 	rootCmd.PersistentFlags().IntSliceVar(&cobraConfig.STUNPorts, "stun-ports", []int{3478}, "ports for the embedded STUN server (can be specified multiple times or comma-separated)")
 	rootCmd.PersistentFlags().StringVar(&cobraConfig.STUNLogLevel, "stun-log-level", "info", "log level for STUN server (panic, fatal, error, warn, info, debug, trace)")
@@ -160,15 +155,8 @@ func execute(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("setup metrics: %v", err)
 	}
 
-	trustedProxies, err := trustedproxy.Parse(cobraConfig.TrustedProxies)
-	if err != nil {
-		log.Debugf("failed to parse trusted proxies: %s", err)
-		return fmt.Errorf("failed to parse trusted proxies: %s", err)
-	}
-
 	srvListenerCfg := server.ListenerConfig{
-		Address:        cobraConfig.ListenAddress,
-		TrustedProxies: trustedProxies,
+		Address: cobraConfig.ListenAddress,
 	}
 
 	tlsConfig, tlsSupport, err := handleTLSConfig(cobraConfig)
@@ -185,7 +173,7 @@ func execute(cmd *cobra.Command, args []string) error {
 	}
 
 	hashedSecret := sha256.Sum256([]byte(cobraConfig.AuthSecret))
-	authenticator := auth.NewTimedHMACValidator(hashedSecret[:])
+	authenticator := auth.NewTimedHMACValidator(hashedSecret[:], 24*time.Hour)
 
 	cfg := server.Config{
 		Meter:          metricsServer.Meter,

relay/server/handshake.go

@@ -5,8 +5,14 @@ import (
 	"fmt"
 	"time"
 
+	log "github.com/sirupsen/logrus"
+
 	"github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/shared/relay/messages"
+	//nolint:staticcheck
+	"github.com/netbirdio/netbird/shared/relay/messages/address"
+	//nolint:staticcheck
+	authmsg "github.com/netbirdio/netbird/shared/relay/messages/auth"
 )
 
 const (
@@ -17,30 +23,55 @@ const (
 
 type Validator interface {
 	Validate(any) error
+	// Deprecated: Use Validate instead.
+	ValidateHelloMsgType(any) error
 }
 
-// preparedMsg contains the marshalled success response message
+// preparedMsg contains the marshalled success response messages
 type preparedMsg struct {
-	responseAuthMsg []byte
+	responseHelloMsg []byte
+	responseAuthMsg  []byte
 }
 
 func newPreparedMsg(instanceURL string) (*preparedMsg, error) {
+	rhm, err := marshalResponseHelloMsg(instanceURL)
+	if err != nil {
+		return nil, err
+	}
+
 	ram, err := messages.MarshalAuthResponse(instanceURL)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal auth response msg: %w", err)
 	}
 
 	return &preparedMsg{
-		responseAuthMsg: ram,
+		responseHelloMsg: rhm,
+		responseAuthMsg:  ram,
 	}, nil
 }
 
+func marshalResponseHelloMsg(instanceURL string) ([]byte, error) {
+	addr := &address.Address{URL: instanceURL}
+	addrData, err := addr.Marshal()
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal response address: %w", err)
+	}
+
+	//nolint:staticcheck
+	responseMsg, err := messages.MarshalHelloResponse(addrData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal hello response: %w", err)
+	}
+	return responseMsg, nil
+}
+
 type handshake struct {
 	conn        listener.Conn
 	validator   Validator
 	preparedMsg *preparedMsg
 
-	peerID *messages.PeerID
+	handshakeMethodAuth bool
+	peerID              *messages.PeerID
 }
 
 func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, error) {
@@ -62,11 +93,17 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err
 		return nil, fmt.Errorf("determine message type from %s: %w", h.conn.RemoteAddr(), err)
 	}
 
-	if msgType != messages.MsgTypeAuth {
+	var peerID *messages.PeerID
+	switch msgType {
+	//nolint:staticcheck
+	case messages.MsgTypeHello:
+		peerID, err = h.handleHelloMsg(buf)
+	case messages.MsgTypeAuth:
+		h.handshakeMethodAuth = true
+		peerID, err = h.handleAuthMsg(buf)
+	default:
 		return nil, fmt.Errorf("invalid message type %d from %s", msgType, h.conn.RemoteAddr())
 	}
-
-	peerID, err := h.handleAuthMsg(buf)
 	if err != nil {
 		return peerID, err
 	}
@@ -75,17 +112,46 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err
 }
 
 func (h *handshake) handshakeResponse(ctx context.Context) error {
-	if _, err := h.conn.Write(ctx, h.preparedMsg.responseAuthMsg); err != nil {
+	var responseMsg []byte
+	if h.handshakeMethodAuth {
+		responseMsg = h.preparedMsg.responseAuthMsg
+	} else {
+		responseMsg = h.preparedMsg.responseHelloMsg
+	}
+
+	if _, err := h.conn.Write(ctx, responseMsg); err != nil {
 		return fmt.Errorf("handshake response write to %s (%s): %w", h.peerID, h.conn.RemoteAddr(), err)
 	}
 
 	return nil
 }
 
+func (h *handshake) handleHelloMsg(buf []byte) (*messages.PeerID, error) {
+	//nolint:staticcheck
+	peerID, authData, err := messages.UnmarshalHelloMsg(buf)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal hello message: %w", err)
+	}
+
+	log.Warnf("peer %s (%s) is using deprecated initial message type", peerID, h.conn.RemoteAddr())
+
+	authMsg, err := authmsg.UnmarshalMsg(authData)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal auth message: %w", err)
+	}
+
+	//nolint:staticcheck
+	if err := h.validator.ValidateHelloMsgType(authMsg.AdditionalData); err != nil {
+		return nil, fmt.Errorf("validate %s (%s): %w", peerID, h.conn.RemoteAddr(), err)
+	}
+
+	return peerID, nil
+}
+
 func (h *handshake) handleAuthMsg(buf []byte) (*messages.PeerID, error) {
 	rawPeerID, authPayload, err := messages.UnmarshalAuthMsg(buf)
 	if err != nil {
-		return nil, fmt.Errorf("unmarshal auth message: %w", err)
+		return nil, fmt.Errorf("unmarshal hello message: %w", err)
 	}
 
 	if err := h.validator.Validate(authPayload); err != nil {

relay/server/listener/quic/listener.go

@@ -51,10 +51,7 @@ func (l *Listener) Listen(acceptFn func(conn relaylistener.Conn)) error {
 
 		log.Infof("QUIC client connected from: %s", session.RemoteAddr())
 		conn := NewConn(session)
-		// Run the accept handler (which performs the pre-auth handshake) in its
-		// own goroutine so a slow or stalled handshake cannot block accepting
-		// further connections.
-		go acceptFn(conn)
+		acceptFn(conn)
 	}
 }
 

relay/server/listener/ws/listener.go

@@ -15,7 +15,6 @@ import (
 	"github.com/netbirdio/netbird/relay/protocol"
 	relaylistener "github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/shared/relay"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 const (
@@ -28,9 +27,6 @@ type Listener struct {
 	Address string
 	// TLSConfig is the TLS configuration for the server.
 	TLSConfig *tls.Config
-	// TrustedProxies is the set of upstream proxies whose X-Real-Ip/X-Real-Port
-	// headers are trusted. Headers from any other immediate peer are ignored.
-	TrustedProxies *trustedproxy.List
 
 	server   *http.Server
 	acceptFn func(conn relaylistener.Conn)
@@ -79,7 +75,7 @@ func (l *Listener) Shutdown(ctx context.Context) error {
 }
 
 func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
-	connRemoteAddr := remoteAddr(r, l.TrustedProxies)
+	connRemoteAddr := remoteAddr(r)
 
 	acceptOptions := &websocket.AcceptOptions{
 		OriginPatterns: []string{"*"},
@@ -106,17 +102,9 @@ func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
 	l.acceptFn(conn)
 }
 
-func remoteAddr(r *http.Request, trustedProxies *trustedproxy.List) string {
-	realIP := r.Header.Get("X-Real-Ip")
-	realPort := r.Header.Get("X-Real-Port")
-	if realIP == "" || realPort == "" {
+func remoteAddr(r *http.Request) string {
+	if r.Header.Get("X-Real-Ip") == "" || r.Header.Get("X-Real-Port") == "" {
 		return r.RemoteAddr
 	}
-
-	if !trustedProxies.IsTrusted(r.RemoteAddr) {
-		log.Debugf("ignoring X-Real-Ip header from untrusted peer %s", r.RemoteAddr)
-		return r.RemoteAddr
-	}
-
-	return net.JoinHostPort(realIP, realPort)
+	return net.JoinHostPort(r.Header.Get("X-Real-Ip"), r.Header.Get("X-Real-Port"))
 }

relay/server/server.go

@@ -15,17 +15,14 @@ import (
 	"github.com/netbirdio/netbird/relay/server/listener/quic"
 	"github.com/netbirdio/netbird/relay/server/listener/ws"
 	quictls "github.com/netbirdio/netbird/shared/relay/tls"
-	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 // ListenerConfig is the configuration for the listener.
 // Address: the address to bind the listener to. It could be an address behind a reverse proxy.
 // TLSConfig: the TLS configuration for the listener.
-// TrustedProxies: upstream proxy prefixes whose forwarding headers (X-Real-Ip/X-Real-Port) are trusted.
 type ListenerConfig struct {
-	Address        string
-	TLSConfig      *tls.Config
-	TrustedProxies *trustedproxy.List
+	Address   string
+	TLSConfig *tls.Config
 }
 
 // Server is the main entry point for the relay server.
@@ -65,9 +62,8 @@ func NewServer(config Config) (*Server, error) {
 // Listen starts the relay server.
 func (r *Server) Listen(cfg ListenerConfig) error {
 	wSListener := &ws.Listener{
-		Address:        cfg.Address,
-		TLSConfig:      cfg.TLSConfig,
-		TrustedProxies: cfg.TrustedProxies,
+		Address:   cfg.Address,
+		TLSConfig: cfg.TLSConfig,
 	}
 
 	r.listenerMux.Lock()

shared/management/http/api/openapi.yml

@@ -5164,10 +5164,6 @@ components:
           type: boolean
           description: Whether upstream TLS certificate verification is skipped when the proxy dials this provider's URL. Intended for self-hosted / internal gateways behind a private or self-signed certificate.
           example: false
-        metadata_disabled:
-          type: boolean
-          description: Whether identity metadata injection is disabled for this provider. When enabled (the default), the proxy stamps the caller's user and authorizing group onto upstream requests as provider-specific metadata (e.g. AWS Bedrock's X-Amzn-Bedrock-Request-Metadata header). Set true to suppress it.
-          example: false
         created_at:
           type: string
           format: date-time
@@ -5188,7 +5184,6 @@ components:
         - models
         - enabled
         - skip_tls_verification
-        - metadata_disabled
         - created_at
         - updated_at
     AgentNetworkProviderRequest:
@@ -5245,10 +5240,6 @@ components:
           type: boolean
           description: Skip upstream TLS certificate verification when the proxy dials this provider's URL. For self-hosted / internal gateways behind a private or self-signed certificate. Defaults to false. When omitted on update, the stored value is left unchanged.
           example: false
-        metadata_disabled:
-          type: boolean
-          description: Disable identity metadata injection (the caller's user + authorizing group) for this provider. Defaults to false (metadata is injected). When omitted on update, the stored value is left unchanged.
-          example: false
       required:
         - provider_id
         - name

shared/management/http/api/types.gen.go

@@ -2227,9 +2227,6 @@ type AgentNetworkProvider struct {
 	// IdentityHeaderUserId Wire header name the proxy stamps with the caller's display identity (user email or peer name) when the catalog entry's HeaderPair is `customizable`. Empty disables stamping for this dimension. Ignored when the catalog entry has a fixed HeaderPair (e.g. LiteLLM, Portkey). Used today by Bifrost: typical values are `x-bf-lh-netbird_user_id` (always-on log metadata) or `x-bf-dim-netbird_user_id` (Prometheus / OTEL — requires the label to be pre-declared in the gateway's `client.prometheus_labels` config).
 	IdentityHeaderUserId *string `json:"identity_header_user_id,omitempty"`
 
-	// MetadataDisabled Whether identity metadata injection is disabled for this provider. When enabled (the default), the proxy stamps the caller's user and authorizing group onto upstream requests as provider-specific metadata (e.g. AWS Bedrock's X-Amzn-Bedrock-Request-Metadata header). Set true to suppress it.
-	MetadataDisabled bool `json:"metadata_disabled"`
-
 	// Models Models exposed through this endpoint, with the operator's per-1k input/output prices. Empty means all catalog models are allowed at catalog prices.
 	Models []AgentNetworkProviderModel `json:"models"`
 
@@ -2281,9 +2278,6 @@ type AgentNetworkProviderRequest struct {
 	// IdentityHeaderUserId Wire header name for the caller's display identity. See AgentNetworkProvider.identity_header_user_id. When omitted on a request, the stored value is left unchanged; pass an empty string explicitly to clear it (which disables stamping for this dimension).
 	IdentityHeaderUserId *string `json:"identity_header_user_id,omitempty"`
 
-	// MetadataDisabled Disable identity metadata injection (the caller's user + authorizing group) for this provider. Defaults to false (metadata is injected). When omitted on update, the stored value is left unchanged.
-	MetadataDisabled *bool `json:"metadata_disabled,omitempty"`
-
 	// Models Models exposed through this endpoint, with the operator's per-1k input/output prices. Empty means all catalog models are allowed at catalog prices.
 	Models *[]AgentNetworkProviderModel `json:"models,omitempty"`
 

shared/relay/auth/allow/allow_all.go

@@ -8,3 +8,7 @@ type Auth struct {
 func (a *Auth) Validate(any) error {
 	return nil
 }
+
+func (a *Auth) ValidateHelloMsgType(any) error {
+	return nil
+}

shared/relay/auth/hmac/token.go

@@ -1,8 +1,10 @@
 package hmac
 
 import (
+	"bytes"
 	"crypto/hmac"
 	"encoding/base64"
+	"encoding/gob"
 	"fmt"
 	"hash"
 	"strconv"
@@ -16,6 +18,14 @@ type Token struct {
 	Signature string
 }
 
+func unmarshalToken(payload []byte) (Token, error) {
+	var creds Token
+	buffer := bytes.NewBuffer(payload)
+	decoder := gob.NewDecoder(buffer)
+	err := decoder.Decode(&creds)
+	return creds, err
+}
+
 // TimedHMAC generates a token with TTL and uses a pre-shared secret known to the relay server
 type TimedHMAC struct {
 	secret     string

shared/relay/auth/hmac/validator.go

@@ -0,0 +1,33 @@
+package hmac
+
+import (
+	"crypto/sha256"
+	"fmt"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type TimedHMACValidator struct {
+	*TimedHMAC
+}
+
+func NewTimedHMACValidator(secret string, duration time.Duration) *TimedHMACValidator {
+	ta := NewTimedHMAC(secret, duration)
+	return &TimedHMACValidator{
+		ta,
+	}
+}
+
+func (a *TimedHMACValidator) Validate(credentials any) error {
+	b, ok := credentials.([]byte)
+	if !ok {
+		return fmt.Errorf("invalid credentials type")
+	}
+	c, err := unmarshalToken(b)
+	if err != nil {
+		log.Debugf("failed to unmarshal token: %s", err)
+		return err
+	}
+	return a.TimedHMAC.Validate(sha256.New, c)
+}

shared/relay/auth/validator.go

@@ -1,19 +1,28 @@
 package auth
 
 import (
+	"time"
+
+	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
 	authv2 "github.com/netbirdio/netbird/shared/relay/auth/hmac/v2"
 )
 
 type TimedHMACValidator struct {
 	authenticatorV2 *authv2.Validator
+	authenticator   *auth.TimedHMACValidator
 }
 
-func NewTimedHMACValidator(secret []byte) *TimedHMACValidator {
+func NewTimedHMACValidator(secret []byte, duration time.Duration) *TimedHMACValidator {
 	return &TimedHMACValidator{
 		authenticatorV2: authv2.NewValidator(secret),
+		authenticator:   auth.NewTimedHMACValidator(string(secret), duration),
 	}
 }
 
 func (a *TimedHMACValidator) Validate(credentials any) error {
 	return a.authenticatorV2.Validate(credentials)
 }
+
+func (a *TimedHMACValidator) ValidateHelloMsgType(credentials any) error {
+	return a.authenticator.Validate(credentials)
+}

shared/relay/client/fallback_opener.go

@@ -1,165 +0,0 @@
-package client
-
-import (
-	"context"
-	"errors"
-	"net"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	raceTotalTimeout  = 40 * time.Second
-	raceFallbackDelay = 10 * time.Second
-)
-
-type raceAttempt struct {
-	conn net.Conn
-	err  error
-}
-
-type raceOutcome struct {
-	conn net.Conn
-	err  error
-	done bool
-}
-
-type connRace struct {
-	opener            *FallbackOpener
-	peerKey           string
-	remoteRelayServer RelayServer
-	preferForeign     bool
-
-	raceCtx         context.Context
-	otherCtx        context.Context
-	cancelPreferred context.CancelFunc
-	cancelOther     context.CancelFunc
-	results         chan raceAttempt
-	fallbackTimer   *time.Timer
-
-	otherStarted bool
-	settled      int
-	lastErr      error
-}
-
-type FallbackOpener struct {
-	home         *Client
-	foreignStore *ForeignRelaysStore
-}
-
-func NewFallbackOpener(home *Client, foreignStore *ForeignRelaysStore) *FallbackOpener {
-	return &FallbackOpener{
-		home:         home,
-		foreignStore: foreignStore,
-	}
-}
-
-func (r *FallbackOpener) Run(ctx context.Context, peerKey string, remoteRelayServer RelayServer, preferForeign bool) (net.Conn, error) {
-	raceCtx, cancel := context.WithTimeout(ctx, raceTotalTimeout)
-	defer cancel()
-
-	preferredCtx, cancelPreferred := context.WithCancel(raceCtx)
-	otherCtx, cancelOther := context.WithCancel(raceCtx)
-
-	race := &connRace{
-		opener:            r,
-		peerKey:           peerKey,
-		remoteRelayServer: remoteRelayServer,
-		preferForeign:     preferForeign,
-		raceCtx:           raceCtx,
-		otherCtx:          otherCtx,
-		cancelPreferred:   cancelPreferred,
-		cancelOther:       cancelOther,
-		results:           make(chan raceAttempt, 2),
-		fallbackTimer:     time.NewTimer(raceFallbackDelay),
-	}
-	defer race.fallbackTimer.Stop()
-
-	go func() {
-		race.results <- r.open(preferredCtx, peerKey, remoteRelayServer, preferForeign)
-	}()
-
-	for {
-		select {
-		case <-race.fallbackTimer.C:
-			race.startOther()
-		case res := <-race.results:
-			if o := race.handleResult(res); o.done {
-				return o.conn, o.err
-			}
-		case <-raceCtx.Done():
-			return race.onTimeout()
-		}
-	}
-}
-
-func (c *connRace) startOther() {
-	if c.otherStarted {
-		return
-	}
-	c.otherStarted = true
-	c.fallbackTimer.Stop()
-	go func() {
-		c.results <- c.opener.open(c.otherCtx, c.peerKey, c.remoteRelayServer, !c.preferForeign)
-	}()
-}
-
-func (c *connRace) handleResult(res raceAttempt) raceOutcome {
-	if (res.err == nil && res.conn != nil) || errors.Is(res.err, ErrConnAlreadyExists) {
-		c.settled++
-		c.stop()
-		return raceOutcome{conn: res.conn, err: res.err, done: true}
-	}
-
-	c.lastErr = res.err
-	c.settled++
-	if !c.otherStarted {
-		c.startOther()
-		return raceOutcome{}
-	}
-	if c.settled == 2 {
-		c.cancelPreferred()
-		c.cancelOther()
-		return raceOutcome{err: c.lastErr, done: true}
-	}
-	return raceOutcome{}
-}
-
-func (c *connRace) onTimeout() (net.Conn, error) {
-	c.stop()
-	if c.lastErr != nil {
-		return nil, c.lastErr
-	}
-	return nil, c.raceCtx.Err()
-}
-
-func (c *connRace) stop() {
-	c.cancelPreferred()
-	c.cancelOther()
-	go c.opener.drainLoser(c.results, c.settled, c.otherStarted)
-}
-
-func (r *FallbackOpener) open(ctx context.Context, peerKey string, remoteRelayServer RelayServer, foreign bool) raceAttempt {
-	if foreign {
-		conn, err := r.foreignStore.OpenConn(ctx, peerKey, remoteRelayServer)
-		return raceAttempt{conn: conn, err: err}
-	}
-	conn, err := r.home.OpenConn(ctx, peerKey)
-	return raceAttempt{conn: conn, err: err}
-}
-
-func (r *FallbackOpener) drainLoser(results chan raceAttempt, settled int, otherStarted bool) {
-	started := 1
-	if otherStarted {
-		started = 2
-	}
-	for i := settled; i < started; i++ {
-		res := <-results
-		if res.conn != nil {
-			if err := res.conn.Close(); err != nil {
-				log.Debugf("failed to close losing relay connection: %v", err)
-			}
-		}
-	}
-}

shared/relay/client/fallback_opener_test.go

@@ -1,175 +0,0 @@
-package client
-
-import (
-	"context"
-	"errors"
-	"net"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/relay/server"
-)
-
-type fakeConn struct {
-	net.Conn
-	closed chan struct{}
-}
-
-func newFakeConn() *fakeConn {
-	return &fakeConn{closed: make(chan struct{})}
-}
-
-func (c *fakeConn) Close() error {
-	close(c.closed)
-	return nil
-}
-
-func newTestConnRace(t *testing.T) *connRace {
-	t.Helper()
-	raceCtx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-	_, cancelPreferred := context.WithCancel(raceCtx)
-	otherCtx, cancelOther := context.WithCancel(raceCtx)
-	timer := time.NewTimer(time.Hour)
-	timer.Stop()
-	return &connRace{
-		opener:          &FallbackOpener{},
-		peerKey:         "peerKey",
-		raceCtx:         raceCtx,
-		otherCtx:        otherCtx,
-		cancelPreferred: cancelPreferred,
-		cancelOther:     cancelOther,
-		results:         make(chan raceAttempt, 2),
-		fallbackTimer:   timer,
-	}
-}
-
-func TestHandleResult_PreferredSucceeds(t *testing.T) {
-	c := newTestConnRace(t)
-
-	conn := newFakeConn()
-	o := c.handleResult(raceAttempt{conn: conn})
-
-	require.True(t, o.done)
-	require.NoError(t, o.err)
-	require.Same(t, net.Conn(conn), o.conn)
-	require.False(t, c.otherStarted, "fallback must not start once the preferred attempt wins")
-}
-
-func TestHandleResult_ConnAlreadyExistsIsSuccess(t *testing.T) {
-	c := newTestConnRace(t)
-
-	o := c.handleResult(raceAttempt{err: ErrConnAlreadyExists})
-
-	require.True(t, o.done)
-	require.ErrorIs(t, o.err, ErrConnAlreadyExists)
-	require.False(t, c.otherStarted)
-}
-
-func TestHandleResult_PreferredFailsStartsOther(t *testing.T) {
-	c := newTestConnRace(t)
-	// The fallback attempt opens against a stalling listener so startOther's
-	// goroutine blocks on Connect until raceCtx is cancelled by t.Cleanup.
-	serverAddr, _ := stallingRelayListener(t)
-	c.opener.foreignStore = NewForeignRelaysStore(c.raceCtx, hmacTokenStore, "alice", 1280, newTransportFallback(), func(string) {}, keepUnusedServerTime)
-	c.remoteRelayServer = RelayServer{Addr: serverAddr}
-	c.preferForeign = false
-
-	o := c.handleResult(raceAttempt{err: errors.New("boom")})
-
-	require.False(t, o.done, "a single failure must not settle the race")
-	require.True(t, c.otherStarted, "the fallback attempt must start after the preferred one fails")
-	require.EqualError(t, c.lastErr, "boom")
-}
-
-func TestHandleResult_BothFailReturnsLastErr(t *testing.T) {
-	c := newTestConnRace(t)
-	c.otherStarted = true
-	c.settled = 1
-	c.lastErr = errors.New("first")
-
-	o := c.handleResult(raceAttempt{err: errors.New("second")})
-
-	require.True(t, o.done)
-	require.EqualError(t, o.err, "second")
-}
-
-func TestOnTimeout_PrefersLastErr(t *testing.T) {
-	c := newTestConnRace(t)
-	c.lastErr = errors.New("dial failed")
-
-	_, err := c.onTimeout()
-	require.EqualError(t, err, "dial failed")
-}
-
-func TestOnTimeout_FallsBackToCtxErr(t *testing.T) {
-	c := newTestConnRace(t)
-	raceCtx, cancel := context.WithCancel(context.Background())
-	cancel()
-	c.raceCtx = raceCtx
-
-	_, err := c.onTimeout()
-	require.ErrorIs(t, err, context.Canceled)
-}
-
-func TestDrainLoser_ClosesLateWinner(t *testing.T) {
-	r := &FallbackOpener{}
-	results := make(chan raceAttempt, 2)
-
-	loser := newFakeConn()
-	results <- raceAttempt{conn: loser}
-
-	done := make(chan struct{})
-	go func() {
-		defer close(done)
-		r.drainLoser(results, 1, true)
-	}()
-
-	select {
-	case <-loser.closed:
-	case <-time.After(2 * time.Second):
-		t.Fatal("losing connection was not closed")
-	}
-	select {
-	case <-done:
-	case <-time.After(2 * time.Second):
-		t.Fatal("drainLoser did not return")
-	}
-}
-
-func TestDrainLoser_NoOtherAttempt(t *testing.T) {
-	r := &FallbackOpener{}
-	results := make(chan raceAttempt)
-
-	done := make(chan struct{})
-	go func() {
-		defer close(done)
-		r.drainLoser(results, 1, false)
-	}()
-
-	select {
-	case <-done:
-	case <-time.After(2 * time.Second):
-		t.Fatal("drainLoser blocked with no outstanding attempt")
-	}
-}
-
-func startTestRelayServer(t *testing.T, addr string) string {
-	t.Helper()
-
-	srv, err := server.NewServer(newManagerTestServerConfig(addr))
-	require.NoError(t, err)
-
-	errChan := make(chan error, 1)
-	go func() {
-		if err := srv.Listen(server.ListenerConfig{Address: addr}); err != nil {
-			errChan <- err
-		}
-	}()
-	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
-
-	require.NoError(t, waitForServerToStart(errChan))
-	return addr
-}

shared/relay/client/foreign_relays_store.go

@@ -1,155 +0,0 @@
-package client
-
-import (
-	"context"
-	"net"
-	"sync"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-	"golang.org/x/sync/singleflight"
-
-	relayAuth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
-)
-
-type foreignRelay struct {
-	client  *Client
-	created time.Time
-	inUse   int
-}
-
-type ForeignRelaysStore struct {
-	mu      sync.RWMutex
-	clients map[string]*foreignRelay
-
-	group singleflight.Group
-
-	ctx                  context.Context
-	tokenStore           *relayAuth.TokenStore
-	peerID               string
-	mtu                  uint16
-	transportFallback    *transportFallback
-	onDisconnect         func(string)
-	keepUnusedServerTime time.Duration
-}
-
-func NewForeignRelaysStore(ctx context.Context, tokenStore *relayAuth.TokenStore, peerID string, mtu uint16, transportFallback *transportFallback, onDisconnect func(string), keepUnusedServerTime time.Duration) *ForeignRelaysStore {
-	return &ForeignRelaysStore{
-		clients:              make(map[string]*foreignRelay),
-		ctx:                  ctx,
-		tokenStore:           tokenStore,
-		peerID:               peerID,
-		mtu:                  mtu,
-		transportFallback:    transportFallback,
-		onDisconnect:         onDisconnect,
-		keepUnusedServerTime: keepUnusedServerTime,
-	}
-}
-
-func (f *ForeignRelaysStore) OpenConn(ctx context.Context, peerKey string, remoteRelayServer RelayServer) (net.Conn, error) {
-	fr, err := f.acquire(remoteRelayServer)
-	if err != nil {
-		return nil, err
-	}
-	defer f.release(fr)
-
-	return fr.client.OpenConn(ctx, peerKey)
-}
-
-func (f *ForeignRelaysStore) acquire(remoteRelayServer RelayServer) (*foreignRelay, error) {
-	f.mu.Lock()
-	if fr, ok := f.clients[remoteRelayServer.Addr]; ok {
-		fr.inUse++
-		f.mu.Unlock()
-		return fr, nil
-	}
-	f.mu.Unlock()
-
-	v, err, _ := f.group.Do(remoteRelayServer.Addr, func() (any, error) {
-		f.mu.RLock()
-		fr, ok := f.clients[remoteRelayServer.Addr]
-		f.mu.RUnlock()
-		if ok {
-			return fr, nil
-		}
-
-		relayClient := NewClientWithServerIP(remoteRelayServer.Addr, remoteRelayServer.IP, f.tokenStore, f.peerID, f.mtu)
-		relayClient.SetTransportFallback(f.transportFallback)
-		if err := relayClient.Connect(f.ctx); err != nil {
-			return nil, err
-		}
-		relayClient.SetOnDisconnectListener(f.onDisconnect)
-
-		f.mu.Lock()
-		fr = &foreignRelay{client: relayClient, created: time.Now()}
-		f.clients[remoteRelayServer.Addr] = fr
-		f.mu.Unlock()
-		return fr, nil
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	fr := v.(*foreignRelay)
-	f.mu.Lock()
-	if cur, ok := f.clients[remoteRelayServer.Addr]; !ok || cur != fr {
-		f.mu.Unlock()
-		return f.acquire(remoteRelayServer)
-	}
-	fr.inUse++
-	f.mu.Unlock()
-	return fr, nil
-}
-
-func (f *ForeignRelaysStore) release(fr *foreignRelay) {
-	f.mu.Lock()
-	fr.inUse--
-	f.mu.Unlock()
-}
-
-func (f *ForeignRelaysStore) evict(serverAddress string) {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	if _, ok := f.clients[serverAddress]; ok {
-		delete(f.clients, serverAddress)
-		log.Debugf("evicted disconnected foreign relay client: %s", serverAddress)
-	}
-}
-
-func (f *ForeignRelaysStore) cleanupUnused() {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-
-	for addr, fr := range f.clients {
-		if time.Since(fr.created) <= f.keepUnusedServerTime {
-			continue
-		}
-		if fr.inUse > 0 {
-			continue
-		}
-		if fr.client.HasConns() {
-			continue
-		}
-		fr.client.SetOnDisconnectListener(nil)
-		go func() {
-			_ = fr.client.Close()
-		}()
-		log.Debugf("clean up unused relay server connection: %s", addr)
-		delete(f.clients, addr)
-	}
-}
-
-func (f *ForeignRelaysStore) states() []RelayConnState {
-	f.mu.RLock()
-	clients := make([]*Client, 0, len(f.clients))
-	for _, fr := range f.clients {
-		clients = append(clients, fr.client)
-	}
-	f.mu.RUnlock()
-
-	states := make([]RelayConnState, 0, len(clients))
-	for _, c := range clients {
-		states = append(states, relayConnState(c))
-	}
-	return states
-}

shared/relay/client/foreign_relays_store_test.go

@@ -1,184 +0,0 @@
-package client
-
-import (
-	"context"
-	"sync"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-)
-
-func newTestForeignStore(t *testing.T, ctx context.Context) *ForeignRelaysStore {
-	t.Helper()
-	return NewForeignRelaysStore(ctx, hmacTokenStore, "alice", 1280, newTransportFallback(), func(string) {}, keepUnusedServerTime)
-}
-
-func TestForeignStore_AcquireDedupsConcurrentOpens(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52601")
-	server := RelayServer{Addr: "rel://" + addr}
-
-	store := newTestForeignStore(t, ctx)
-
-	const n = 8
-	var wg sync.WaitGroup
-	results := make([]*foreignRelay, n)
-	for i := 0; i < n; i++ {
-		wg.Add(1)
-		go func(i int) {
-			defer wg.Done()
-			fr, err := store.acquire(server)
-			require.NoError(t, err)
-			results[i] = fr
-		}(i)
-	}
-	wg.Wait()
-
-	first := results[0]
-	require.NotNil(t, first)
-	for _, fr := range results {
-		require.Same(t, first, fr, "all acquires must share the same foreign relay")
-	}
-
-	store.mu.RLock()
-	require.Len(t, store.clients, 1, "only one client entry must be stored")
-	require.Equal(t, n, first.inUse, "every acquire must be counted")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_AcquireReleaseRefcount(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52602")
-	server := RelayServer{Addr: "rel://" + addr}
-
-	store := newTestForeignStore(t, ctx)
-
-	fr, err := store.acquire(server)
-	require.NoError(t, err)
-	_, err = store.acquire(server)
-	require.NoError(t, err)
-
-	store.mu.RLock()
-	require.Equal(t, 2, fr.inUse)
-	store.mu.RUnlock()
-
-	store.release(fr)
-	store.mu.RLock()
-	require.Equal(t, 1, fr.inUse)
-	require.Len(t, store.clients, 1, "release must not evict the client")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_AcquireConnectFailure(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	t.Cleanup(cancel)
-
-	store := newTestForeignStore(t, ctx)
-
-	// Nothing is listening on this port, so Connect fails.
-	_, err := store.acquire(RelayServer{Addr: "rel://127.0.0.1:1"})
-	require.Error(t, err)
-
-	store.mu.RLock()
-	require.Empty(t, store.clients, "a failed connect must not leave a client behind")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_Evict(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	store := newTestForeignStore(t, ctx)
-	store.clients["rel://a"] = &foreignRelay{}
-	store.clients["rel://b"] = &foreignRelay{}
-
-	store.evict("rel://a")
-	store.evict("rel://missing")
-
-	require.NotContains(t, store.clients, "rel://a")
-	require.Contains(t, store.clients, "rel://b")
-}
-
-func TestForeignStore_CleanupUnused_KeepsRecent(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52603")
-	store := newTestForeignStore(t, ctx)
-
-	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
-	require.NoError(t, err)
-	store.release(fr)
-
-	store.cleanupUnused()
-
-	store.mu.RLock()
-	require.Len(t, store.clients, 1, "a freshly created client must be kept")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_CleanupUnused_KeepsInUse(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52604")
-	store := newTestForeignStore(t, ctx)
-
-	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
-	require.NoError(t, err)
-
-	store.mu.Lock()
-	fr.created = time.Now().Add(-2 * keepUnusedServerTime)
-	store.mu.Unlock()
-
-	store.cleanupUnused()
-
-	store.mu.RLock()
-	require.Len(t, store.clients, 1, "an in-use client must be kept even when aged")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_CleanupUnused_EvictsAgedIdle(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52605")
-	store := newTestForeignStore(t, ctx)
-
-	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
-	require.NoError(t, err)
-	store.release(fr)
-
-	store.mu.Lock()
-	fr.created = time.Now().Add(-2 * keepUnusedServerTime)
-	store.mu.Unlock()
-
-	require.False(t, fr.client.HasConns(), "no peer connections were opened")
-
-	store.cleanupUnused()
-
-	store.mu.RLock()
-	require.Empty(t, store.clients, "an aged idle client must be evicted")
-	store.mu.RUnlock()
-}
-
-func TestForeignStore_States(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	addr := startTestRelayServer(t, "127.0.0.1:52606")
-	store := newTestForeignStore(t, ctx)
-
-	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
-	require.NoError(t, err)
-	store.release(fr)
-
-	states := store.states()
-	require.Len(t, states, 1)
-	require.NotEmpty(t, states[0].URL)
-}

shared/relay/client/manager.go

@@ -22,6 +22,27 @@ var (
 	ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
 )
 
+// RelayTrack hold the relay clients for the foreign relay servers.
+// With the mutex can ensure we can open new connection in case the relay connection has been established with
+// the relay server.
+type RelayTrack struct {
+	sync.RWMutex
+	relayClient *Client
+	err         error
+	created     time.Time
+	// ready is closed once the dial started by openConnVia finishes (relayClient
+	// or err is set). Callers reusing a track wait on this instead of the track
+	// lock, so the dial never runs under rt.Lock.
+	ready chan struct{}
+}
+
+func NewRelayTrack() *RelayTrack {
+	return &RelayTrack{
+		created: time.Now(),
+		ready:   make(chan struct{}),
+	}
+}
+
 type OnServerCloseListener func()
 
 // ManagerOption configures a Manager at construction time.
@@ -38,11 +59,6 @@ type RelayConnState struct {
 	Err error
 }
 
-type RelayServer struct {
-	Addr string
-	IP   netip.Addr
-}
-
 // WithMaxBackoffInterval caps the exponential backoff between reconnect
 // attempts to the home relay. A non-positive value keeps the default.
 func WithMaxBackoffInterval(d time.Duration) ManagerOption {
@@ -67,7 +83,8 @@ type Manager struct {
 	relayClientMu  sync.RWMutex
 	reconnectGuard *Guard
 
-	foreign *ForeignRelaysStore
+	relayClients      map[string]*RelayTrack
+	relayClientsMutex sync.RWMutex
 
 	onDisconnectedListeners map[string]*list.List
 	onReconnectedListenerFn func()
@@ -103,6 +120,7 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 			ConnectionTimeout: defaultConnectionTimeout,
 			TransportFallback: tf,
 		},
+		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),
 		cleanupInterval:         relayCleanupInterval,
 		keepUnusedServerTime:    keepUnusedServerTime,
@@ -110,7 +128,6 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 	for _, opt := range opts {
 		opt(m)
 	}
-	m.foreign = NewForeignRelaysStore(ctx, tokenStore, peerID, mtu, tf, m.onServerDisconnected, m.keepUnusedServerTime)
 	m.serverPicker.ServerURLs.Store(serverURLs)
 	m.reconnectGuard = NewGuard(m.serverPicker, m.maxBackoffInterval)
 	return m
@@ -142,26 +159,40 @@ func (m *Manager) Serve() error {
 	return err
 }
 
-func (m *Manager) OpenConn(ctx context.Context, remoteRelayServer RelayServer, peerKey string, preferForeign bool) (net.Conn, error) {
+// OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
+// established via the relay server. If the peer is on a different relay server, the manager will establish a new
+// connection to the relay server. It returns back with a net.Conn what represent the remote peer connection.
+//
+// serverIP, when valid and serverAddress is foreign, is used as a dial target if the FQDN-based dial fails.
+// Ignored for the local home-server path. TLS verification still uses the FQDN via SNI.
+func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) {
 	m.relayClientMu.RLock()
-	relayClient := m.relayClient
-	m.relayClientMu.RUnlock()
+	defer m.relayClientMu.RUnlock()
 
-	if relayClient == nil {
+	if m.relayClient == nil {
 		return nil, ErrRelayClientNotConnected
 	}
 
-	foreign, err := m.isForeignServer(relayClient, remoteRelayServer.Addr)
+	foreign, err := m.isForeignServer(serverAddress)
 	if err != nil {
 		return nil, err
 	}
 
+	var (
+		netConn net.Conn
+	)
 	if !foreign {
-		return relayClient.OpenConn(ctx, peerKey)
+		log.Debugf("open peer connection via permanent server: %s", peerKey)
+		netConn, err = m.relayClient.OpenConn(ctx, peerKey)
+	} else {
+		log.Debugf("open peer connection via foreign server: %s", serverAddress)
+		netConn, err = m.openConnVia(ctx, serverAddress, peerKey, serverIP)
+	}
+	if err != nil {
+		return nil, err
 	}
 
-	opener := NewFallbackOpener(relayClient, m.foreign)
-	return opener.Run(ctx, peerKey, remoteRelayServer, preferForeign)
+	return netConn, err
 }
 
 // Ready returns true if the home Relay client is connected to the relay server.
@@ -192,7 +223,7 @@ func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServ
 		return ErrRelayClientNotConnected
 	}
 
-	foreign, err := m.isForeignServer(m.relayClient, serverAddress)
+	foreign, err := m.isForeignServer(serverAddress)
 	if err != nil {
 		return err
 	}
@@ -256,7 +287,26 @@ func (m *Manager) RelayStates() []RelayConnState {
 		states = append(states, st)
 	}
 
-	states = append(states, m.foreign.states()...)
+	// Snapshot the tracks, then query each outside the map lock: a track can be
+	// held by an in-progress Connect, and blocking on it must not stall other
+	// relay operations.
+	m.relayClientsMutex.RLock()
+	tracks := make([]*RelayTrack, 0, len(m.relayClients))
+	for _, rt := range m.relayClients {
+		tracks = append(tracks, rt)
+	}
+	m.relayClientsMutex.RUnlock()
+
+	// Only connected foreign relays carry state; a failed connect is evicted
+	// immediately (openConnVia), so there is no error state to surface.
+	for _, rt := range tracks {
+		rt.RLock()
+		rc := rt.relayClient
+		rt.RUnlock()
+		if rc != nil {
+			states = append(states, relayConnState(rc))
+		}
+	}
 
 	return states
 }
@@ -277,6 +327,76 @@ func (m *Manager) UpdateToken(token *relayAuth.Token) error {
 	return m.tokenStore.UpdateToken(token)
 }
 
+func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) {
+	// check if already has a connection to the desired relay server
+	m.relayClientsMutex.RLock()
+	rt, ok := m.relayClients[serverAddress]
+	m.relayClientsMutex.RUnlock()
+	if ok {
+		return m.openConnOnTrack(ctx, rt, peerKey)
+	}
+
+	// if not, establish a new connection but check it again (because changed the lock type) before starting the
+	// connection
+	m.relayClientsMutex.Lock()
+	rt, ok = m.relayClients[serverAddress]
+	if ok {
+		m.relayClientsMutex.Unlock()
+		return m.openConnOnTrack(ctx, rt, peerKey)
+	}
+
+	// Publish the track and release the map lock BEFORE dialing, so the dial does
+	// not run under rt.Lock (which would block RelayStates and the cleanup loop
+	// for the full dial). Concurrent callers find this track and wait on rt.ready.
+	rt = NewRelayTrack()
+	m.relayClients[serverAddress] = rt
+	m.relayClientsMutex.Unlock()
+
+	relayClient := NewClientWithServerIP(serverAddress, serverIP, m.tokenStore, m.peerID, m.mtu)
+	relayClient.SetTransportFallback(m.transportFallback)
+	err := relayClient.Connect(m.ctx)
+	if err != nil {
+		rt.Lock()
+		rt.err = err
+		rt.Unlock()
+		close(rt.ready)
+		m.relayClientsMutex.Lock()
+		delete(m.relayClients, serverAddress)
+		m.relayClientsMutex.Unlock()
+		return nil, err
+	}
+	// if connection closed then delete the relay client from the list
+	relayClient.SetOnDisconnectListener(m.onServerDisconnected)
+	rt.Lock()
+	rt.relayClient = relayClient
+	rt.Unlock()
+	close(rt.ready)
+
+	return relayClient.OpenConn(ctx, peerKey)
+}
+
+// openConnOnTrack opens a peer connection through an existing relay track,
+// waiting for the dial started by another openConnVia call to finish. It waits
+// on rt.ready rather than the track lock, so it neither holds nor contends the
+// track lock across the dial.
+func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (net.Conn, error) {
+	select {
+	case <-rt.ready:
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+
+	rt.RLock()
+	defer rt.RUnlock()
+	if rt.err != nil {
+		return nil, rt.err
+	}
+	if rt.relayClient == nil {
+		return nil, ErrRelayClientNotConnected
+	}
+	return rt.relayClient.OpenConn(ctx, peerKey)
+}
+
 func (m *Manager) onServerConnected() {
 	m.listenerLock.Lock()
 	defer m.listenerLock.Unlock()
@@ -302,12 +422,21 @@ func (m *Manager) onServerDisconnected(serverAddress string) {
 	m.relayClientMu.Unlock()
 
 	if !isHome {
-		m.foreign.evict(serverAddress)
+		m.evictForeignRelay(serverAddress)
 	}
 
 	m.notifyOnDisconnectListeners(serverAddress)
 }
 
+func (m *Manager) evictForeignRelay(serverAddress string) {
+	m.relayClientsMutex.Lock()
+	defer m.relayClientsMutex.Unlock()
+	if _, ok := m.relayClients[serverAddress]; ok {
+		delete(m.relayClients, serverAddress)
+		log.Debugf("evicted disconnected foreign relay client: %s", serverAddress)
+	}
+}
+
 func (m *Manager) listenGuardEvent(ctx context.Context) {
 	for {
 		select {
@@ -330,8 +459,8 @@ func (m *Manager) storeClient(client *Client) {
 	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
 }
 
-func (m *Manager) isForeignServer(relayClient *Client, address string) (bool, error) {
-	rAddr, err := relayClient.ServerInstanceURL()
+func (m *Manager) isForeignServer(address string) (bool, error) {
+	rAddr, err := m.relayClient.ServerInstanceURL()
 	if err != nil {
 		return false, fmt.Errorf("relay client not connected")
 	}
@@ -346,11 +475,50 @@ func (m *Manager) startCleanupLoop() {
 		case <-m.ctx.Done():
 			return
 		case <-ticker.C:
-			m.foreign.cleanupUnused()
+			m.cleanUpUnusedRelays()
 		}
 	}
 }
 
+func (m *Manager) cleanUpUnusedRelays() {
+	m.relayClientsMutex.Lock()
+	defer m.relayClientsMutex.Unlock()
+
+	for addr, rt := range m.relayClients {
+		rt.Lock()
+		// if the connection failed to the server the relay client will be nil
+		// but the instance will be kept in the relayClients until the next locking
+		if rt.err != nil {
+			rt.Unlock()
+			continue
+		}
+
+		// dial still in progress (openConnVia publishes the track before Connect
+		// completes and no longer holds rt.Lock during it), nothing to clean up.
+		if rt.relayClient == nil {
+			rt.Unlock()
+			continue
+		}
+
+		if time.Since(rt.created) <= m.keepUnusedServerTime {
+			rt.Unlock()
+			continue
+		}
+
+		if rt.relayClient.HasConns() {
+			rt.Unlock()
+			continue
+		}
+		rt.relayClient.SetOnDisconnectListener(nil)
+		go func() {
+			_ = rt.relayClient.Close()
+		}()
+		log.Debugf("clean up unused relay server connection: %s", addr)
+		delete(m.relayClients, addr)
+		rt.Unlock()
+	}
+}
+
 func (m *Manager) addListener(serverAddress string, onClosedListener OnServerCloseListener) {
 	m.listenerLock.Lock()
 	defer m.listenerLock.Unlock()

shared/relay/client/manager_cleanup_test.go

@@ -2,14 +2,17 @@ package client
 
 import (
 	"context"
+	"net/netip"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial drives a real, hanging foreign
-// relay dial and asserts the foreign store cleanup does not stall behind it.
+// relay dial and asserts cleanUpUnusedRelays does not stall behind it.
 func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
-	serverAddr, accepted := stallingRelayListener(t)
+	serverAddr := stallingRelayListener(t)
 
 	mCtx, mCancel := context.WithCancel(context.Background())
 	t.Cleanup(mCancel)
@@ -19,32 +22,39 @@ func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	dialDone := make(chan struct{})
 	go func() {
 		defer close(dialDone)
-		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
+		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
 	}()
 
-	select {
-	case <-accepted:
-	case <-time.After(5 * time.Second):
-		t.Fatal("relay dial did not reach the listener")
-	}
+	// The track appears in the map once the dial is in flight.
+	require.Eventually(t, func() bool {
+		m.relayClientsMutex.RLock()
+		defer m.relayClientsMutex.RUnlock()
+		_, ok := m.relayClients[serverAddr]
+		return ok
+	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
 
 	cleanupDone := make(chan struct{})
 	go func() {
 		defer close(cleanupDone)
-		m.foreign.cleanupUnused()
+		m.cleanUpUnusedRelays()
 	}()
 
 	select {
 	case <-cleanupDone:
 	case <-time.After(2 * time.Second):
-		t.Fatal("cleanupUnused blocked on an in-progress relay dial")
+		t.Fatal("cleanUpUnusedRelays blocked on an in-progress relay dial while holding the relay map lock")
 	}
 
+	m.relayClientsMutex.RLock()
+	_, stillTracked := m.relayClients[serverAddr]
+	m.relayClientsMutex.RUnlock()
+	require.True(t, stillTracked, "an in-progress relay dial must not be evicted by cleanup")
+
 	// Release the hanging dial so the goroutine can exit cleanly.
 	mCancel()
 	select {
 	case <-dialDone:
 	case <-time.After(5 * time.Second):
-		t.Fatal("foreign OpenConn did not return after context cancellation")
+		t.Fatal("openConnVia did not return after context cancellation")
 	}
 }

shared/relay/client/manager_relaystates_test.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"net"
+	"net/netip"
 	"sync"
 	"testing"
 	"time"
@@ -12,16 +13,13 @@ import (
 
 // stallingRelayListener accepts TCP connections and holds them open without ever
 // responding, so a relay handshake dialed against it blocks until its context is
-// cancelled. accepted is signalled once per incoming connection so a caller can
-// wait until a dial has actually reached the listener. It returns the
-// "rel://host:port" URL to dial.
-func stallingRelayListener(t *testing.T) (string, <-chan struct{}) {
+// cancelled. It returns the "rel://host:port" URL to dial.
+func stallingRelayListener(t *testing.T) string {
 	t.Helper()
 
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 
-	accepted := make(chan struct{}, 1)
 	var mu sync.Mutex
 	var conns []net.Conn
 	go func() {
@@ -33,10 +31,6 @@ func stallingRelayListener(t *testing.T) (string, <-chan struct{}) {
 			mu.Lock()
 			conns = append(conns, c)
 			mu.Unlock()
-			select {
-			case accepted <- struct{}{}:
-			default:
-			}
 		}
 	}()
 	t.Cleanup(func() {
@@ -48,14 +42,14 @@ func stallingRelayListener(t *testing.T) (string, <-chan struct{}) {
 		mu.Unlock()
 	})
 
-	return "rel://" + ln.Addr().String(), accepted
+	return "rel://" + ln.Addr().String()
 }
 
 // TestRelayStates_DoesNotBlockOnRealHangingDial is a regression test for
 // RelayStates() called by a "status -d command" hanging behind an in-progress
-// foreign relay dial.
+// relay dial.
 func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
-	serverAddr, accepted := stallingRelayListener(t)
+	serverAddr := stallingRelayListener(t)
 
 	mCtx, mCancel := context.WithCancel(context.Background())
 	t.Cleanup(mCancel)
@@ -65,14 +59,15 @@ func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	dialDone := make(chan struct{})
 	go func() {
 		defer close(dialDone)
-		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
+		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
 	}()
 
-	select {
-	case <-accepted:
-	case <-time.After(5 * time.Second):
-		t.Fatal("relay dial did not reach the listener")
-	}
+	require.Eventually(t, func() bool {
+		m.relayClientsMutex.RLock()
+		defer m.relayClientsMutex.RUnlock()
+		_, ok := m.relayClients[serverAddr]
+		return ok
+	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
 
 	done := make(chan []RelayConnState, 1)
 	go func() {
@@ -91,6 +86,6 @@ func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	select {
 	case <-dialDone:
 	case <-time.After(5 * time.Second):
-		t.Fatal("foreign OpenConn did not return after context cancellation")
+		t.Fatal("openConnVia did not return after context cancellation")
 	}
 }

shared/relay/client/manager_serverip_test.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"io"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -84,7 +85,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 	t.Run("no server IP, dial fails", func(t *testing.T) {
 		dialCtx, dialCancel := context.WithTimeout(ctx, 5*time.Second)
 		defer dialCancel()
-		_, err := mgrAlice.OpenConn(dialCtx, RelayServer{Addr: brokenFQDN}, "bob", true)
+		_, err := mgrAlice.OpenConn(dialCtx, brokenFQDN, "bob", netip.Addr{})
 		if err == nil {
 			t.Fatalf("expected OpenConn to fail without server IP, got success")
 		}
@@ -94,7 +95,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 		// Bob waits for Alice's incoming peer connection on his side.
 		bobSideCh := make(chan error, 1)
 		go func() {
-			conn, err := mgrBob.OpenConn(ctx, RelayServer{Addr: bobRealAddr}, "alice", false)
+			conn, err := mgrBob.OpenConn(ctx, bobRealAddr, "alice", netip.Addr{})
 			if err != nil {
 				bobSideCh <- err
 				return
@@ -112,7 +113,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 			bobSideCh <- nil
 		}()
 
-		aliceConn, err := mgrAlice.OpenConn(ctx, RelayServer{Addr: brokenFQDN, IP: bobAdvertisedIP}, "bob", true)
+		aliceConn, err := mgrAlice.OpenConn(ctx, brokenFQDN, "bob", bobAdvertisedIP)
 		if err != nil {
 			t.Fatalf("alice OpenConn with server IP: %s", err)
 		}

shared/relay/client/manager_test.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"fmt"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -105,11 +106,11 @@ func TestForeignConn(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to get relay address: %s", err)
 	}
-	connAliceToBob, err := clientAlice.OpenConn(ctx, RelayServer{Addr: bobsSrvAddr}, "bob", true)
+	connAliceToBob, err := clientAlice.OpenConn(ctx, bobsSrvAddr, "bob", netip.Addr{})
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
-	connBobToAlice, err := clientBob.OpenConn(ctx, RelayServer{Addr: bobsSrvAddr}, "alice", false)
+	connBobToAlice, err := clientBob.OpenConn(ctx, bobsSrvAddr, "alice", netip.Addr{})
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -209,7 +210,7 @@ func TestForeginConnClose(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
-	conn, err := mgr.OpenConn(ctx, RelayServer{Addr: toURL(srvCfg2)[0]}, "bob", true)
+	conn, err := mgr.OpenConn(ctx, toURL(srvCfg2)[0], "bob", netip.Addr{})
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -301,7 +302,7 @@ func TestForeignAutoClose(t *testing.T) {
 	}
 
 	t.Log("open connection to another peer")
-	if _, err = mgr.OpenConn(ctx, RelayServer{Addr: foreignServerURL}, "anotherpeer", true); err == nil {
+	if _, err = mgr.OpenConn(ctx, foreignServerURL, "anotherpeer", netip.Addr{}); err == nil {
 		t.Fatalf("should have failed to open connection to another peer")
 	}
 
@@ -371,7 +372,7 @@ func TestAutoReconnect(t *testing.T) {
 	if err != nil {
 		t.Errorf("failed to get relay address: %s", err)
 	}
-	conn, err := clientAlice.OpenConn(ctx, RelayServer{Addr: ra}, "bob", false)
+	conn, err := clientAlice.OpenConn(ctx, ra, "bob", netip.Addr{})
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -391,7 +392,7 @@ func TestAutoReconnect(t *testing.T) {
 	}
 
 	log.Infof("reopent the connection")
-	_, err = clientAlice.OpenConn(ctx, RelayServer{Addr: ra}, "bob", false)
+	_, err = clientAlice.OpenConn(ctx, ra, "bob", netip.Addr{})
 	if err != nil {
 		t.Errorf("failed to open channel: %s", err)
 	}
@@ -453,7 +454,7 @@ func TestNotifierDoubleAdd(t *testing.T) {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
 
-	conn1, err := clientAlice.OpenConn(ctx, RelayServer{Addr: clientAlice.ServerURLs()[0]}, "bob", false)
+	conn1, err := clientAlice.OpenConn(ctx, clientAlice.ServerURLs()[0], "bob", netip.Addr{})
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}

shared/relay/messages/address/address.go

@@ -0,0 +1,21 @@
+// Deprecated: This package is deprecated and will be removed in a future release.
+package address
+
+import (
+	"bytes"
+	"encoding/gob"
+	"fmt"
+)
+
+type Address struct {
+	URL string
+}
+
+func (addr *Address) Marshal() ([]byte, error) {
+	var buf bytes.Buffer
+	enc := gob.NewEncoder(&buf)
+	if err := enc.Encode(addr); err != nil {
+		return nil, fmt.Errorf("encode Address: %w", err)
+	}
+	return buf.Bytes(), nil
+}

shared/relay/messages/auth/auth.go

@@ -0,0 +1,43 @@
+// Deprecated: This package is deprecated and will be removed in a future release.
+package auth
+
+import (
+	"bytes"
+	"encoding/gob"
+	"fmt"
+)
+
+type Algorithm int
+
+const (
+	AlgoUnknown Algorithm = iota
+	AlgoHMACSHA256
+	AlgoHMACSHA512
+)
+
+func (a Algorithm) String() string {
+	switch a {
+	case AlgoHMACSHA256:
+		return "HMAC-SHA256"
+	case AlgoHMACSHA512:
+		return "HMAC-SHA512"
+	default:
+		return "Unknown"
+	}
+}
+
+type Msg struct {
+	AuthAlgorithm  Algorithm
+	AdditionalData []byte
+}
+
+func UnmarshalMsg(data []byte) (*Msg, error) {
+	var msg *Msg
+
+	buf := bytes.NewBuffer(data)
+	dec := gob.NewDecoder(buf)
+	if err := dec.Decode(&msg); err != nil {
+		return nil, fmt.Errorf("decode Msg: %w", err)
+	}
+	return msg, nil
+}

shared/relay/messages/message.go

@@ -14,10 +14,9 @@ const (
 	CurrentProtocolVersion = 1
 
 	MsgTypeUnknown MsgType = 0
-	// MsgTypeHello and MsgTypeHelloResponse are the removed legacy handshake
-	// message types. They are retained only to reserve wire values 1 and 2 so
-	// the values are never reused; the server rejects both.
-	MsgTypeHello         = 1
+	// Deprecated: Use MsgTypeAuth instead.
+	MsgTypeHello = 1
+	// Deprecated: Use MsgTypeAuthResponse instead.
 	MsgTypeHelloResponse = 2
 	MsgTypeTransport     = 3
 	MsgTypeClose         = 4
@@ -43,6 +42,10 @@ const (
 	offsetAuthPeerID    = sizeOfProtoHeader + sizeOfMagicByte
 	headerTotalSizeAuth = sizeOfProtoHeader + headerSizeAuth
 
+	// hello message
+	headerSizeHello     = sizeOfMagicByte + peerIDSize
+	headerSizeHelloResp = 0
+
 	// transport
 	headerSizeTransport      = peerIDSize
 	offsetTransportID        = sizeOfProtoHeader
@@ -110,6 +113,7 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) {
 	msgType := MsgType(msg[1])
 	switch msgType {
 	case
+		MsgTypeHello,
 		MsgTypeAuth,
 		MsgTypeTransport,
 		MsgTypeClose,
@@ -131,6 +135,7 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	msgType := MsgType(msg[1])
 	switch msgType {
 	case
+		MsgTypeHelloResponse,
 		MsgTypeAuthResponse,
 		MsgTypeTransport,
 		MsgTypeClose,
@@ -143,6 +148,67 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	}
 }
 
+// Deprecated: Use MarshalAuthMsg instead.
+// MarshalHelloMsg initial hello message
+// The Hello message is the first message sent by a client after establishing a connection with the Relay server. This
+// message is used to authenticate the client with the server. The authentication is done using an HMAC method.
+// The protocol does not limit to use HMAC, it can be any other method. If the authentication failed the server will
+// close the network connection without any response.
+func MarshalHelloMsg(peerID PeerID, additions []byte) ([]byte, error) {
+	msg := make([]byte, sizeOfProtoHeader+sizeOfMagicByte, sizeOfProtoHeader+headerSizeHello+len(additions))
+
+	msg[0] = byte(CurrentProtocolVersion)
+	msg[1] = byte(MsgTypeHello)
+
+	copy(msg[sizeOfProtoHeader:sizeOfProtoHeader+sizeOfMagicByte], magicHeader)
+
+	msg = append(msg, peerID[:]...)
+	msg = append(msg, additions...)
+
+	return msg, nil
+}
+
+// Deprecated: Use UnmarshalAuthMsg instead.
+// UnmarshalHelloMsg extracts peerID and the additional data from the hello message. The Additional data is used to
+// authenticate the client with the server.
+func UnmarshalHelloMsg(msg []byte) (*PeerID, []byte, error) {
+	if len(msg) < sizeOfProtoHeader+headerSizeHello {
+		return nil, nil, ErrInvalidMessageLength
+	}
+	if !bytes.Equal(msg[sizeOfProtoHeader:sizeOfProtoHeader+sizeOfMagicByte], magicHeader) {
+		return nil, nil, errors.New("invalid magic header")
+	}
+
+	peerID := PeerID(msg[sizeOfProtoHeader+sizeOfMagicByte : sizeOfProtoHeader+headerSizeHello])
+
+	return &peerID, msg[headerSizeHello:], nil
+}
+
+// Deprecated: Use MarshalAuthResponse instead.
+// MarshalHelloResponse creates a response message to the hello message.
+// In case of success connection the server response with a Hello Response message. This message contains the server's
+// instance URL. This URL will be used by choose the common Relay server in case if the peers are in different Relay
+// servers.
+func MarshalHelloResponse(additionalData []byte) ([]byte, error) {
+	msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+headerSizeHelloResp+len(additionalData))
+
+	msg[0] = byte(CurrentProtocolVersion)
+	msg[1] = byte(MsgTypeHelloResponse)
+
+	msg = append(msg, additionalData...)
+
+	return msg, nil
+}
+
+// Deprecated: Use UnmarshalAuthResponse instead.
+// UnmarshalHelloResponse extracts the additional data from the hello response message.
+func UnmarshalHelloResponse(msg []byte) ([]byte, error) {
+	if len(msg) < sizeOfProtoHeader+headerSizeHelloResp {
+		return nil, ErrInvalidMessageLength
+	}
+	return msg, nil
+}
+
 // MarshalAuthMsg initial authentication message
 // The Auth message is the first message sent by a client after establishing a connection with the Relay server. This
 // message is used to authenticate the client with the server. The authentication is done using an HMAC method.

shared/relay/messages/message_test.go

@@ -4,11 +4,28 @@ import (
 	"testing"
 )
 
-func TestDetermineClientMessageTypeRejectsHello(t *testing.T) {
-	// The reserved legacy Hello message (type 1) must be rejected by the server.
-	msg := []byte{byte(CurrentProtocolVersion), byte(MsgTypeHello)}
-	if _, err := DetermineClientMessageType(msg); err == nil {
-		t.Fatalf("expected hello message type to be rejected")
+func TestMarshalHelloMsg(t *testing.T) {
+	peerID := HashID("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
+	msg, err := MarshalHelloMsg(peerID, nil)
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+
+	msgType, err := DetermineClientMessageType(msg)
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+
+	if msgType != MsgTypeHello {
+		t.Errorf("expected %d, got %d", MsgTypeHello, msgType)
+	}
+
+	receivedPeerID, _, err := UnmarshalHelloMsg(msg)
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	if receivedPeerID.String() != peerID.String() {
+		t.Errorf("expected %s, got %s", peerID, receivedPeerID)
 	}
 }
 

shared/testing_helpers/populate_fields.go

@@ -1,101 +0,0 @@
-package testing_helpers
-
-import (
-	"fmt"
-	"net/netip"
-	"reflect"
-)
-
-type PopulateFields struct {
-	CustomFieldSetters map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error)
-	TagsToSkip         map[string]string
-}
-
-func NewPopulateFields() *PopulateFields {
-	return &PopulateFields{CustomFieldSetters: defaultCustomFieldSetters(), TagsToSkip: make(map[string]string)}
-}
-
-func (p *PopulateFields) WithCustomFieldSetter(t reflect.Type, f func(this *PopulateFields, field reflect.Value) (int, error)) *PopulateFields {
-	p.CustomFieldSetters[t] = f
-	return p
-}
-
-func (p *PopulateFields) WithSkippedTag(tag, value string) *PopulateFields {
-	p.TagsToSkip[tag] = value
-	return p
-}
-
-func (p *PopulateFields) PopulateAll(v reflect.Value) (int, error) {
-	typ := v.Type()
-	totalExportedFields := 0
-	for i := 0; i < typ.NumField(); i++ {
-		f := typ.Field(i)
-		if f.PkgPath != "" { // unexported
-			continue
-		}
-
-		if p.skippedTagPresent(f.Tag) {
-			continue
-		}
-
-		numOfExportedFields, err := p.setNonZero(v.Field(i))
-		totalExportedFields += numOfExportedFields
-		if err != nil {
-			return totalExportedFields, err
-		}
-	}
-	return totalExportedFields, nil
-}
-
-// setNonZero assigns a deterministic non-zero value to a field based on its kind,
-// recursing into nested structs and populating one element of slice fields.
-func (p *PopulateFields) setNonZero(field reflect.Value) (int, error) {
-	if f, ok := p.CustomFieldSetters[field.Type()]; ok {
-		return f(p, field)
-	}
-
-	switch field.Kind() {
-	case reflect.String:
-		field.SetString("non-zero")
-	case reflect.Bool:
-		field.SetBool(true)
-	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
-		field.SetInt(7)
-	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
-		field.SetUint(7)
-	case reflect.Float32, reflect.Float64:
-		field.SetFloat(7)
-	case reflect.Struct:
-		n, err := p.PopulateAll(field)
-		return n + 1, err
-	case reflect.Slice:
-		s := reflect.MakeSlice(field.Type(), 1, 1)
-		_, err := p.setNonZero(s.Index(0))
-		if err != nil {
-			return 0, err
-		}
-		field.Set(s)
-	default:
-		return 0, fmt.Errorf("unhandled field kind %s; extend setNonZero", field.Kind())
-	}
-
-	return 1, nil
-}
-
-func defaultCustomFieldSetters() map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error) {
-	return map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error){
-		reflect.TypeOf(netip.Prefix{}): func(_ *PopulateFields, field reflect.Value) (int, error) {
-			field.Set(reflect.ValueOf(netip.MustParsePrefix("10.0.0.0/24")))
-			return 1, nil
-		},
-	}
-}
-
-func (p *PopulateFields) skippedTagPresent(t reflect.StructTag) bool {
-	for tag, value := range p.TagsToSkip {
-		if v := t.Get(tag); v == value {
-			return true
-		}
-	}
-	return false
-}

trustedproxy/trustedproxy.go

@@ -1,132 +0,0 @@
-package trustedproxy
-
-import (
-	"fmt"
-	"net/netip"
-	"strings"
-)
-
-// List holds a parsed set of trusted upstream proxy prefixes and answers trust
-// questions against it. The zero value (and a nil *List) is a valid, empty list
-// that never trusts any address, so callers can use it without a nil check.
-type List struct {
-	prefixes []netip.Prefix
-}
-
-// Parse parses a comma-separated list of CIDR prefixes or bare IPs into a List.
-// Bare IPs are converted to single-host prefixes (/32 or /128). An empty input
-// yields an empty List that trusts nothing.
-func Parse(raw string) (*List, error) {
-	if raw == "" {
-		return &List{}, nil
-	}
-
-	parts := strings.Split(raw, ",")
-	prefixes := make([]netip.Prefix, 0, len(parts))
-	for _, part := range parts {
-		part = strings.TrimSpace(part)
-		if part == "" {
-			continue
-		}
-
-		prefix, err := netip.ParsePrefix(part)
-		if err == nil {
-			prefixes = append(prefixes, prefix)
-			continue
-		}
-
-		addr, addrErr := netip.ParseAddr(part)
-		if addrErr != nil {
-			return nil, fmt.Errorf("parse trusted proxy %q: not a valid CIDR or IP: %w", part, addrErr)
-		}
-
-		bits := 32
-		if addr.Is6() {
-			bits = 128
-		}
-		prefixes = append(prefixes, netip.PrefixFrom(addr, bits))
-	}
-	return &List{prefixes: prefixes}, nil
-}
-
-// FromPrefixes wraps an already-parsed set of prefixes in a List.
-func FromPrefixes(prefixes []netip.Prefix) *List {
-	return &List{prefixes: prefixes}
-}
-
-// Empty reports whether the list contains no prefixes.
-func (l *List) Empty() bool {
-	return l == nil || len(l.prefixes) == 0
-}
-
-// IsTrusted reports whether the given host:port or bare IP falls within the list.
-func (l *List) IsTrusted(remoteAddr string) bool {
-	if l.Empty() {
-		return false
-	}
-	return l.Contains(ExtractHostIP(remoteAddr))
-}
-
-// Contains reports whether the given address falls within any trusted prefix.
-func (l *List) Contains(addr netip.Addr) bool {
-	if l.Empty() || !addr.IsValid() {
-		return false
-	}
-	for _, prefix := range l.prefixes {
-		if prefix.Contains(addr) {
-			return true
-		}
-	}
-	return false
-}
-
-// ResolveClientIP extracts the real client IP from X-Forwarded-For using the
-// list. It walks the XFF chain right-to-left, skipping IPs that match trusted
-// prefixes; the first untrusted IP is the real client. If the list is empty or
-// remoteAddr is not trusted, it returns the remoteAddr IP directly, ignoring any
-// forwarding headers.
-func (l *List) ResolveClientIP(remoteAddr, xff string) netip.Addr {
-	remoteIP := ExtractHostIP(remoteAddr)
-
-	if l.Empty() || !l.Contains(remoteIP) {
-		return remoteIP
-	}
-
-	if xff == "" {
-		return remoteIP
-	}
-
-	parts := strings.Split(xff, ",")
-	for i := len(parts) - 1; i >= 0; i-- {
-		ip := strings.TrimSpace(parts[i])
-		if ip == "" {
-			continue
-		}
-		addr, err := netip.ParseAddr(ip)
-		if err != nil {
-			continue
-		}
-		addr = addr.Unmap()
-		if !l.Contains(addr) {
-			return addr
-		}
-	}
-
-	if first := strings.TrimSpace(parts[0]); first != "" {
-		if addr, err := netip.ParseAddr(first); err == nil {
-			return addr.Unmap()
-		}
-	}
-	return remoteIP
-}
-
-// ExtractHostIP parses the IP from a host:port string and returns it unmapped.
-func ExtractHostIP(hostPort string) netip.Addr {
-	if ap, err := netip.ParseAddrPort(hostPort); err == nil {
-		return ap.Addr().Unmap()
-	}
-	if addr, err := netip.ParseAddr(hostPort); err == nil {
-		return addr.Unmap()
-	}
-	return netip.Addr{}
-}

trustedproxy/trustedproxy_test.go

@@ -1,216 +0,0 @@
-package trustedproxy
-
-import (
-	"net/netip"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func TestParse(t *testing.T) {
-	tests := []struct {
-		name    string
-		raw     string
-		want    []netip.Prefix
-		wantErr bool
-	}{
-		{
-			name: "empty string returns empty list",
-			raw:  "",
-			want: nil,
-		},
-		{
-			name: "single CIDR",
-			raw:  "10.0.0.0/8",
-			want: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
-		},
-		{
-			name: "single bare IPv4",
-			raw:  "1.2.3.4",
-			want: []netip.Prefix{netip.MustParsePrefix("1.2.3.4/32")},
-		},
-		{
-			name: "single bare IPv6",
-			raw:  "::1",
-			want: []netip.Prefix{netip.MustParsePrefix("::1/128")},
-		},
-		{
-			name: "comma-separated CIDRs",
-			raw:  "10.0.0.0/8, 192.168.1.0/24",
-			want: []netip.Prefix{
-				netip.MustParsePrefix("10.0.0.0/8"),
-				netip.MustParsePrefix("192.168.1.0/24"),
-			},
-		},
-		{
-			name: "mixed CIDRs and bare IPs",
-			raw:  "10.0.0.0/8, 1.2.3.4, fd00::/8",
-			want: []netip.Prefix{
-				netip.MustParsePrefix("10.0.0.0/8"),
-				netip.MustParsePrefix("1.2.3.4/32"),
-				netip.MustParsePrefix("fd00::/8"),
-			},
-		},
-		{
-			name: "whitespace around entries",
-			raw:  "  10.0.0.0/8 , 192.168.0.0/16  ",
-			want: []netip.Prefix{
-				netip.MustParsePrefix("10.0.0.0/8"),
-				netip.MustParsePrefix("192.168.0.0/16"),
-			},
-		},
-		{
-			name: "trailing comma produces no extra entry",
-			raw:  "10.0.0.0/8,",
-			want: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
-		},
-		{
-			name:    "invalid entry",
-			raw:     "not-an-ip",
-			wantErr: true,
-		},
-		{
-			name:    "partially invalid",
-			raw:     "10.0.0.0/8, garbage",
-			wantErr: true,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := Parse(tt.raw)
-			if tt.wantErr {
-				require.Error(t, err)
-				return
-			}
-			require.NoError(t, err)
-			assert.Equal(t, tt.want, got.prefixes)
-		})
-	}
-}
-
-func TestListIsTrusted(t *testing.T) {
-	list, err := Parse("10.0.0.0/8, 192.168.1.0/24, fd00::/8")
-	require.NoError(t, err)
-
-	tests := []struct {
-		name string
-		addr string
-		list *List
-		want bool
-	}{
-		{"nil list", "10.0.0.1", nil, false},
-		{"empty list", "10.0.0.1", &List{}, false},
-		{"IP within /8 prefix", "10.1.2.3", list, true},
-		{"IP within /24 prefix", "192.168.1.100", list, true},
-		{"IP outside all prefixes", "203.0.113.50", list, false},
-		{"boundary IP just outside prefix", "192.168.2.1", list, false},
-		{"unparsable IP", "not-an-ip", list, false},
-		{"IPv6 in trusted range", "fd00::1", list, true},
-		{"IPv6 outside range", "2001:db8::1", list, false},
-		{"empty string", "", list, false},
-		{"host:port within prefix", "10.1.2.3:9999", list, true},
-		{"host:port outside prefix", "203.0.113.50:9999", list, false},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert.Equal(t, tt.want, tt.list.IsTrusted(tt.addr))
-		})
-	}
-}
-
-func TestListResolveClientIP(t *testing.T) {
-	trusted, err := Parse("10.0.0.0/8, 172.16.0.0/12")
-	require.NoError(t, err)
-
-	tests := []struct {
-		name       string
-		remoteAddr string
-		xff        string
-		list       *List
-		want       netip.Addr
-	}{
-		{
-			name:       "empty list returns RemoteAddr",
-			remoteAddr: "203.0.113.50:9999",
-			xff:        "1.2.3.4",
-			list:       &List{},
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "nil list returns RemoteAddr",
-			remoteAddr: "203.0.113.50:9999",
-			xff:        "1.2.3.4",
-			list:       nil,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "untrusted RemoteAddr ignores XFF",
-			remoteAddr: "203.0.113.50:9999",
-			xff:        "1.2.3.4, 10.0.0.1",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "trusted RemoteAddr with single client in XFF",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "203.0.113.50",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "trusted RemoteAddr walks past trusted entries in XFF",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "203.0.113.50, 10.0.0.2, 172.16.0.5",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "trusted RemoteAddr with empty XFF falls back to RemoteAddr",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "",
-			list:       trusted,
-			want:       netip.MustParseAddr("10.0.0.1"),
-		},
-		{
-			name:       "all XFF IPs trusted returns leftmost",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "10.0.0.2, 172.16.0.1, 10.0.0.3",
-			list:       trusted,
-			want:       netip.MustParseAddr("10.0.0.2"),
-		},
-		{
-			name:       "XFF with whitespace",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        " 203.0.113.50 , 10.0.0.2 ",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "XFF with empty segments",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "203.0.113.50,,10.0.0.2",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "multi-hop with mixed trust",
-			remoteAddr: "10.0.0.1:5000",
-			xff:        "8.8.8.8, 203.0.113.50, 172.16.0.1",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-		{
-			name:       "RemoteAddr without port",
-			remoteAddr: "10.0.0.1",
-			xff:        "203.0.113.50",
-			list:       trusted,
-			want:       netip.MustParseAddr("203.0.113.50"),
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert.Equal(t, tt.want, tt.list.ResolveClientIP(tt.remoteAddr, tt.xff))
-		})
-	}
-}