[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/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/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/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

@@ -335,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=

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/labelgen/labelgen.go

@@ -3,10 +3,9 @@ package labelgen
 
 import (
 	"fmt"
+	"math/rand"
 	"sort"
 	"sync"
-
-	"github.com/netbirdio/netbird/management/server/util"
 )
 
 // pickAttempts caps the random retries before falling back to the
@@ -41,15 +40,16 @@ func uniqueWords() []string {
 // PickUnique selects a label not already in `taken`. It tries up to
 // pickAttempts random picks; on exhaustion it scans the deduplicated
 // wordlist for any remaining free entry, and if none is left appends
-// `-<fallbackSuffix>` to a random word and returns.
-func PickUnique(taken map[string]struct{}, fallbackSuffix string) string {
+// `-<fallbackSuffix>` to a deterministic word and returns. The caller
+// is responsible for seeding rng (math/rand).
+func PickUnique(rng *rand.Rand, taken map[string]struct{}, fallbackSuffix string) string {
 	pool := uniqueWords()
 	if len(pool) == 0 {
 		return fallbackSuffix
 	}
 
 	for i := 0; i < pickAttempts; i++ {
-		w := pool[util.RandIntn(len(pool))]
+		w := pool[rng.Intn(len(pool))]
 		if _, ok := taken[w]; !ok {
 			return w
 		}
@@ -61,6 +61,6 @@ func PickUnique(taken map[string]struct{}, fallbackSuffix string) string {
 		}
 	}
 
-	w := pool[util.RandIntn(len(pool))]
+	w := pool[rng.Intn(len(pool))]
 	return fmt.Sprintf("%s-%s", w, fallbackSuffix)
 }

management/internals/modules/agentnetwork/labelgen/labelgen_test.go

@@ -1,7 +1,7 @@
 package labelgen
 
 import (
-	"slices"
+	"math/rand"
 	"strings"
 	"testing"
 
@@ -9,12 +9,19 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TestPickUnique_ReturnsWordFromPool confirms a pick against an empty
-// taken set is always drawn verbatim from the wordlist.
-func TestPickUnique_ReturnsWordFromPool(t *testing.T) {
-	got := PickUnique(map[string]struct{}{}, "abcd")
+// TestPickUnique_DeterministicWithSeededRng locks the property the
+// caller relies on: same seed + same taken set → same pick. Without
+// that, the bootstrap flow can't reproduce a label across retries.
+func TestPickUnique_DeterministicWithSeededRng(t *testing.T) {
+	taken := map[string]struct{}{}
 
-	assert.True(t, slices.Contains(uniqueWords(), got), "Pick %q must be drawn from the wordlist", got)
+	rngA := rand.New(rand.NewSource(42))
+	rngB := rand.New(rand.NewSource(42))
+
+	a := PickUnique(rngA, taken, "abcd")
+	b := PickUnique(rngB, taken, "abcd")
+
+	assert.Equal(t, a, b, "Same seed and taken set must produce identical pick")
 }
 
 // TestPickUnique_AvoidsTakenWordsWhenMostAreReserved seeds taken with
@@ -39,7 +46,8 @@ func TestPickUnique_AvoidsTakenWordsWhenMostAreReserved(t *testing.T) {
 		taken[w] = struct{}{}
 	}
 
-	got := PickUnique(taken, "abcd")
+	rng := rand.New(rand.NewSource(7))
+	got := PickUnique(rng, taken, "abcd")
 
 	_, isFree := free[got]
 	assert.True(t, isFree, "PickUnique must return one of the free words; got %q", got)
@@ -57,7 +65,8 @@ func TestPickUnique_FallsBackWhenAllReserved(t *testing.T) {
 		taken[w] = struct{}{}
 	}
 
-	got := PickUnique(taken, "abcd")
+	rng := rand.New(rand.NewSource(99))
+	got := PickUnique(rng, taken, "abcd")
 
 	assert.True(t, strings.HasSuffix(got, "-abcd"), "Exhausted pool must produce <word>-<suffix>; got %q", got)
 

management/internals/modules/agentnetwork/manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand"
 	"slices"
 	"strings"
 	"sync"
@@ -122,6 +123,11 @@ type managerImpl struct {
 	// accountID, then by synthesised service ID.
 	reconcileMu    sync.Mutex
 	reconcileCache map[string]map[string]*proto.ProxyMapping
+
+	// labelRngMu guards labelRng. PickUnique consumes math/rand.Source
+	// state; concurrent provider creates would otherwise race.
+	labelRngMu sync.Mutex
+	labelRng   *rand.Rand
 }
 
 // NewManager constructs the persistent Agent Network manager. The
@@ -141,6 +147,7 @@ func NewManager(
 		permissionsManager: permissionsManager,
 		proxyController:    proxyController,
 		reconcileCache:     make(map[string]map[string]*proto.ProxyMapping),
+		labelRng:           rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
 }
 
@@ -646,7 +653,9 @@ func (m *managerImpl) bootstrapSettingsIfNeeded(ctx context.Context, accountID,
 		suffix = suffix[:4]
 	}
 
-	subdomain := labelgen.PickUnique(taken, suffix)
+	m.labelRngMu.Lock()
+	subdomain := labelgen.PickUnique(m.labelRng, taken, suffix)
+	m.labelRngMu.Unlock()
 
 	now := time.Now().UTC()
 	settings := &types.Settings{

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.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand"
 	"net"
 	"net/netip"
 	"os"
@@ -62,7 +63,7 @@ const (
 type userLoggedInOnce bool
 
 func cacheEntryExpiration() time.Duration {
-	r := util.RandIntn(int(nbcache.DefaultIDPCacheExpirationMax.Milliseconds()-nbcache.DefaultIDPCacheExpirationMin.Milliseconds())) + int(nbcache.DefaultIDPCacheExpirationMin.Milliseconds())
+	r := rand.Intn(int(nbcache.DefaultIDPCacheExpirationMax.Milliseconds()-nbcache.DefaultIDPCacheExpirationMin.Milliseconds())) + int(nbcache.DefaultIDPCacheExpirationMin.Milliseconds())
 	return time.Duration(r) * time.Millisecond
 }
 
@@ -2454,7 +2455,8 @@ func (am *DefaultAccountManager) ensureIPv6Subnet(ctx context.Context, transacti
 		return transaction.UpdateAccountNetworkV6(ctx, accountID, network.NetV6)
 	}
 	if network.NetV6.IP == nil {
-		network.NetV6 = types.AllocateIPv6Subnet()
+		r := rand.New(rand.NewSource(time.Now().UnixNano()))
+		network.NetV6 = types.AllocateIPv6Subnet(r)
 
 		// Sync settings to match the allocated subnet so SaveAccountSettings persists it.
 		ones, _ := network.NetV6.Mask.Size()

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/group_ipv6_test.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"math/rand"
 	"testing"
 	"time"
 
@@ -27,7 +28,7 @@ func TestGroupIPv6Assignment(t *testing.T) {
 	require.NoError(t, err)
 
 	// Allocate IPv6 subnet for the account
-	account.Network.NetV6 = types.AllocateIPv6Subnet()
+	account.Network.NetV6 = types.AllocateIPv6Subnet(rand.New(rand.NewSource(time.Now().UnixNano())))
 	require.NoError(t, am.Store.SaveAccount(ctx, account))
 
 	// Create setup key

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/idp/util.go

@@ -2,12 +2,11 @@ package idp
 
 import (
 	"encoding/json"
+	"math/rand"
 	"net/url"
 	"os"
 	"strings"
 	"time"
-
-	"github.com/netbirdio/netbird/management/server/util"
 )
 
 var (
@@ -34,32 +33,31 @@ func GeneratePassword(passwordLength, minSpecialChar, minNum, minUpperCase int)
 
 	//Set special character
 	for i := 0; i < minSpecialChar; i++ {
-		random := util.RandIntn(len(specialCharSet))
+		random := rand.Intn(len(specialCharSet))
 		password.WriteString(string(specialCharSet[random]))
 	}
 
 	//Set numeric
 	for i := 0; i < minNum; i++ {
-		random := util.RandIntn(len(numberSet))
+		random := rand.Intn(len(numberSet))
 		password.WriteString(string(numberSet[random]))
 	}
 
 	//Set uppercase
 	for i := 0; i < minUpperCase; i++ {
-		random := util.RandIntn(len(upperCharSet))
+		random := rand.Intn(len(upperCharSet))
 		password.WriteString(string(upperCharSet[random]))
 	}
 
 	remainingLength := passwordLength - minSpecialChar - minNum - minUpperCase
 	for i := 0; i < remainingLength; i++ {
-		random := util.RandIntn(len(allCharSet))
+		random := rand.Intn(len(allCharSet))
 		password.WriteString(string(allCharSet[random]))
 	}
 	inRune := []rune(password.String())
-	for i := len(inRune) - 1; i > 0; i-- {
-		j := util.RandIntn(i + 1)
+	rand.Shuffle(len(inRune), func(i, j int) {
 		inRune[i], inRune[j] = inRune[j], inRune[i]
-	}
+	})
 	return string(inRune)
 }
 

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")

management/server/types/network.go

@@ -1,13 +1,14 @@
 package types
 
 import (
-	"crypto/rand"
 	"encoding/binary"
 	"fmt"
+	"math/rand"
 	"net"
 	"net/netip"
 	"slices"
 	"sync"
+	"time"
 
 	"github.com/c-robinson/iplib"
 	"github.com/rs/xid"
@@ -136,12 +137,14 @@ func NewNetwork() *Network {
 	n := iplib.NewNet4(net.ParseIP("100.64.0.0"), NetSize)
 	sub, _ := n.Subnet(SubnetSize)
 
-	intn := util.RandIntn(len(sub))
+	s := rand.NewSource(time.Now().UnixNano())
+	r := rand.New(s)
+	intn := r.Intn(len(sub))
 
 	return &Network{
 		Identifier: xid.New().String(),
 		Net:        sub[intn].IPNet,
-		NetV6:      AllocateIPv6Subnet(),
+		NetV6:      AllocateIPv6Subnet(r),
 		Dns:        "",
 		Serial:     0,
 	}
@@ -151,13 +154,18 @@ func NewNetwork() *Network {
 // The format follows RFC 4193 section 3.1: fd + 40-bit Global ID + 16-bit Subnet ID.
 // The Global ID and Subnet ID are randomized (simplified from the SHA-1 algorithm
 // in section 3.2.2), giving 2^56 possible /64 subnets across all accounts.
-func AllocateIPv6Subnet() net.IPNet {
+func AllocateIPv6Subnet(r *rand.Rand) net.IPNet {
 	ip := make(net.IP, 16)
 	ip[0] = 0xfd
-	// Bytes 1-5: 40-bit random Global ID, bytes 6-7: 16-bit random Subnet ID
-	if _, err := rand.Read(ip[1:8]); err != nil {
-		panic(err)
-	}
+	// Bytes 1-5: 40-bit random Global ID
+	ip[1] = byte(r.Intn(256))
+	ip[2] = byte(r.Intn(256))
+	ip[3] = byte(r.Intn(256))
+	ip[4] = byte(r.Intn(256))
+	ip[5] = byte(r.Intn(256))
+	// Bytes 6-7: 16-bit random Subnet ID
+	ip[6] = byte(r.Intn(256))
+	ip[7] = byte(r.Intn(256))
 
 	return net.IPNet{
 		IP:   ip,
@@ -209,10 +217,11 @@ func AllocatePeerIP(prefix netip.Prefix, takenIps []netip.Addr) (netip.Addr, err
 		taken[binary.BigEndian.Uint32(ab[:])] = struct{}{}
 	}
 
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
 	maxAttempts := (int(totalIPs) - len(taken)) / 100
 
 	for i := 0; i < maxAttempts; i++ {
-		offset := uint32(util.RandIntn(int(totalIPs-2))) + 1
+		offset := uint32(rng.Intn(int(totalIPs-2))) + 1
 		candidate := baseIP + offset
 		if _, exists := taken[candidate]; !exists {
 			return uint32ToIP(candidate), nil
@@ -236,7 +245,8 @@ func AllocateRandomPeerIP(prefix netip.Prefix) (netip.Addr, error) {
 	hostBits := 32 - prefix.Bits()
 	totalIPs := uint32(1 << hostBits)
 
-	offset := uint32(util.RandIntn(int(totalIPs-2))) + 1
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+	offset := uint32(rng.Intn(int(totalIPs-2))) + 1
 
 	candidate := baseIP + offset
 	return uint32ToIP(candidate), nil
@@ -252,26 +262,23 @@ func AllocateRandomPeerIPv6(prefix netip.Prefix) (netip.Addr, error) {
 
 	ip := prefix.Addr().As16()
 
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+
 	// Determine which byte the host bits start in
 	firstHostByte := ones / 8
 	// If the prefix doesn't end on a byte boundary, handle the partial byte
 	partialBits := ones % 8
 
-	var rnd [16]byte
-	if _, err := rand.Read(rnd[firstHostByte:]); err != nil {
-		return netip.Addr{}, err
-	}
-
 	if partialBits > 0 {
 		// Keep the network bits in the partial byte, randomize the rest
 		hostMask := byte(0xff >> partialBits)
-		ip[firstHostByte] = (ip[firstHostByte] & ^hostMask) | (rnd[firstHostByte] & hostMask)
+		ip[firstHostByte] = (ip[firstHostByte] & ^hostMask) | (byte(rng.Intn(256)) & hostMask)
 		firstHostByte++
 	}
 
 	// Randomize remaining full host bytes
 	for i := firstHostByte; i < 16; i++ {
-		ip[i] = rnd[i]
+		ip[i] = byte(rng.Intn(256))
 	}
 
 	// Avoid all-zeros and all-ones host parts by checking only host bits.

management/server/util/util.go

@@ -1,20 +1,5 @@
 package util
 
-import (
-	"crypto/rand"
-	"math/big"
-)
-
-// RandIntn returns a uniformly distributed int in [0, n) sourced from
-// crypto/rand. It panics if n <= 0 or the platform randomness source fails.
-func RandIntn(n int) int {
-	v, err := rand.Int(rand.Reader, big.NewInt(int64(n)))
-	if err != nil {
-		panic(err)
-	}
-	return int(v.Int64())
-}
-
 // Difference returns the elements in `a` that aren't in `b`.
 func Difference(a, b []string) []string {
 	mb := make(map[string]struct{}, len(b))

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) {

relay/cmd/root.go

@@ -173,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)
 	}
 }
 

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/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
-}