fix(traffic): harden UDP session lifecycle (#286)

* fix(traffic): harden UDP session lifecycle

* fix(traffic): await UDP runtime shutdown

commit: 574ac6767ce73ee088baca1bc7e09634531e1018
author: bnema · 38d ago · 2026-09-14 20:48 UTC
committer: GitHub · 38d ago · 2026-09-14 20:48 UTC
parent: d69854d
signature: GPG UnverifiedB5690EEEBB952194 The key is unknown

2 files · +699 -87

M internal/adapters/in/traffic/udp.go +188 -87


@@ -4,6 +4,7 @@ import (
4	4		"context"
5	5		"errors"
6	6		"net"
	7	+	"net/netip"
7	8		"strconv"
8	9		"sync"
9	10		"sync/atomic"

@@ -28,28 +29,42 @@ type udpEntryPointRuntime struct {
28	29		started atomic.Bool
29	30		closed atomic.Bool
30	31
31		-	ctx context.Context
32		-	cancel context.CancelFunc
33		-	done chan struct{}
34		-	doneOnce sync.Once
	32	+	ctx context.Context
	33	+	cancel context.CancelFunc
	34	+	// runWG tracks every goroutine owned by this runtime (read loop, datagram
	35	+	// workers, expiry loop and per-session backend loops). Counts are only added
	36	+	// while holding mu and before the owning goroutine starts; stop() flips
	37	+	// accepting under the same lock before it ever waits, so no Add can race a
	38	+	// Wait.
	39	+	runWG sync.WaitGroup
35	40
36	41		datagrams chan udpDatagram
37	42
38		-	mu sync.Mutex
39		-	sessions map[string]*udpSession
	43	+	// dialContext dials the upstream backend. It is a private instance field so
	44	+	// tests can drive in-flight dial interleavings deterministically; production
	45	+	// code always uses the standard UDP dialer set at construction and never
	46	+	// mutates it after the runtime starts.
	47	+	dialContext func(ctx context.Context, network, address string) (net.Conn, error)
	48	+
	49	+	mu sync.Mutex
	50	+	accepting bool
	51	+	sessions map[netip.AddrPort]*udpSession
40	52		}
41	53
42	54		type udpDatagram struct {
43		-	clientAddr net.Addr
	55	+	clientAddr netip.AddrPort
44	56		packet []byte
45	57		}
46	58
47	59		type udpSession struct {
48		-	clientAddr net.Addr
	60	+	clientAddr netip.AddrPort
	61	+	// remoteAddr is the immutable packet destination precomputed once at
	62	+	// publication, so backendLoop never has to allocate or format an address
	63	+	// per packet. It is shared read-only across the session's goroutines.
	64	+	remoteAddr *net.UDPAddr
49	65		backend net.Conn
50	66		backendRef domain.TrafficBackend
51	67		lastSeen atomic.Int64
52		-	done chan struct{}
53	68		once sync.Once
54	69		}
55	70

@@ -62,9 +77,12 @@ func newUDPEntryPointRuntime(parentCtx context.Context, manager *Manager, entryP
62	77		trusted: trusted,
63	78		ctx: ctx,
64	79		cancel: cancel,
65		-	done: make(chan struct{}),
	80	+	accepting: true,
66	81		datagrams: make(chan udpDatagram, udpDatagramBacklog),
67		-	sessions: map[string]*udpSession{},
	82	+	sessions: map[netip.AddrPort]*udpSession{},
	83	+	dialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
	84	+	return (&net.Dialer{}).DialContext(ctx, network, address)
	85	+	},
68	86		}
69	87		}
70	88

@@ -72,6 +90,13 @@ func (r *udpEntryPointRuntime) start() {
72	90		if !r.started.CompareAndSwap(false, true) {
73	91		return
74	92		}
	93	+	r.mu.Lock()
	94	+	if !r.accepting {
	95	+	r.mu.Unlock()
	96	+	return
	97	+	}
	98	+	r.runWG.Add(udpWorkerCount + 2) // read loop, expiry loop and datagram workers
	99	+	r.mu.Unlock()
75	100		entryPoint := r.entryPointSnapshot()
76	101		trafficInfo(r.ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Str("protocol", string(entryPoint.Protocol)).Msg("started udp traffic entrypoint")
77	102		for range udpWorkerCount {

@@ -82,7 +107,7 @@ func (r *udpEntryPointRuntime) start() {
82	107		}
83	108
84	109		func (r *udpEntryPointRuntime) readLoop() {
85		-	defer r.closeDone()
	110	+	defer r.runWG.Done()
86	111		buf := make([]byte, udpBufferSize)
87	112		for {
88	113		n, clientAddr, err := r.packetConn.ReadFrom(buf)

@@ -93,7 +118,12 @@ func (r *udpEntryPointRuntime) readLoop() {
93	118		r.counters.totalErrors.Add(1)
94	119		continue
95	120		}
96		-	if !trustedRemoteAddr(r.trustedSnapshot(), clientAddr) {
	121	+	key, ok := udpAddrPort(clientAddr)
	122	+	if !ok {
	123	+	r.counters.totalRefused.Add(1)
	124	+	continue
	125	+	}
	126	+	if !trustedAddrPort(r.trustedSnapshot(), key) {
97	127		r.counters.totalRefused.Add(1)
98	128		continue
99	129		}

@@ -108,7 +138,7 @@ func (r *udpEntryPointRuntime) readLoop() {
108	138		}
109	139		packet := append([]byte(nil), buf[:n]...)
110	140		select {
111		-	case r.datagrams <- udpDatagram{clientAddr: clientAddr, packet: packet}:
	141	+	case r.datagrams <- udpDatagram{clientAddr: key, packet: packet}:
112	142		case <-r.ctx.Done():
113	143		return
114	144		default:

@@ -118,6 +148,7 @@ func (r *udpEntryPointRuntime) readLoop() {
118	148		}
119	149
120	150		func (r *udpEntryPointRuntime) datagramWorker() {
	151	+	defer r.runWG.Done()
121	152		for {
122	153		select {
123	154		case <-r.ctx.Done():

@@ -133,9 +164,9 @@ func (r *udpEntryPointRuntime) datagramWorker() {
133	164		}
134	165		}
135	166
136		-	func (r *udpEntryPointRuntime) handleDatagram(clientAddr net.Addr, packet []byte) {
	167	+	func (r *udpEntryPointRuntime) handleDatagram(clientAddr netip.AddrPort, packet []byte) {
137	168		options := effectiveUDPOptions(snapshotUDPOptions(r.manager.snapshot.Load()))
138		-	session, ok := r.session(clientAddr.String(), clientAddr, options)
	169	+	session, ok := r.session(clientAddr, options)
139	170		if !ok {
140	171		return
141	172		}

@@ -143,16 +174,20 @@ func (r *udpEntryPointRuntime) handleDatagram(clientAddr net.Addr, packet []byte
143	174		r.counters.bytesIn.Add(int64(n))
144	175		if err != nil {
145	176		r.counters.totalErrors.Add(1)
146		-	r.removeSession(clientAddr.String())
	177	+	r.removeSession(session)
147	178		}
148	179		}
149	180
150		-	func (r udpEntryPointRuntime) session(key string, clientAddr net.Addr, options domain.UDPOptions) (udpSession, bool) {
	181	+	func (r udpEntryPointRuntime) session(clientAddr netip.AddrPort, options domain.UDPOptions) (udpSession, bool) {
151	182		r.mu.Lock()
152		-	if session := r.sessions[key]; session != nil {
153		-	session.touch()
	183	+	if existing := r.sessions[clientAddr]; existing != nil {
	184	+	existing.touch()
	185	+	r.mu.Unlock()
	186	+	return existing, true
	187	+	}
	188	+	if !r.accepting {
154	189		r.mu.Unlock()
155		-	return session, true
	190	+	return nil, false
156	191		}
157	192		r.mu.Unlock()
158	193

@@ -170,33 +205,60 @@ func (r *udpEntryPointRuntime) session(key string, clientAddr net.Addr, options
170	205		}
171	206		r.mu.Unlock()
172	207		dialCtx, cancel := context.WithTimeout(r.ctx, udpDialTimeout(options))
173		-	backendConn, err := (&net.Dialer{}).DialContext(dialCtx, "udp", net.JoinHostPort(backend.Host, strconv.Itoa(backend.Port)))
	208	+	backendConn, err := r.dialContext(dialCtx, "udp", net.JoinHostPort(backend.Host, strconv.Itoa(backend.Port)))
174	209		cancel()
175	210		if err != nil {
176	211		r.counters.totalErrors.Add(1)
177	212		return nil, false
178	213		}
179		-	session := &udpSession{clientAddr: clientAddr, backend: backendConn, backendRef: backend, done: make(chan struct{})}
	214	+	session := &udpSession{
	215	+	clientAddr: clientAddr,
	216	+	remoteAddr: net.UDPAddrFromAddrPort(clientAddr),
	217	+	backend: backendConn,
	218	+	backendRef: backend,
	219	+	}
180	220		session.touch()
	221	+	return r.publishSession(session, options)
	222	+	}
181	223
	224	+	// publishSession runs the locked post-dial admission, re-dedup and publication
	225	+	// step. It must only be called after the backend dial completes. It returns the
	226	+	// session callers should use:
	227	+	// - the newly dialed session once it is published,
	228	+	// - an existing session that won a concurrent creation race, with the newly
	229	+	// dialed backend closed and the existing session touched,
	230	+	// - nil when admission has closed or the session limit was reached, with the
	231	+	// newly dialed backend closed.
	232	+	//
	233	+	// The runWG count for a published session is added while holding mu and before
	234	+	// its goroutine starts, so it can never race stop()'s Wait (see the runWG
	235	+	// field comment). The backendLoop goroutine starts after the lock is released.
	236	+	func (r udpEntryPointRuntime) publishSession(session udpSession, options domain.UDPOptions) (*udpSession, bool) {
182	237		r.mu.Lock()
183		-	if existing := r.sessions[key]; existing != nil {
	238	+	if !r.accepting {
	239	+	r.mu.Unlock()
	240	+	_ = session.backend.Close()
	241	+	return nil, false
	242	+	}
	243	+	if existing := r.sessions[session.clientAddr]; existing != nil {
184	244		r.mu.Unlock()
185		-	_ = backendConn.Close()
	245	+	_ = session.backend.Close()
	246	+	existing.touch()
186	247		return existing, true
187	248		}
188	249		if options.MaxSessions > 0 && len(r.sessions) >= options.MaxSessions {
189	250		r.mu.Unlock()
190		-	_ = backendConn.Close()
	251	+	_ = session.backend.Close()
191	252		r.counters.totalRefused.Add(1)
192	253		return nil, false
193	254		}
194		-	r.sessions[key] = session
195		-	r.mu.Unlock()
196		-
	255	+	r.sessions[session.clientAddr] = session
	256	+	r.runWG.Add(1)
197	257		r.counters.activeUDPSessions.Add(1)
198	258		r.counters.totalAccepted.Add(1)
199		-	go r.backendLoop(key, session)
	259	+	r.mu.Unlock()
	260	+
	261	+	go r.backendLoop(session)
200	262		return session, true
201	263		}
202	264

@@ -211,8 +273,9 @@ func udpDialTimeout(options domain.UDPOptions) time.Duration {
211	273		return 5 * time.Second
212	274		}
213	275
214		-	func (r udpEntryPointRuntime) backendLoop(key string, session udpSession) {
215		-	defer r.removeSession(key)
	276	+	func (r udpEntryPointRuntime) backendLoop(session udpSession) {
	277	+	defer r.runWG.Done()
	278	+	defer r.removeSession(session)
216	279		buf := make([]byte, udpBufferSize)
217	280		for {
218	281		n, err := session.backend.Read(buf)

@@ -222,7 +285,7 @@ func (r udpEntryPointRuntime) backendLoop(key string, session udpSession) {
222	285		}
223	286		return
224	287		}
225		-	written, err := r.packetConn.WriteTo(buf[:n], session.clientAddr)
	288	+	written, err := r.packetConn.WriteTo(buf[:n], session.remoteAddr)
226	289		r.counters.bytesOut.Add(int64(written))
227	290		if err != nil {
228	291		if !r.isClosed() {

@@ -234,6 +297,7 @@ func (r udpEntryPointRuntime) backendLoop(key string, session udpSession) {
234	297		}
235	298
236	299		func (r *udpEntryPointRuntime) expireLoop() {
	300	+	defer r.runWG.Done()
237	301		ticker := time.NewTicker(25 * time.Millisecond)
238	302		defer ticker.Stop()
239	303		for {

@@ -266,14 +330,22 @@ func (r *udpEntryPointRuntime) expireIdleSessions(idleTimeout time.Duration) {
266	330		}
267	331		}
268	332
269		-	func (r *udpEntryPointRuntime) removeSession(key string) {
	333	+	// removeSession drops session from the session map and closes its backend only
	334	+	// when it is still the session currently registered for its client address. A
	335	+	// stale remover (expiry, backend read/write failure, reload or CIDR cleanup) can
	336	+	// therefore never evict or close a replacement published under the same key.
	337	+	// The backend socket is closed outside the lock.
	338	+	func (r udpEntryPointRuntime) removeSession(session udpSession) {
	339	+	if session == nil {
	340	+	return
	341	+	}
270	342		r.mu.Lock()
271		-	session := r.sessions[key]
272		-	if session != nil {
273		-	delete(r.sessions, key)
	343	+	current := r.sessions[session.clientAddr]
	344	+	if current == session {
	345	+	delete(r.sessions, session.clientAddr)
274	346		}
275	347		r.mu.Unlock()
276		-	if session == nil {
	348	+	if current != session {
277	349		return
278	350		}
279	351		session.close()

@@ -310,29 +382,26 @@ func (r *udpEntryPointRuntime) stop(ctx context.Context, drainTimeout time.Durat
310	382		if r.closed.CompareAndSwap(false, true) {
311	383		entryPoint := r.entryPointSnapshot()
312	384		trafficInfo(ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Msg("stopping udp traffic entrypoint")
	385	+	// Closing admission under mu serializes with start() and session(),
	386	+	// which are the only places that add to runWG. Once this critical
	387	+	// section completes no further Add can occur, so waitRuntime may Wait
	388	+	// safely.
	389	+	r.mu.Lock()
	390	+	r.accepting = false
	391	+	r.mu.Unlock()
313	392		r.cancel()
314	393		_ = r.packetConn.Close()
315		-	if !r.started.Load() {
316		-	r.closeDone()
317		-	}
318	394		}
319		-	select {
320		-	case <-r.done:
321		-	case <-ctx.Done():
322		-	r.closeSessions()
323		-	return
324		-	}
325		-	if r.waitSessions(ctx, drainTimeout) {
	395	+	if r.waitRuntime(ctx, drainTimeout) {
326	396		trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
327	397		return
328	398		}
329	399		trafficDebug(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Dur("drain_timeout", drainTimeout).Msg("forcing udp traffic entrypoint drain")
330	400		r.closeSessions()
331		-	select {
332		-	case <-r.sessionsDone():
333		-	trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
334		-	case <-ctx.Done():
	401	+	if !r.waitRuntime(ctx, drainTimeout) {
	402	+	r.runWG.Wait()
335	403		}
	404	+	trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
336	405		}
337	406
338	407		func (r *udpEntryPointRuntime) drainSessionsAfter(drainTimeout time.Duration) {

@@ -358,56 +427,40 @@ func (r udpEntryPointRuntime) drainSessionsMatchingAfter(match func(udpSession
358	427		}()
359	428		}
360	429
361		-	func (r *udpEntryPointRuntime) waitSessions(ctx context.Context, drainTimeout time.Duration) bool {
	430	+	func (r *udpEntryPointRuntime) waitRuntime(ctx context.Context, drainTimeout time.Duration) bool {
362	431		if drainTimeout <= 0 {
363	432		drainTimeout = defaultUDPOptions().DrainTimeout
364	433		}
365	434		drainCtx, cancel := context.WithTimeout(ctx, drainTimeout)
366	435		defer cancel()
	436	+	done := make(chan struct{})
	437	+	go func() {
	438	+	r.runWG.Wait()
	439	+	close(done)
	440	+	}()
367	441		select {
368		-	case <-r.sessionsDone():
	442	+	case <-done:
369	443		return true
370	444		case <-drainCtx.Done():
371	445		return false
372	446		}
373	447		}
374	448
375		-	func (r *udpEntryPointRuntime) closeDone() {
376		-	r.doneOnce.Do(func() { close(r.done) })
377		-	}
378		-
379		-	func (r *udpEntryPointRuntime) sessionsDone() <-chan struct{} {
380		-	done := make(chan struct{})
381		-	go func() {
382		-	for {
383		-	r.mu.Lock()
384		-	count := len(r.sessions)
385		-	r.mu.Unlock()
386		-	if count == 0 {
387		-	close(done)
388		-	return
389		-	}
390		-	time.Sleep(10 * time.Millisecond)
391		-	}
392		-	}()
393		-	return done
394		-	}
395		-
396	449		func (r *udpEntryPointRuntime) closeSessions() {
397	450		r.closeSessionsMatching(func(*udpSession) bool { return true })
398	451		}
399	452
400	453		func (r udpEntryPointRuntime) closeSessionsMatching(match func(udpSession) bool) {
401	454		r.mu.Lock()
402		-	keys := make([]string, 0, len(r.sessions))
403		-	for key, session := range r.sessions {
	455	+	sessions := make([]*udpSession, 0, len(r.sessions))
	456	+	for _, session := range r.sessions {
404	457		if match(session) {
405		-	keys = append(keys, key)
	458	+	sessions = append(sessions, session)
406	459		}
407	460		}
408	461		r.mu.Unlock()
409		-	for _, key := range keys {
410		-	r.removeSession(key)
	462	+	for _, session := range sessions {
	463	+	r.removeSession(session)
411	464		}
412	465		}
413	466

@@ -424,15 +477,15 @@ func (r *udpEntryPointRuntime) updateEntryPoint(entryPoint domain.EntryPoint, tr
424	477		r.mu.Lock()
425	478		r.entryPoint = entryPoint
426	479		r.trusted = trusted
427		-	stale := make([]string, 0)
428		-	for key, session := range r.sessions {
429		-	if !trustedRemoteAddr(trusted, session.clientAddr) {
430		-	stale = append(stale, key)
	480	+	stale := make([]*udpSession, 0)
	481	+	for _, session := range r.sessions {
	482	+	if !trustedAddrPort(trusted, session.clientAddr) {
	483	+	stale = append(stale, session)
431	484		}
432	485		}
433	486		r.mu.Unlock()
434		-	for _, key := range stale {
435		-	r.removeSession(key)
	487	+	for _, session := range stale {
	488	+	r.removeSession(session)
436	489		}
437	490		}
438	491

@@ -450,11 +503,59 @@ func (r udpEntryPointRuntime) trustedSnapshot() []net.IPNet {
450	503
451	504		func (r *udpEntryPointRuntime) isClosed() bool { return r.closed.Load() }
452	505
	506	+	// udpAddrPort converts a packet source address into a typed, comparable session
	507	+	// key. IPv4-mapped addresses are normalized to their IPv4 form so that both
	508	+	// representations share a single session; IPv6 zones are preserved. Invalid
	509	+	// addresses are rejected before any trusted-list handling so they can never be
	510	+	// admitted as a zero-valued session key.
	511	+	func udpAddrPort(addr net.Addr) (netip.AddrPort, bool) {
	512	+	if udpAddr, ok := addr.(*net.UDPAddr); ok {
	513	+	key := normalizeUDPAddrPort(udpAddr.AddrPort())
	514	+	return key, key.IsValid()
	515	+	}
	516	+	parsed, err := netip.ParseAddrPort(addr.String())
	517	+	if err != nil {
	518	+	return netip.AddrPort{}, false
	519	+	}
	520	+	key := normalizeUDPAddrPort(parsed)
	521	+	return key, key.IsValid()
	522	+	}
	523	+
	524	+	func normalizeUDPAddrPort(addrPort netip.AddrPort) netip.AddrPort {
	525	+	if !addrPort.IsValid() {
	526	+	return addrPort
	527	+	}
	528	+	return netip.AddrPortFrom(addrPort.Addr().Unmap(), addrPort.Port())
	529	+	}
	530	+
	531	+	// trustedAddrPort reports whether addr is within the trusted CIDRs. An empty
	532	+	// trusted list admits every client.
	533	+	//
	534	+	// Matching is intentionally by IP only: the IPv6 zone identifies the local
	535	+	// interface a packet arrived on, not the peer, so a zoned link-local client
	536	+	// matches the same CIDR as its zoneless form. Zones remain part of the session
	537	+	// key (see udpAddrPort) and are dropped only for trust evaluation. The address
	538	+	// is unmapped before matching so a v4-mapped client matches an IPv4 CIDR.
	539	+	func trustedAddrPort(trusted []*net.IPNet, addr netip.AddrPort) bool {
	540	+	if !addr.IsValid() {
	541	+	return false
	542	+	}
	543	+	if len(trusted) == 0 {
	544	+	return true
	545	+	}
	546	+	ip := addr.Addr().Unmap().AsSlice()
	547	+	for _, network := range trusted {
	548	+	if network.Contains(ip) {
	549	+	return true
	550	+	}
	551	+	}
	552	+	return false
	553	+	}
	554	+
453	555		func (s *udpSession) touch() { s.lastSeen.Store(time.Now().UnixNano()) }
454	556
455	557		func (s *udpSession) close() {
456	558		s.once.Do(func() {
457	559		_ = s.backend.Close()
458		-	close(s.done)
459	560		})
460	561		}

 	"context"
 	"errors"
 	"net"
+	"net/netip"
 	"strconv"
 	"sync"
 	"sync/atomic"
 	started atomic.Bool
 	closed  atomic.Bool
-	ctx      context.Context
+	ctx    context.Context
-	cancel   context.CancelFunc
+	cancel context.CancelFunc
-	done     chan struct{}
+	// runWG tracks every goroutine owned by this runtime (read loop, datagram
-	doneOnce sync.Once
+	// workers, expiry loop and per-session backend loops). Counts are only added
+	// while holding mu and before the owning goroutine starts; stop() flips
+	// accepting under the same lock before it ever waits, so no Add can race a
+	// Wait.
+	runWG sync.WaitGroup
 	datagrams chan udpDatagram
-	mu       sync.Mutex
+	// dialContext dials the upstream backend. It is a private instance field so
-	sessions map[string]*udpSession
+	// tests can drive in-flight dial interleavings deterministically; production
+	// code always uses the standard UDP dialer set at construction and never
+	// mutates it after the runtime starts.
+	dialContext func(ctx context.Context, network, address string) (net.Conn, error)
+	mu        sync.Mutex
+	accepting bool
+	sessions  map[netip.AddrPort]*udpSession
 }
 type udpDatagram struct {
-	clientAddr net.Addr
+	clientAddr netip.AddrPort
 	packet     []byte
 }
 type udpSession struct {
-	clientAddr net.Addr
+	clientAddr netip.AddrPort
+	// remoteAddr is the immutable packet destination precomputed once at
+	// publication, so backendLoop never has to allocate or format an address
+	// per packet. It is shared read-only across the session's goroutines.
+	remoteAddr *net.UDPAddr
 	backend    net.Conn
 	backendRef domain.TrafficBackend
 	lastSeen   atomic.Int64
-	done       chan struct{}
 	once       sync.Once
 }
 		trusted:    trusted,
 		ctx:        ctx,
 		cancel:     cancel,
-		done:       make(chan struct{}),
+		accepting:  true,
 		datagrams:  make(chan udpDatagram, udpDatagramBacklog),
-		sessions:   map[string]*udpSession{},
+		sessions:   map[netip.AddrPort]*udpSession{},
+		dialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
+			return (&net.Dialer{}).DialContext(ctx, network, address)
+		},
 	}
 }
 	if !r.started.CompareAndSwap(false, true) {
 		return
 	}
+	r.mu.Lock()
+	if !r.accepting {
+		r.mu.Unlock()
+		return
+	}
+	r.runWG.Add(udpWorkerCount + 2) // read loop, expiry loop and datagram workers
+	r.mu.Unlock()
 	entryPoint := r.entryPointSnapshot()
 	trafficInfo(r.ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Str("protocol", string(entryPoint.Protocol)).Msg("started udp traffic entrypoint")
 	for range udpWorkerCount {
 }
 func (r *udpEntryPointRuntime) readLoop() {
-	defer r.closeDone()
+	defer r.runWG.Done()
 	buf := make([]byte, udpBufferSize)
 	for {
 		n, clientAddr, err := r.packetConn.ReadFrom(buf)
 			r.counters.totalErrors.Add(1)
 			continue
 		}
-		if !trustedRemoteAddr(r.trustedSnapshot(), clientAddr) {
+		key, ok := udpAddrPort(clientAddr)
+		if !ok {
+			r.counters.totalRefused.Add(1)
+			continue
+		}
+		if !trustedAddrPort(r.trustedSnapshot(), key) {
 			r.counters.totalRefused.Add(1)
 			continue
 		}
 		}
 		packet := append([]byte(nil), buf[:n]...)
 		select {
-		case r.datagrams <- udpDatagram{clientAddr: clientAddr, packet: packet}:
+		case r.datagrams <- udpDatagram{clientAddr: key, packet: packet}:
 		case <-r.ctx.Done():
 			return
 		default:
 }
 func (r *udpEntryPointRuntime) datagramWorker() {
+	defer r.runWG.Done()
 	for {
 		select {
 		case <-r.ctx.Done():
 	}
 }
-func (r *udpEntryPointRuntime) handleDatagram(clientAddr net.Addr, packet []byte) {
+func (r *udpEntryPointRuntime) handleDatagram(clientAddr netip.AddrPort, packet []byte) {
 	options := effectiveUDPOptions(snapshotUDPOptions(r.manager.snapshot.Load()))
-	session, ok := r.session(clientAddr.String(), clientAddr, options)
+	session, ok := r.session(clientAddr, options)
 	if !ok {
 		return
 	}
 	r.counters.bytesIn.Add(int64(n))
 	if err != nil {
 		r.counters.totalErrors.Add(1)
-		r.removeSession(clientAddr.String())
+		r.removeSession(session)
 	}
 }
-func (r *udpEntryPointRuntime) session(key string, clientAddr net.Addr, options domain.UDPOptions) (*udpSession, bool) {
+func (r *udpEntryPointRuntime) session(clientAddr netip.AddrPort, options domain.UDPOptions) (*udpSession, bool) {
 	r.mu.Lock()
-	if session := r.sessions[key]; session != nil {
+	if existing := r.sessions[clientAddr]; existing != nil {
-		session.touch()
+		existing.touch()
+		r.mu.Unlock()
+		return existing, true
+	}
+	if !r.accepting {
 		r.mu.Unlock()
-		return session, true
+		return nil, false
 	}
 	r.mu.Unlock()
 	}
 	r.mu.Unlock()
 	dialCtx, cancel := context.WithTimeout(r.ctx, udpDialTimeout(options))
-	backendConn, err := (&net.Dialer{}).DialContext(dialCtx, "udp", net.JoinHostPort(backend.Host, strconv.Itoa(backend.Port)))
+	backendConn, err := r.dialContext(dialCtx, "udp", net.JoinHostPort(backend.Host, strconv.Itoa(backend.Port)))
 	cancel()
 	if err != nil {
 		r.counters.totalErrors.Add(1)
 		return nil, false
 	}
-	session := &udpSession{clientAddr: clientAddr, backend: backendConn, backendRef: backend, done: make(chan struct{})}
+	session := &udpSession{
+		clientAddr: clientAddr,
+		remoteAddr: net.UDPAddrFromAddrPort(clientAddr),
+		backend:    backendConn,
+		backendRef: backend,
+	}
 	session.touch()
+	return r.publishSession(session, options)
+}
+// publishSession runs the locked post-dial admission, re-dedup and publication
+// step. It must only be called after the backend dial completes. It returns the
+// session callers should use:
+//   - the newly dialed session once it is published,
+//   - an existing session that won a concurrent creation race, with the newly
+//     dialed backend closed and the existing session touched,
+//   - nil when admission has closed or the session limit was reached, with the
+//     newly dialed backend closed.
+//
+// The runWG count for a published session is added while holding mu and before
+// its goroutine starts, so it can never race stop()'s Wait (see the runWG
+// field comment). The backendLoop goroutine starts after the lock is released.
+func (r *udpEntryPointRuntime) publishSession(session *udpSession, options domain.UDPOptions) (*udpSession, bool) {
 	r.mu.Lock()
-	if existing := r.sessions[key]; existing != nil {
+	if !r.accepting {
+		r.mu.Unlock()
+		_ = session.backend.Close()
+		return nil, false
+	}
+	if existing := r.sessions[session.clientAddr]; existing != nil {
 		r.mu.Unlock()
-		_ = backendConn.Close()
+		_ = session.backend.Close()
+		existing.touch()
 		return existing, true
 	}
 	if options.MaxSessions > 0 && len(r.sessions) >= options.MaxSessions {
 		r.mu.Unlock()
-		_ = backendConn.Close()
+		_ = session.backend.Close()
 		r.counters.totalRefused.Add(1)
 		return nil, false
 	}
-	r.sessions[key] = session
+	r.sessions[session.clientAddr] = session
-	r.mu.Unlock()
+	r.runWG.Add(1)
 	r.counters.activeUDPSessions.Add(1)
 	r.counters.totalAccepted.Add(1)
-	go r.backendLoop(key, session)
+	r.mu.Unlock()
+	go r.backendLoop(session)
 	return session, true
 }
 	return 5 * time.Second
 }
-func (r *udpEntryPointRuntime) backendLoop(key string, session *udpSession) {
+func (r *udpEntryPointRuntime) backendLoop(session *udpSession) {
-	defer r.removeSession(key)
+	defer r.runWG.Done()
+	defer r.removeSession(session)
 	buf := make([]byte, udpBufferSize)
 	for {
 		n, err := session.backend.Read(buf)
 			}
 			return
 		}
-		written, err := r.packetConn.WriteTo(buf[:n], session.clientAddr)
+		written, err := r.packetConn.WriteTo(buf[:n], session.remoteAddr)
 		r.counters.bytesOut.Add(int64(written))
 		if err != nil {
 			if !r.isClosed() {
 }
 func (r *udpEntryPointRuntime) expireLoop() {
+	defer r.runWG.Done()
 	ticker := time.NewTicker(25 * time.Millisecond)
 	defer ticker.Stop()
 	for {
 	}
 }
-func (r *udpEntryPointRuntime) removeSession(key string) {
+// removeSession drops session from the session map and closes its backend only
+// when it is still the session currently registered for its client address. A
+// stale remover (expiry, backend read/write failure, reload or CIDR cleanup) can
+// therefore never evict or close a replacement published under the same key.
+// The backend socket is closed outside the lock.
+func (r *udpEntryPointRuntime) removeSession(session *udpSession) {
+	if session == nil {
+		return
+	}
 	r.mu.Lock()
-	session := r.sessions[key]
+	current := r.sessions[session.clientAddr]
-	if session != nil {
+	if current == session {
-		delete(r.sessions, key)
+		delete(r.sessions, session.clientAddr)
 	}
 	r.mu.Unlock()
-	if session == nil {
+	if current != session {
 		return
 	}
 	session.close()
 	if r.closed.CompareAndSwap(false, true) {
 		entryPoint := r.entryPointSnapshot()
 		trafficInfo(ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Msg("stopping udp traffic entrypoint")
+		// Closing admission under mu serializes with start() and session(),
+		// which are the only places that add to runWG. Once this critical
+		// section completes no further Add can occur, so waitRuntime may Wait
+		// safely.
+		r.mu.Lock()
+		r.accepting = false
+		r.mu.Unlock()
 		r.cancel()
 		_ = r.packetConn.Close()
-		if !r.started.Load() {
-			r.closeDone()
-		}
 	}
-	select {
+	if r.waitRuntime(ctx, drainTimeout) {
-	case <-r.done:
-	case <-ctx.Done():
-		r.closeSessions()
-		return
-	}
-	if r.waitSessions(ctx, drainTimeout) {
 		trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
 		return
 	}
 	trafficDebug(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Dur("drain_timeout", drainTimeout).Msg("forcing udp traffic entrypoint drain")
 	r.closeSessions()
-	select {
+	if !r.waitRuntime(ctx, drainTimeout) {
-	case <-r.sessionsDone():
+		r.runWG.Wait()
-		trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
-	case <-ctx.Done():
 	}
+	trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
 }
 func (r *udpEntryPointRuntime) drainSessionsAfter(drainTimeout time.Duration) {
 	}()
 }
-func (r *udpEntryPointRuntime) waitSessions(ctx context.Context, drainTimeout time.Duration) bool {
+func (r *udpEntryPointRuntime) waitRuntime(ctx context.Context, drainTimeout time.Duration) bool {
 	if drainTimeout <= 0 {
 		drainTimeout = defaultUDPOptions().DrainTimeout
 	}
 	drainCtx, cancel := context.WithTimeout(ctx, drainTimeout)
 	defer cancel()
+	done := make(chan struct{})
+	go func() {
+		r.runWG.Wait()
+		close(done)
+	}()
 	select {
-	case <-r.sessionsDone():
+	case <-done:
 		return true
 	case <-drainCtx.Done():
 		return false
 	}
 }
-func (r *udpEntryPointRuntime) closeDone() {
-	r.doneOnce.Do(func() { close(r.done) })
-}
-func (r *udpEntryPointRuntime) sessionsDone() <-chan struct{} {
-	done := make(chan struct{})
-	go func() {
-		for {
-			r.mu.Lock()
-			count := len(r.sessions)
-			r.mu.Unlock()
-			if count == 0 {
-				close(done)
-				return
-			}
-			time.Sleep(10 * time.Millisecond)
-		}
-	}()
-	return done
-}
 func (r *udpEntryPointRuntime) closeSessions() {
 	r.closeSessionsMatching(func(*udpSession) bool { return true })
 }
 func (r *udpEntryPointRuntime) closeSessionsMatching(match func(*udpSession) bool) {
 	r.mu.Lock()
-	keys := make([]string, 0, len(r.sessions))
+	sessions := make([]*udpSession, 0, len(r.sessions))
-	for key, session := range r.sessions {
+	for _, session := range r.sessions {
 		if match(session) {
-			keys = append(keys, key)
+			sessions = append(sessions, session)
 		}
 	}
 	r.mu.Unlock()
-	for _, key := range keys {
+	for _, session := range sessions {
-		r.removeSession(key)
+		r.removeSession(session)
 	}
 }
 	r.mu.Lock()
 	r.entryPoint = entryPoint
 	r.trusted = trusted
-	stale := make([]string, 0)
+	stale := make([]*udpSession, 0)
-	for key, session := range r.sessions {
+	for _, session := range r.sessions {
-		if !trustedRemoteAddr(trusted, session.clientAddr) {
+		if !trustedAddrPort(trusted, session.clientAddr) {
-			stale = append(stale, key)
+			stale = append(stale, session)
 		}
 	}
 	r.mu.Unlock()
-	for _, key := range stale {
+	for _, session := range stale {
-		r.removeSession(key)
+		r.removeSession(session)
 	}
 }
 func (r *udpEntryPointRuntime) isClosed() bool { return r.closed.Load() }
+// udpAddrPort converts a packet source address into a typed, comparable session
+// key. IPv4-mapped addresses are normalized to their IPv4 form so that both
+// representations share a single session; IPv6 zones are preserved. Invalid
+// addresses are rejected before any trusted-list handling so they can never be
+// admitted as a zero-valued session key.
+func udpAddrPort(addr net.Addr) (netip.AddrPort, bool) {
+	if udpAddr, ok := addr.(*net.UDPAddr); ok {
+		key := normalizeUDPAddrPort(udpAddr.AddrPort())
+		return key, key.IsValid()
+	}
+	parsed, err := netip.ParseAddrPort(addr.String())
+	if err != nil {
+		return netip.AddrPort{}, false
+	}
+	key := normalizeUDPAddrPort(parsed)
+	return key, key.IsValid()
+}
+func normalizeUDPAddrPort(addrPort netip.AddrPort) netip.AddrPort {
+	if !addrPort.IsValid() {
+		return addrPort
+	}
+	return netip.AddrPortFrom(addrPort.Addr().Unmap(), addrPort.Port())
+}
+// trustedAddrPort reports whether addr is within the trusted CIDRs. An empty
+// trusted list admits every client.
+//
+// Matching is intentionally by IP only: the IPv6 zone identifies the local
+// interface a packet arrived on, not the peer, so a zoned link-local client
+// matches the same CIDR as its zoneless form. Zones remain part of the session
+// key (see udpAddrPort) and are dropped only for trust evaluation. The address
+// is unmapped before matching so a v4-mapped client matches an IPv4 CIDR.
+func trustedAddrPort(trusted []*net.IPNet, addr netip.AddrPort) bool {
+	if !addr.IsValid() {
+		return false
+	}
+	if len(trusted) == 0 {
+		return true
+	}
+	ip := addr.Addr().Unmap().AsSlice()
+	for _, network := range trusted {
+		if network.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
 func (s *udpSession) touch() { s.lastSeen.Store(time.Now().UnixNano()) }
 func (s *udpSession) close() {
 	s.once.Do(func() {
 		_ = s.backend.Close()
-		close(s.done)
 	})
 }

M internal/adapters/in/traffic/udp_test.go +511 -0


@@ -3,6 +3,8 @@ package traffic
3	3		import (
4	4		"context"
5	5		"net"
	6	+	"net/netip"
	7	+	"sync"
6	8		"testing"
7	9		"time"
8	10

@@ -106,6 +108,391 @@ func TestUDPReloadReplacesSessionWhenBackendChanges(t *testing.T) {
106	108		assertUDPRoundTripWant(t, newConn, "third", "b:third")
107	109		}
108	110
	111	+	func TestUDPStaleRemovalKeepsReplacementSession(t *testing.T) {
	112	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4100")
	113	+	stale := newTestUDPSession(t, clientAddr)
	114	+	replacement := newTestUDPSession(t, clientAddr)
	115	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
	116	+	runtime.mu.Lock()
	117	+	runtime.sessions[clientAddr] = replacement
	118	+	runtime.counters.activeUDPSessions.Store(1)
	119	+	runtime.mu.Unlock()
	120	+
	121	+	// A stale remover carries the pointer of the session that used to own the
	122	+	// client key. Expiry, backend read/write failures, reload and CIDR cleanup
	123	+	// all funnel through removeSession, so none of them may evict the session
	124	+	// that replaced it.
	125	+	runtime.removeSession(stale)
	126	+
	127	+	runtime.mu.Lock()
	128	+	current, ok := runtime.sessions[clientAddr]
	129	+	runtime.mu.Unlock()
	130	+	require.True(t, ok, "replacement session must stay registered")
	131	+	require.Same(t, replacement, current)
	132	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load(), "stale removal must not decrement the active counter")
	133	+	require.False(t, testUDPSessionClosed(replacement), "replacement backend must stay open")
	134	+
	135	+	runtime.removeSession(replacement)
	136	+
	137	+	runtime.mu.Lock()
	138	+	_, ok = runtime.sessions[clientAddr]
	139	+	runtime.mu.Unlock()
	140	+	require.False(t, ok, "removing the current session must unregister it")
	141	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load(), "removing the current session must decrement once")
	142	+	require.True(t, testUDPSessionClosed(replacement))
	143	+
	144	+	// Removing again must be idempotent: no double decrement and no panic.
	145	+	runtime.removeSession(replacement)
	146	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	147	+	}
	148	+
	149	+	func TestUDPExpiryClosesOnlyRegisteredIdleSessionOnce(t *testing.T) {
	150	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4200")
	151	+	// stale was replaced under the same key and is no longer registered.
	152	+	stale := newTestUDPSession(t, clientAddr)
	153	+	stale.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
	154	+	current := newTestUDPSession(t, clientAddr)
	155	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
	156	+	runtime.mu.Lock()
	157	+	runtime.sessions[clientAddr] = current
	158	+	runtime.counters.activeUDPSessions.Store(1)
	159	+	runtime.mu.Unlock()
	160	+
	161	+	runtime.expireIdleSessions(time.Minute)
	162	+	require.False(t, testUDPSessionClosed(current), "fresh registered session must not expire")
	163	+	require.False(t, testUDPSessionClosed(stale), "an unregistered session must never be closed by expiry")
	164	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load())
	165	+
	166	+	current.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
	167	+	runtime.expireIdleSessions(time.Minute)
	168	+
	169	+	runtime.mu.Lock()
	170	+	_, ok := runtime.sessions[clientAddr]
	171	+	runtime.mu.Unlock()
	172	+	require.False(t, ok)
	173	+	require.True(t, testUDPSessionClosed(current))
	174	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load(), "expiry must decrement exactly once")
	175	+
	176	+	// The backend loop of the expired session also removes it on return; the
	177	+	// identity check keeps the counter coherent and must not double decrement.
	178	+	runtime.removeSession(current)
	179	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	180	+	}
	181	+
	182	+	func TestUDPAdmissionGateBlocksPublicationAfterClosingBegins(t *testing.T) {
	183	+	backend := startUDPEchoServer(t)
	184	+	runtime, _ := newTestUDPRuntime(t, backend.address)
	185	+	// Model the instant stop() closes admission but before it cancels the
	186	+	// runtime context: a datagram worker that reaches session() now must not
	187	+	// publish a new session.
	188	+	runtime.mu.Lock()
	189	+	runtime.accepting = false
	190	+	runtime.mu.Unlock()
	191	+	runtime.closed.Store(true)
	192	+
	193	+	session, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4300"), effectiveUDPOptions(domain.UDPOptions{}))
	194	+	require.False(t, ok)
	195	+	require.Nil(t, session)
	196	+
	197	+	runtime.mu.Lock()
	198	+	require.Empty(t, runtime.sessions)
	199	+	runtime.mu.Unlock()
	200	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	201	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	202	+	}
	203	+
	204	+	func TestUDPPublishSessionRejectsCompletedDialAfterAdmissionClosure(t *testing.T) {
	205	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
	206	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4800")
	207	+	dialed := newTestUDPSession(t, clientAddr)
	208	+	runtime.mu.Lock()
	209	+	runtime.accepting = false
	210	+	runtime.mu.Unlock()
	211	+	runtime.closed.Store(true)
	212	+
	213	+	session, ok := runtime.publishSession(dialed, effectiveUDPOptions(domain.UDPOptions{}))
	214	+	require.False(t, ok)
	215	+	require.Nil(t, session)
	216	+	require.True(t, testUDPSessionClosed(dialed), "a rejected dial must have its backend closed")
	217	+	runtime.mu.Lock()
	218	+	require.Empty(t, runtime.sessions)
	219	+	runtime.mu.Unlock()
	220	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	221	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load())
	222	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	223	+	}
	224	+
	225	+	func TestUDPPublishSessionReusesExistingAndTouchesIt(t *testing.T) {
	226	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
	227	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4900")
	228	+	existing := newTestUDPSession(t, clientAddr)
	229	+	existing.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
	230	+	dialed := newTestUDPSession(t, clientAddr)
	231	+	runtime.mu.Lock()
	232	+	runtime.sessions[clientAddr] = existing
	233	+	runtime.counters.activeUDPSessions.Store(1)
	234	+	runtime.mu.Unlock()
	235	+
	236	+	session, ok := runtime.publishSession(dialed, effectiveUDPOptions(domain.UDPOptions{}))
	237	+	require.True(t, ok)
	238	+	require.Same(t, existing, session)
	239	+	require.True(t, testUDPSessionClosed(dialed), "the losing dial must have its backend closed")
	240	+	require.False(t, testUDPSessionClosed(existing))
	241	+	require.Greater(t, existing.lastSeen.Load(), time.Now().Add(-time.Minute).UnixNano(), "reuse must touch the existing session")
	242	+	runtime.mu.Lock()
	243	+	current := runtime.sessions[clientAddr]
	244	+	runtime.mu.Unlock()
	245	+	require.Same(t, existing, current)
	246	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load(), "reuse must not add a session")
	247	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load(), "reuse must not count as a new accept")
	248	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	249	+	}
	250	+
	251	+	func TestUDPRemoveSessionNilIsSafe(t *testing.T) {
	252	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
	253	+	require.NotPanics(t, func() { runtime.removeSession(nil) })
	254	+	runtime.mu.Lock()
	255	+	require.Empty(t, runtime.sessions)
	256	+	runtime.mu.Unlock()
	257	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	258	+	}
	259	+
	260	+	func TestUDPReadLoopRejectsInvalidSourceAddress(t *testing.T) {
	261	+	backend := startUDPEchoServer(t)
	262	+	packetConn := newScriptedPacketConn()
	263	+	runtime, _ := newTestUDPRuntimeWithPacketConn(t, backend.address, packetConn)
	264	+	runtime.start()
	265	+
	266	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("zero"), addr: &net.UDPAddr{}}
	267	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("junk"), addr: testUDPNetAddr{network: "udp", address: "not-an-address"}}
	268	+	require.Eventually(t, func() bool { return runtime.counters.totalRefused.Load() == 2 }, time.Second, 5*time.Millisecond)
	269	+	runtime.mu.Lock()
	270	+	require.Empty(t, runtime.sessions, "invalid source addresses must never open a session")
	271	+	runtime.mu.Unlock()
	272	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	273	+
	274	+	// A valid datagram after the invalid ones proves the read loop keeps serving.
	275	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("ok"), addr: &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 43210}}
	276	+	require.Eventually(t, func() bool { return runtime.counters.activeUDPSessions.Load() == 1 }, time.Second, 5*time.Millisecond)
	277	+	require.Equal(t, int64(1), runtime.counters.totalAccepted.Load())
	278	+	}
	279	+
	280	+	func TestUDPStopBeforeStartLeavesNoRuntimeGoroutines(t *testing.T) {
	281	+	backend := startUDPEchoServer(t)
	282	+	runtime, _ := newTestUDPRuntime(t, backend.address)
	283	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	284	+	defer cancel()
	285	+
	286	+	runtime.stop(ctx, 50*time.Millisecond)
	287	+	// Starting a runtime whose admission already closed must not launch workers.
	288	+	runtime.start()
	289	+
	290	+	runtime.mu.Lock()
	291	+	require.Empty(t, runtime.sessions)
	292	+	runtime.mu.Unlock()
	293	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	294	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	295	+	}
	296	+
	297	+	func TestUDPShutdownLeavesNoSessionsOrRuntimeGoroutines(t *testing.T) {
	298	+	backend := startUDPEchoServer(t)
	299	+	graph := udpGraph(t, freeUDPAddress(t), backend.address)
	300	+	graph.Options.UDP.DrainTimeout = 50 * time.Millisecond
	301	+	manager := NewManager()
	302	+	require.NoError(t, manager.Apply(context.Background(), &graph))
	303	+
	304	+	conn := dialUDP(t, graph.EntryPoints[0].Address)
	305	+	defer conn.Close()
	306	+	assertUDPRoundTrip(t, conn, "held")
	307	+	require.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 1 }, time.Second, 10*time.Millisecond)
	308	+
	309	+	manager.mu.Lock()
	310	+	runtime := manager.udpListeners["udp"]
	311	+	manager.mu.Unlock()
	312	+	require.NotNil(t, runtime)
	313	+
	314	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	315	+	defer cancel()
	316	+	require.NoError(t, manager.Shutdown(ctx))
	317	+
	318	+	runtime.mu.Lock()
	319	+	require.Empty(t, runtime.sessions)
	320	+	runtime.mu.Unlock()
	321	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	322	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	323	+	}
	324	+
	325	+	func TestUDPConcurrentCreationDuringShutdownLeavesNoSessions(t *testing.T) {
	326	+	for range 5 {
	327	+	backend := startUDPEchoServer(t)
	328	+	runtime, _ := newTestUDPRuntime(t, backend.address)
	329	+	runtime.start()
	330	+
	331	+	var creators sync.WaitGroup
	332	+	for worker := range 8 {
	333	+	creators.Add(1)
	334	+	go func(worker int) {
	335	+	defer creators.Done()
	336	+	for attempt := range 25 {
	337	+	key := netip.AddrPortFrom(netip.AddrFrom4([4]byte{127, 0, 0, 1}), uint16(5000+worker*100+attempt))
	338	+	runtime.session(key, effectiveUDPOptions(domain.UDPOptions{}))
	339	+	}
	340	+	}(worker)
	341	+	}
	342	+
	343	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	344	+	runtime.stop(ctx, 50*time.Millisecond)
	345	+	cancel()
	346	+	creators.Wait()
	347	+
	348	+	runtime.mu.Lock()
	349	+	require.Empty(t, runtime.sessions, "no session may be published after closing began")
	350	+	runtime.mu.Unlock()
	351	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	352	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	353	+	}
	354	+	}
	355	+
	356	+	// TestUDPInFlightDialDoesNotPublishAfterAdmissionClosure blocks the backend dial
	357	+	// on the runtime's private dial seam, closes admission exactly as stop() does,
	358	+	// then lets the dial complete: session() must reject the completed dial without
	359	+	// publishing a session or a backend goroutine. The seam is instance-local, so no
	360	+	// package-global state is touched.
	361	+	func TestUDPInFlightDialDoesNotPublishAfterAdmissionClosure(t *testing.T) {
	362	+	runtime, _ := newTestUDPRuntime(t, "127.0.0.1:15353")
	363	+	release := make(chan struct{})
	364	+	dialStarted := make(chan struct{}, 1)
	365	+	liveConn := newLiveUDPConn(t)
	366	+	runtime.dialContext = func(ctx context.Context, _, _ string) (net.Conn, error) {
	367	+	select {
	368	+	case dialStarted <- struct{}{}:
	369	+	default:
	370	+	}
	371	+	select {
	372	+	case <-release:
	373	+	return liveConn, nil
	374	+	case <-ctx.Done():
	375	+	return nil, ctx.Err()
	376	+	}
	377	+	}
	378	+
	379	+	type result struct {
	380	+	session *udpSession
	381	+	ok bool
	382	+	}
	383	+	results := make(chan result, 1)
	384	+	go func() {
	385	+	session, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4700"), effectiveUDPOptions(domain.UDPOptions{}))
	386	+	results <- result{session: session, ok: ok}
	387	+	}()
	388	+
	389	+	select {
	390	+	case <-dialStarted:
	391	+	case <-time.After(5 * time.Second):
	392	+	t.Fatal("backend dial never reached the dial seam")
	393	+	}
	394	+
	395	+	runtime.mu.Lock()
	396	+	runtime.accepting = false
	397	+	runtime.mu.Unlock()
	398	+	runtime.closed.Store(true)
	399	+	close(release)
	400	+
	401	+	select {
	402	+	case got := <-results:
	403	+	require.False(t, got.ok, "a dial completing after admission closure must not open a session")
	404	+	require.Nil(t, got.session)
	405	+	case <-time.After(5 * time.Second):
	406	+	t.Fatal("in-flight session creation did not finish")
	407	+	}
	408	+
	409	+	runtime.mu.Lock()
	410	+	require.Empty(t, runtime.sessions)
	411	+	runtime.mu.Unlock()
	412	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
	413	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load())
	414	+	requireUDPRuntimeGoroutinesDone(t, runtime)
	415	+	}
	416	+
	417	+	func TestUDPAddrPortNormalizesIPv4MappedAndPreservesZones(t *testing.T) {
	418	+	plainV4, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9000})
	419	+	require.True(t, ok)
	420	+	mappedV4, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("::ffff:127.0.0.1"), Port: 9000})
	421	+	require.True(t, ok)
	422	+	wantV4 := netip.MustParseAddrPort("127.0.0.1:9000")
	423	+	require.Equal(t, wantV4, plainV4)
	424	+	require.Equal(t, wantV4, mappedV4, "IPv4-mapped addresses must share the IPv4 session key")
	425	+
	426	+	zoned, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("fe80::1"), Zone: "eth0", Port: 9001})
	427	+	require.True(t, ok)
	428	+	require.Equal(t, netip.MustParseAddrPort("[fe80::1%eth0]:9001"), zoned)
	429	+	otherZone, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("fe80::1"), Zone: "eth1", Port: 9001})
	430	+	require.True(t, ok)
	431	+	require.NotEqual(t, zoned, otherZone, "IPv6 zones must remain part of the key")
	432	+
	433	+	parsed, ok := udpAddrPort(testUDPNetAddr{network: "udp", address: "[2001:db8::1]:9002"})
	434	+	require.True(t, ok)
	435	+	require.Equal(t, netip.MustParseAddrPort("[2001:db8::1]:9002"), parsed)
	436	+
	437	+	_, ok = udpAddrPort(testUDPNetAddr{network: "udp", address: "not-an-address"})
	438	+	require.False(t, ok)
	439	+	}
	440	+
	441	+	func TestUDPAddrPortRejectsInvalidAddresses(t *testing.T) {
	442	+	_, ok := udpAddrPort(&net.UDPAddr{})
	443	+	require.False(t, ok, "zero UDPAddr must be rejected")
	444	+	_, ok = udpAddrPort(&net.UDPAddr{IP: nil, Port: 53})
	445	+	require.False(t, ok, "UDPAddr without an IP must be rejected")
	446	+	_, ok = udpAddrPort(testUDPNetAddr{network: "udp", address: "not-an-address"})
	447	+	require.False(t, ok)
	448	+
	449	+	// The trusted-list path must not admit an invalid address either, even when
	450	+	// the list is empty and would otherwise trust every client.
	451	+	require.False(t, trustedAddrPort(nil, netip.AddrPort{}))
	452	+	}
	453	+
	454	+	func TestUDPSessionPrecomputesImmutableRemoteAddr(t *testing.T) {
	455	+	backend := startUDPEchoServer(t)
	456	+	runtime, _ := newTestUDPRuntime(t, backend.address)
	457	+	options := effectiveUDPOptions(domain.UDPOptions{})
	458	+
	459	+	v4, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4600"), options)
	460	+	require.True(t, ok)
	461	+	require.NotNil(t, v4.remoteAddr)
	462	+	require.Equal(t, "127.0.0.1:4600", v4.remoteAddr.String())
	463	+	runtime.removeSession(v4)
	464	+
	465	+	zoned, ok := runtime.session(netip.MustParseAddrPort("[fe80::1%eth0]:4601"), options)
	466	+	require.True(t, ok)
	467	+	require.NotNil(t, zoned.remoteAddr)
	468	+	require.Equal(t, netip.MustParseAddrPort("[fe80::1%eth0]:4601"), zoned.remoteAddr.AddrPort(), "precomputed destination must preserve the IPv6 zone")
	469	+	runtime.removeSession(zoned)
	470	+	}
	471	+
	472	+	func TestUDPTrustedAddrPortMatchesCIDRs(t *testing.T) {
	473	+	ipv4Trusted, err := parseTrustedCIDRs([]string{"127.0.0.0/8"})
	474	+	require.NoError(t, err)
	475	+	require.True(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("127.0.0.1:1")))
	476	+	require.True(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("[::ffff:127.0.0.1]:1")), "v4-mapped client must match IPv4 CIDR")
	477	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("192.0.2.1:1")))
	478	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("[2001:db8::1]:1")), "IPv6 client must not match IPv4 CIDR")
	479	+
	480	+	ipv6Trusted, err := parseTrustedCIDRs([]string{"2001:db8::/32"})
	481	+	require.NoError(t, err)
	482	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1]:1")))
	483	+	// Intentional UDP semantic: CIDR trust is evaluated on the IP only. The
	484	+	// zone names the local interface a packet arrived on, not the peer, so a
	485	+	// zoned and zoneless client of the same address are equally trusted. Zones
	486	+	// stay part of the session key; do not "fix" this to compare zones.
	487	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1%eth0]:1")), "zone must not affect CIDR matching")
	488	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1%eth1]:1")), "any zone on a trusted address is trusted")
	489	+	require.False(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:dead::2%eth0]:1")), "an untrusted address is untrusted regardless of zone")
	490	+	require.False(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("127.0.0.1:1")))
	491	+
	492	+	require.True(t, trustedAddrPort(nil, netip.MustParseAddrPort("203.0.113.9:1")), "empty trusted list admits every client")
	493	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.AddrPort{}), "invalid client address is not trusted")
	494	+	}
	495	+
109	496		func assertUDPRejected(t testing.T, manager Manager, address string, refused int64) {
110	497		t.Helper()
111	498		conn := dialUDP(t, address)

@@ -330,6 +717,130 @@ func assertUDPRoundTripWant(t testing.T, conn net.UDPConn, message string, wan
330	717		assert.Equal(t, want, string(buf[:n]))
331	718		}
332	719
	720	+	// newTestUDPRuntime builds a UDP runtime backed by a valid graph snapshot
	721	+	// without binding through Manager.Apply, so tests can drive publication,
	722	+	// expiry and shutdown interleavings directly.
	723	+	func newTestUDPRuntime(t testing.T, backendAddress string) (udpEntryPointRuntime, *Manager) {
	724	+	t.Helper()
	725	+	packetConn, err := net.ListenPacket("udp", "127.0.0.1:0")
	726	+	require.NoError(t, err)
	727	+	return newTestUDPRuntimeWithPacketConn(t, backendAddress, packetConn)
	728	+	}
	729	+
	730	+	// newTestUDPRuntimeWithPacketConn builds a runtime over a caller-supplied
	731	+	// PacketConn so the read loop can be driven deterministically.
	732	+	func newTestUDPRuntimeWithPacketConn(t testing.T, backendAddress string, packetConn net.PacketConn) (udpEntryPointRuntime, *Manager) {
	733	+	t.Helper()
	734	+	graph := udpGraph(t, freeUDPAddress(t), backendAddress)
	735	+	manager := NewManager()
	736	+	manager.snapshot.Store(&graph)
	737	+	runtime := newUDPEntryPointRuntime(context.Background(), manager, graph.EntryPoints[0], packetConn, nil)
	738	+	t.Cleanup(func() {
	739	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	740	+	defer cancel()
	741	+	runtime.stop(ctx, 50*time.Millisecond)
	742	+	_ = packetConn.Close()
	743	+	})
	744	+	return runtime, manager
	745	+	}
	746	+
	747	+	// scriptedPacketConn feeds the read loop a fixed sequence of packets, including
	748	+	// source addresses that cannot be converted to a session key.
	749	+	type scriptedUDPPacket struct {
	750	+	payload []byte
	751	+	addr net.Addr
	752	+	}
	753	+
	754	+	type scriptedPacketConn struct {
	755	+	packets chan scriptedUDPPacket
	756	+	closed chan struct{}
	757	+	once sync.Once
	758	+	}
	759	+
	760	+	func newScriptedPacketConn() *scriptedPacketConn {
	761	+	return &scriptedPacketConn{packets: make(chan scriptedUDPPacket, 16), closed: make(chan struct{})}
	762	+	}
	763	+
	764	+	func (c *scriptedPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
	765	+	select {
	766	+	case packet := <-c.packets:
	767	+	return copy(p, packet.payload), packet.addr, nil
	768	+	case <-c.closed:
	769	+	return 0, nil, net.ErrClosed
	770	+	}
	771	+	}
	772	+
	773	+	func (c *scriptedPacketConn) WriteTo(p []byte, _ net.Addr) (int, error) { return len(p), nil }
	774	+
	775	+	func (c *scriptedPacketConn) Close() error {
	776	+	c.once.Do(func() { close(c.closed) })
	777	+	return nil
	778	+	}
	779	+
	780	+	func (c *scriptedPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{IP: net.IPv4zero} }
	781	+	func (c *scriptedPacketConn) SetDeadline(time.Time) error { return nil }
	782	+	func (c *scriptedPacketConn) SetReadDeadline(time.Time) error { return nil }
	783	+	func (c *scriptedPacketConn) SetWriteDeadline(time.Time) error { return nil }
	784	+
	785	+	func newTestUDPSession(t testing.T, clientAddr netip.AddrPort) udpSession {
	786	+	t.Helper()
	787	+	session := &udpSession{clientAddr: clientAddr, remoteAddr: net.UDPAddrFromAddrPort(clientAddr), backend: newLiveUDPConn(t)}
	788	+	session.touch()
	789	+	return session
	790	+	}
	791	+
	792	+	// newLiveUDPConn dials a live loopback UDP listener so probe writes never fail
	793	+	// with ICMP port-unreachable errors, and tears both down with the test.
	794	+	func newLiveUDPConn(t *testing.T) net.Conn {
	795	+	t.Helper()
	796	+	listener, err := net.ListenPacket("udp", "127.0.0.1:0")
	797	+	require.NoError(t, err)
	798	+	t.Cleanup(func() { _ = listener.Close() })
	799	+	go func() {
	800	+	buf := make([]byte, 64)
	801	+	for {
	802	+	if _, _, err := listener.ReadFrom(buf); err != nil {
	803	+	return
	804	+	}
	805	+	}
	806	+	}()
	807	+	conn, err := net.Dial("udp", listener.LocalAddr().String())
	808	+	require.NoError(t, err)
	809	+	t.Cleanup(func() { _ = conn.Close() })
	810	+	return conn
	811	+	}
	812	+
	813	+	// testUDPSessionClosed probes the backend with a write. Synthetic sessions are
	814	+	// connected to a live loopback listener (see newLiveUDPConn), so an open socket
	815	+	// always accepts the probe and a closed one always errors, without the ICMP
	816	+	// port-unreachable flakiness of writing to an unbound port.
	817	+	func testUDPSessionClosed(session *udpSession) bool {
	818	+	_, err := session.backend.Write([]byte("probe"))
	819	+	return err != nil
	820	+	}
	821	+
	822	+	func requireUDPRuntimeGoroutinesDone(t testing.T, runtime udpEntryPointRuntime) {
	823	+	t.Helper()
	824	+	done := make(chan struct{})
	825	+	go func() {
	826	+	runtime.runWG.Wait()
	827	+	close(done)
	828	+	}()
	829	+	select {
	830	+	case <-done:
	831	+	case <-time.After(2 * time.Second):
	832	+	t.Fatal("udp runtime goroutines did not complete")
	833	+	}
	834	+	}
	835	+
	836	+	type testUDPNetAddr struct {
	837	+	network string
	838	+	address string
	839	+	}
	840	+
	841	+	func (a testUDPNetAddr) Network() string { return a.network }
	842	+	func (a testUDPNetAddr) String() string { return a.address }
	843	+
333	844		type udpEchoServer struct{ address string }
334	845
335	846		func startUDPEchoServer(t *testing.T) udpEchoServer {


@@ -3,6 +3,8 @@ package traffic
3	import (	3		import (
4	"context"	4		"context"
5	"net"	5		"net"
		6	+	"net/netip"
		7	+	"sync"
6	"testing"	8		"testing"
7	"time"	9		"time"
8		10

@@ -106,6 +108,391 @@ func TestUDPReloadReplacesSessionWhenBackendChanges(t *testing.T) {
106	assertUDPRoundTripWant(t, newConn, "third", "b:third")	108		assertUDPRoundTripWant(t, newConn, "third", "b:third")
107	}	109		}
108		110
		111	+	func TestUDPStaleRemovalKeepsReplacementSession(t *testing.T) {
		112	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4100")
		113	+	stale := newTestUDPSession(t, clientAddr)
		114	+	replacement := newTestUDPSession(t, clientAddr)
		115	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
		116	+	runtime.mu.Lock()
		117	+	runtime.sessions[clientAddr] = replacement
		118	+	runtime.counters.activeUDPSessions.Store(1)
		119	+	runtime.mu.Unlock()
		120	+
		121	+	// A stale remover carries the pointer of the session that used to own the
		122	+	// client key. Expiry, backend read/write failures, reload and CIDR cleanup
		123	+	// all funnel through removeSession, so none of them may evict the session
		124	+	// that replaced it.
		125	+	runtime.removeSession(stale)
		126	+
		127	+	runtime.mu.Lock()
		128	+	current, ok := runtime.sessions[clientAddr]
		129	+	runtime.mu.Unlock()
		130	+	require.True(t, ok, "replacement session must stay registered")
		131	+	require.Same(t, replacement, current)
		132	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load(), "stale removal must not decrement the active counter")
		133	+	require.False(t, testUDPSessionClosed(replacement), "replacement backend must stay open")
		134	+
		135	+	runtime.removeSession(replacement)
		136	+
		137	+	runtime.mu.Lock()
		138	+	_, ok = runtime.sessions[clientAddr]
		139	+	runtime.mu.Unlock()
		140	+	require.False(t, ok, "removing the current session must unregister it")
		141	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load(), "removing the current session must decrement once")
		142	+	require.True(t, testUDPSessionClosed(replacement))
		143	+
		144	+	// Removing again must be idempotent: no double decrement and no panic.
		145	+	runtime.removeSession(replacement)
		146	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		147	+	}
		148	+
		149	+	func TestUDPExpiryClosesOnlyRegisteredIdleSessionOnce(t *testing.T) {
		150	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4200")
		151	+	// stale was replaced under the same key and is no longer registered.
		152	+	stale := newTestUDPSession(t, clientAddr)
		153	+	stale.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
		154	+	current := newTestUDPSession(t, clientAddr)
		155	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
		156	+	runtime.mu.Lock()
		157	+	runtime.sessions[clientAddr] = current
		158	+	runtime.counters.activeUDPSessions.Store(1)
		159	+	runtime.mu.Unlock()
		160	+
		161	+	runtime.expireIdleSessions(time.Minute)
		162	+	require.False(t, testUDPSessionClosed(current), "fresh registered session must not expire")
		163	+	require.False(t, testUDPSessionClosed(stale), "an unregistered session must never be closed by expiry")
		164	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load())
		165	+
		166	+	current.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
		167	+	runtime.expireIdleSessions(time.Minute)
		168	+
		169	+	runtime.mu.Lock()
		170	+	_, ok := runtime.sessions[clientAddr]
		171	+	runtime.mu.Unlock()
		172	+	require.False(t, ok)
		173	+	require.True(t, testUDPSessionClosed(current))
		174	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load(), "expiry must decrement exactly once")
		175	+
		176	+	// The backend loop of the expired session also removes it on return; the
		177	+	// identity check keeps the counter coherent and must not double decrement.
		178	+	runtime.removeSession(current)
		179	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		180	+	}
		181	+
		182	+	func TestUDPAdmissionGateBlocksPublicationAfterClosingBegins(t *testing.T) {
		183	+	backend := startUDPEchoServer(t)
		184	+	runtime, _ := newTestUDPRuntime(t, backend.address)
		185	+	// Model the instant stop() closes admission but before it cancels the
		186	+	// runtime context: a datagram worker that reaches session() now must not
		187	+	// publish a new session.
		188	+	runtime.mu.Lock()
		189	+	runtime.accepting = false
		190	+	runtime.mu.Unlock()
		191	+	runtime.closed.Store(true)
		192	+
		193	+	session, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4300"), effectiveUDPOptions(domain.UDPOptions{}))
		194	+	require.False(t, ok)
		195	+	require.Nil(t, session)
		196	+
		197	+	runtime.mu.Lock()
		198	+	require.Empty(t, runtime.sessions)
		199	+	runtime.mu.Unlock()
		200	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		201	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		202	+	}
		203	+
		204	+	func TestUDPPublishSessionRejectsCompletedDialAfterAdmissionClosure(t *testing.T) {
		205	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
		206	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4800")
		207	+	dialed := newTestUDPSession(t, clientAddr)
		208	+	runtime.mu.Lock()
		209	+	runtime.accepting = false
		210	+	runtime.mu.Unlock()
		211	+	runtime.closed.Store(true)
		212	+
		213	+	session, ok := runtime.publishSession(dialed, effectiveUDPOptions(domain.UDPOptions{}))
		214	+	require.False(t, ok)
		215	+	require.Nil(t, session)
		216	+	require.True(t, testUDPSessionClosed(dialed), "a rejected dial must have its backend closed")
		217	+	runtime.mu.Lock()
		218	+	require.Empty(t, runtime.sessions)
		219	+	runtime.mu.Unlock()
		220	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		221	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load())
		222	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		223	+	}
		224	+
		225	+	func TestUDPPublishSessionReusesExistingAndTouchesIt(t *testing.T) {
		226	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
		227	+	clientAddr := netip.MustParseAddrPort("127.0.0.1:4900")
		228	+	existing := newTestUDPSession(t, clientAddr)
		229	+	existing.lastSeen.Store(time.Now().Add(-time.Hour).UnixNano())
		230	+	dialed := newTestUDPSession(t, clientAddr)
		231	+	runtime.mu.Lock()
		232	+	runtime.sessions[clientAddr] = existing
		233	+	runtime.counters.activeUDPSessions.Store(1)
		234	+	runtime.mu.Unlock()
		235	+
		236	+	session, ok := runtime.publishSession(dialed, effectiveUDPOptions(domain.UDPOptions{}))
		237	+	require.True(t, ok)
		238	+	require.Same(t, existing, session)
		239	+	require.True(t, testUDPSessionClosed(dialed), "the losing dial must have its backend closed")
		240	+	require.False(t, testUDPSessionClosed(existing))
		241	+	require.Greater(t, existing.lastSeen.Load(), time.Now().Add(-time.Minute).UnixNano(), "reuse must touch the existing session")
		242	+	runtime.mu.Lock()
		243	+	current := runtime.sessions[clientAddr]
		244	+	runtime.mu.Unlock()
		245	+	require.Same(t, existing, current)
		246	+	require.Equal(t, int64(1), runtime.counters.activeUDPSessions.Load(), "reuse must not add a session")
		247	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load(), "reuse must not count as a new accept")
		248	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		249	+	}
		250	+
		251	+	func TestUDPRemoveSessionNilIsSafe(t *testing.T) {
		252	+	runtime, _ := newTestUDPRuntime(t, freeUDPAddress(t))
		253	+	require.NotPanics(t, func() { runtime.removeSession(nil) })
		254	+	runtime.mu.Lock()
		255	+	require.Empty(t, runtime.sessions)
		256	+	runtime.mu.Unlock()
		257	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		258	+	}
		259	+
		260	+	func TestUDPReadLoopRejectsInvalidSourceAddress(t *testing.T) {
		261	+	backend := startUDPEchoServer(t)
		262	+	packetConn := newScriptedPacketConn()
		263	+	runtime, _ := newTestUDPRuntimeWithPacketConn(t, backend.address, packetConn)
		264	+	runtime.start()
		265	+
		266	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("zero"), addr: &net.UDPAddr{}}
		267	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("junk"), addr: testUDPNetAddr{network: "udp", address: "not-an-address"}}
		268	+	require.Eventually(t, func() bool { return runtime.counters.totalRefused.Load() == 2 }, time.Second, 5*time.Millisecond)
		269	+	runtime.mu.Lock()
		270	+	require.Empty(t, runtime.sessions, "invalid source addresses must never open a session")
		271	+	runtime.mu.Unlock()
		272	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		273	+
		274	+	// A valid datagram after the invalid ones proves the read loop keeps serving.
		275	+	packetConn.packets <- scriptedUDPPacket{payload: []byte("ok"), addr: &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 43210}}
		276	+	require.Eventually(t, func() bool { return runtime.counters.activeUDPSessions.Load() == 1 }, time.Second, 5*time.Millisecond)
		277	+	require.Equal(t, int64(1), runtime.counters.totalAccepted.Load())
		278	+	}
		279	+
		280	+	func TestUDPStopBeforeStartLeavesNoRuntimeGoroutines(t *testing.T) {
		281	+	backend := startUDPEchoServer(t)
		282	+	runtime, _ := newTestUDPRuntime(t, backend.address)
		283	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		284	+	defer cancel()
		285	+
		286	+	runtime.stop(ctx, 50*time.Millisecond)
		287	+	// Starting a runtime whose admission already closed must not launch workers.
		288	+	runtime.start()
		289	+
		290	+	runtime.mu.Lock()
		291	+	require.Empty(t, runtime.sessions)
		292	+	runtime.mu.Unlock()
		293	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		294	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		295	+	}
		296	+
		297	+	func TestUDPShutdownLeavesNoSessionsOrRuntimeGoroutines(t *testing.T) {
		298	+	backend := startUDPEchoServer(t)
		299	+	graph := udpGraph(t, freeUDPAddress(t), backend.address)
		300	+	graph.Options.UDP.DrainTimeout = 50 * time.Millisecond
		301	+	manager := NewManager()
		302	+	require.NoError(t, manager.Apply(context.Background(), &graph))
		303	+
		304	+	conn := dialUDP(t, graph.EntryPoints[0].Address)
		305	+	defer conn.Close()
		306	+	assertUDPRoundTrip(t, conn, "held")
		307	+	require.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 1 }, time.Second, 10*time.Millisecond)
		308	+
		309	+	manager.mu.Lock()
		310	+	runtime := manager.udpListeners["udp"]
		311	+	manager.mu.Unlock()
		312	+	require.NotNil(t, runtime)
		313	+
		314	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		315	+	defer cancel()
		316	+	require.NoError(t, manager.Shutdown(ctx))
		317	+
		318	+	runtime.mu.Lock()
		319	+	require.Empty(t, runtime.sessions)
		320	+	runtime.mu.Unlock()
		321	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		322	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		323	+	}
		324	+
		325	+	func TestUDPConcurrentCreationDuringShutdownLeavesNoSessions(t *testing.T) {
		326	+	for range 5 {
		327	+	backend := startUDPEchoServer(t)
		328	+	runtime, _ := newTestUDPRuntime(t, backend.address)
		329	+	runtime.start()
		330	+
		331	+	var creators sync.WaitGroup
		332	+	for worker := range 8 {
		333	+	creators.Add(1)
		334	+	go func(worker int) {
		335	+	defer creators.Done()
		336	+	for attempt := range 25 {
		337	+	key := netip.AddrPortFrom(netip.AddrFrom4([4]byte{127, 0, 0, 1}), uint16(5000+worker*100+attempt))
		338	+	runtime.session(key, effectiveUDPOptions(domain.UDPOptions{}))
		339	+	}
		340	+	}(worker)
		341	+	}
		342	+
		343	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		344	+	runtime.stop(ctx, 50*time.Millisecond)
		345	+	cancel()
		346	+	creators.Wait()
		347	+
		348	+	runtime.mu.Lock()
		349	+	require.Empty(t, runtime.sessions, "no session may be published after closing began")
		350	+	runtime.mu.Unlock()
		351	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		352	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		353	+	}
		354	+	}
		355	+
		356	+	// TestUDPInFlightDialDoesNotPublishAfterAdmissionClosure blocks the backend dial
		357	+	// on the runtime's private dial seam, closes admission exactly as stop() does,
		358	+	// then lets the dial complete: session() must reject the completed dial without
		359	+	// publishing a session or a backend goroutine. The seam is instance-local, so no
		360	+	// package-global state is touched.
		361	+	func TestUDPInFlightDialDoesNotPublishAfterAdmissionClosure(t *testing.T) {
		362	+	runtime, _ := newTestUDPRuntime(t, "127.0.0.1:15353")
		363	+	release := make(chan struct{})
		364	+	dialStarted := make(chan struct{}, 1)
		365	+	liveConn := newLiveUDPConn(t)
		366	+	runtime.dialContext = func(ctx context.Context, _, _ string) (net.Conn, error) {
		367	+	select {
		368	+	case dialStarted <- struct{}{}:
		369	+	default:
		370	+	}
		371	+	select {
		372	+	case <-release:
		373	+	return liveConn, nil
		374	+	case <-ctx.Done():
		375	+	return nil, ctx.Err()
		376	+	}
		377	+	}
		378	+
		379	+	type result struct {
		380	+	session *udpSession
		381	+	ok bool
		382	+	}
		383	+	results := make(chan result, 1)
		384	+	go func() {
		385	+	session, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4700"), effectiveUDPOptions(domain.UDPOptions{}))
		386	+	results <- result{session: session, ok: ok}
		387	+	}()
		388	+
		389	+	select {
		390	+	case <-dialStarted:
		391	+	case <-time.After(5 * time.Second):
		392	+	t.Fatal("backend dial never reached the dial seam")
		393	+	}
		394	+
		395	+	runtime.mu.Lock()
		396	+	runtime.accepting = false
		397	+	runtime.mu.Unlock()
		398	+	runtime.closed.Store(true)
		399	+	close(release)
		400	+
		401	+	select {
		402	+	case got := <-results:
		403	+	require.False(t, got.ok, "a dial completing after admission closure must not open a session")
		404	+	require.Nil(t, got.session)
		405	+	case <-time.After(5 * time.Second):
		406	+	t.Fatal("in-flight session creation did not finish")
		407	+	}
		408	+
		409	+	runtime.mu.Lock()
		410	+	require.Empty(t, runtime.sessions)
		411	+	runtime.mu.Unlock()
		412	+	require.Equal(t, int64(0), runtime.counters.activeUDPSessions.Load())
		413	+	require.Equal(t, int64(0), runtime.counters.totalAccepted.Load())
		414	+	requireUDPRuntimeGoroutinesDone(t, runtime)
		415	+	}
		416	+
		417	+	func TestUDPAddrPortNormalizesIPv4MappedAndPreservesZones(t *testing.T) {
		418	+	plainV4, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9000})
		419	+	require.True(t, ok)
		420	+	mappedV4, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("::ffff:127.0.0.1"), Port: 9000})
		421	+	require.True(t, ok)
		422	+	wantV4 := netip.MustParseAddrPort("127.0.0.1:9000")
		423	+	require.Equal(t, wantV4, plainV4)
		424	+	require.Equal(t, wantV4, mappedV4, "IPv4-mapped addresses must share the IPv4 session key")
		425	+
		426	+	zoned, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("fe80::1"), Zone: "eth0", Port: 9001})
		427	+	require.True(t, ok)
		428	+	require.Equal(t, netip.MustParseAddrPort("[fe80::1%eth0]:9001"), zoned)
		429	+	otherZone, ok := udpAddrPort(&net.UDPAddr{IP: net.ParseIP("fe80::1"), Zone: "eth1", Port: 9001})
		430	+	require.True(t, ok)
		431	+	require.NotEqual(t, zoned, otherZone, "IPv6 zones must remain part of the key")
		432	+
		433	+	parsed, ok := udpAddrPort(testUDPNetAddr{network: "udp", address: "[2001:db8::1]:9002"})
		434	+	require.True(t, ok)
		435	+	require.Equal(t, netip.MustParseAddrPort("[2001:db8::1]:9002"), parsed)
		436	+
		437	+	_, ok = udpAddrPort(testUDPNetAddr{network: "udp", address: "not-an-address"})
		438	+	require.False(t, ok)
		439	+	}
		440	+
		441	+	func TestUDPAddrPortRejectsInvalidAddresses(t *testing.T) {
		442	+	_, ok := udpAddrPort(&net.UDPAddr{})
		443	+	require.False(t, ok, "zero UDPAddr must be rejected")
		444	+	_, ok = udpAddrPort(&net.UDPAddr{IP: nil, Port: 53})
		445	+	require.False(t, ok, "UDPAddr without an IP must be rejected")
		446	+	_, ok = udpAddrPort(testUDPNetAddr{network: "udp", address: "not-an-address"})
		447	+	require.False(t, ok)
		448	+
		449	+	// The trusted-list path must not admit an invalid address either, even when
		450	+	// the list is empty and would otherwise trust every client.
		451	+	require.False(t, trustedAddrPort(nil, netip.AddrPort{}))
		452	+	}
		453	+
		454	+	func TestUDPSessionPrecomputesImmutableRemoteAddr(t *testing.T) {
		455	+	backend := startUDPEchoServer(t)
		456	+	runtime, _ := newTestUDPRuntime(t, backend.address)
		457	+	options := effectiveUDPOptions(domain.UDPOptions{})
		458	+
		459	+	v4, ok := runtime.session(netip.MustParseAddrPort("127.0.0.1:4600"), options)
		460	+	require.True(t, ok)
		461	+	require.NotNil(t, v4.remoteAddr)
		462	+	require.Equal(t, "127.0.0.1:4600", v4.remoteAddr.String())
		463	+	runtime.removeSession(v4)
		464	+
		465	+	zoned, ok := runtime.session(netip.MustParseAddrPort("[fe80::1%eth0]:4601"), options)
		466	+	require.True(t, ok)
		467	+	require.NotNil(t, zoned.remoteAddr)
		468	+	require.Equal(t, netip.MustParseAddrPort("[fe80::1%eth0]:4601"), zoned.remoteAddr.AddrPort(), "precomputed destination must preserve the IPv6 zone")
		469	+	runtime.removeSession(zoned)
		470	+	}
		471	+
		472	+	func TestUDPTrustedAddrPortMatchesCIDRs(t *testing.T) {
		473	+	ipv4Trusted, err := parseTrustedCIDRs([]string{"127.0.0.0/8"})
		474	+	require.NoError(t, err)
		475	+	require.True(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("127.0.0.1:1")))
		476	+	require.True(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("[::ffff:127.0.0.1]:1")), "v4-mapped client must match IPv4 CIDR")
		477	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("192.0.2.1:1")))
		478	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.MustParseAddrPort("[2001:db8::1]:1")), "IPv6 client must not match IPv4 CIDR")
		479	+
		480	+	ipv6Trusted, err := parseTrustedCIDRs([]string{"2001:db8::/32"})
		481	+	require.NoError(t, err)
		482	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1]:1")))
		483	+	// Intentional UDP semantic: CIDR trust is evaluated on the IP only. The
		484	+	// zone names the local interface a packet arrived on, not the peer, so a
		485	+	// zoned and zoneless client of the same address are equally trusted. Zones
		486	+	// stay part of the session key; do not "fix" this to compare zones.
		487	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1%eth0]:1")), "zone must not affect CIDR matching")
		488	+	require.True(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:db8::1%eth1]:1")), "any zone on a trusted address is trusted")
		489	+	require.False(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("[2001:dead::2%eth0]:1")), "an untrusted address is untrusted regardless of zone")
		490	+	require.False(t, trustedAddrPort(ipv6Trusted, netip.MustParseAddrPort("127.0.0.1:1")))
		491	+
		492	+	require.True(t, trustedAddrPort(nil, netip.MustParseAddrPort("203.0.113.9:1")), "empty trusted list admits every client")
		493	+	require.False(t, trustedAddrPort(ipv4Trusted, netip.AddrPort{}), "invalid client address is not trusted")
		494	+	}
		495	+
109	func assertUDPRejected(t testing.T, manager Manager, address string, refused int64) {	496		func assertUDPRejected(t testing.T, manager Manager, address string, refused int64) {
110	t.Helper()	497		t.Helper()
111	conn := dialUDP(t, address)	498		conn := dialUDP(t, address)

@@ -330,6 +717,130 @@ func assertUDPRoundTripWant(t testing.T, conn net.UDPConn, message string, wan
330	assert.Equal(t, want, string(buf[:n]))	717		assert.Equal(t, want, string(buf[:n]))
331	}	718		}
332		719
		720	+	// newTestUDPRuntime builds a UDP runtime backed by a valid graph snapshot
		721	+	// without binding through Manager.Apply, so tests can drive publication,
		722	+	// expiry and shutdown interleavings directly.
		723	+	func newTestUDPRuntime(t testing.T, backendAddress string) (udpEntryPointRuntime, *Manager) {
		724	+	t.Helper()
		725	+	packetConn, err := net.ListenPacket("udp", "127.0.0.1:0")
		726	+	require.NoError(t, err)
		727	+	return newTestUDPRuntimeWithPacketConn(t, backendAddress, packetConn)
		728	+	}
		729	+
		730	+	// newTestUDPRuntimeWithPacketConn builds a runtime over a caller-supplied
		731	+	// PacketConn so the read loop can be driven deterministically.
		732	+	func newTestUDPRuntimeWithPacketConn(t testing.T, backendAddress string, packetConn net.PacketConn) (udpEntryPointRuntime, *Manager) {
		733	+	t.Helper()
		734	+	graph := udpGraph(t, freeUDPAddress(t), backendAddress)
		735	+	manager := NewManager()
		736	+	manager.snapshot.Store(&graph)
		737	+	runtime := newUDPEntryPointRuntime(context.Background(), manager, graph.EntryPoints[0], packetConn, nil)
		738	+	t.Cleanup(func() {
		739	+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		740	+	defer cancel()
		741	+	runtime.stop(ctx, 50*time.Millisecond)
		742	+	_ = packetConn.Close()
		743	+	})
		744	+	return runtime, manager
		745	+	}
		746	+
		747	+	// scriptedPacketConn feeds the read loop a fixed sequence of packets, including
		748	+	// source addresses that cannot be converted to a session key.
		749	+	type scriptedUDPPacket struct {
		750	+	payload []byte
		751	+	addr net.Addr
		752	+	}
		753	+
		754	+	type scriptedPacketConn struct {
		755	+	packets chan scriptedUDPPacket
		756	+	closed chan struct{}
		757	+	once sync.Once
		758	+	}
		759	+
		760	+	func newScriptedPacketConn() *scriptedPacketConn {
		761	+	return &scriptedPacketConn{packets: make(chan scriptedUDPPacket, 16), closed: make(chan struct{})}
		762	+	}
		763	+
		764	+	func (c *scriptedPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
		765	+	select {
		766	+	case packet := <-c.packets:
		767	+	return copy(p, packet.payload), packet.addr, nil
		768	+	case <-c.closed:
		769	+	return 0, nil, net.ErrClosed
		770	+	}
		771	+	}
		772	+
		773	+	func (c *scriptedPacketConn) WriteTo(p []byte, _ net.Addr) (int, error) { return len(p), nil }
		774	+
		775	+	func (c *scriptedPacketConn) Close() error {
		776	+	c.once.Do(func() { close(c.closed) })
		777	+	return nil
		778	+	}
		779	+
		780	+	func (c *scriptedPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{IP: net.IPv4zero} }
		781	+	func (c *scriptedPacketConn) SetDeadline(time.Time) error { return nil }
		782	+	func (c *scriptedPacketConn) SetReadDeadline(time.Time) error { return nil }
		783	+	func (c *scriptedPacketConn) SetWriteDeadline(time.Time) error { return nil }
		784	+
		785	+	func newTestUDPSession(t testing.T, clientAddr netip.AddrPort) udpSession {
		786	+	t.Helper()
		787	+	session := &udpSession{clientAddr: clientAddr, remoteAddr: net.UDPAddrFromAddrPort(clientAddr), backend: newLiveUDPConn(t)}
		788	+	session.touch()
		789	+	return session
		790	+	}
		791	+
		792	+	// newLiveUDPConn dials a live loopback UDP listener so probe writes never fail
		793	+	// with ICMP port-unreachable errors, and tears both down with the test.
		794	+	func newLiveUDPConn(t *testing.T) net.Conn {
		795	+	t.Helper()
		796	+	listener, err := net.ListenPacket("udp", "127.0.0.1:0")
		797	+	require.NoError(t, err)
		798	+	t.Cleanup(func() { _ = listener.Close() })
		799	+	go func() {
		800	+	buf := make([]byte, 64)
		801	+	for {
		802	+	if _, _, err := listener.ReadFrom(buf); err != nil {
		803	+	return
		804	+	}
		805	+	}
		806	+	}()
		807	+	conn, err := net.Dial("udp", listener.LocalAddr().String())
		808	+	require.NoError(t, err)
		809	+	t.Cleanup(func() { _ = conn.Close() })
		810	+	return conn
		811	+	}
		812	+
		813	+	// testUDPSessionClosed probes the backend with a write. Synthetic sessions are
		814	+	// connected to a live loopback listener (see newLiveUDPConn), so an open socket
		815	+	// always accepts the probe and a closed one always errors, without the ICMP
		816	+	// port-unreachable flakiness of writing to an unbound port.
		817	+	func testUDPSessionClosed(session *udpSession) bool {
		818	+	_, err := session.backend.Write([]byte("probe"))
		819	+	return err != nil
		820	+	}
		821	+
		822	+	func requireUDPRuntimeGoroutinesDone(t testing.T, runtime udpEntryPointRuntime) {
		823	+	t.Helper()
		824	+	done := make(chan struct{})
		825	+	go func() {
		826	+	runtime.runWG.Wait()
		827	+	close(done)
		828	+	}()
		829	+	select {
		830	+	case <-done:
		831	+	case <-time.After(2 * time.Second):
		832	+	t.Fatal("udp runtime goroutines did not complete")
		833	+	}
		834	+	}
		835	+
		836	+	type testUDPNetAddr struct {
		837	+	network string
		838	+	address string
		839	+	}
		840	+
		841	+	func (a testUDPNetAddr) Network() string { return a.network }
		842	+	func (a testUDPNetAddr) String() string { return a.address }
		843	+
333	type udpEchoServer struct{ address string }	844		type udpEchoServer struct{ address string }
334		845
335	func startUDPEchoServer(t *testing.T) udpEchoServer {	846		func startUDPEchoServer(t *testing.T) udpEchoServer {