1package traffic
2
3import (
4 "context"
5 "errors"
6 "net"
7 "strconv"
8 "sync"
9 "sync/atomic"
10 "time"
11
12 "github.com/bnema/gordon/internal/domain"
13)
14
15const (
16 udpBufferSize = 64 << 10
17 udpWorkerCount = 16
18 udpDatagramBacklog = 1024
19)
20
21type udpEntryPointRuntime struct {
22 manager *Manager
23 entryPoint domain.EntryPoint
24 packetConn net.PacketConn
25 counters trafficCounters
26 trusted []*net.IPNet
27
28 started atomic.Bool
29 closed atomic.Bool
30
31 ctx context.Context
32 cancel context.CancelFunc
33 done chan struct{}
34 doneOnce sync.Once
35
36 datagrams chan udpDatagram
37
38 mu sync.Mutex
39 sessions map[string]*udpSession
40}
41
42type udpDatagram struct {
43 clientAddr net.Addr
44 packet []byte
45}
46
47type udpSession struct {
48 clientAddr net.Addr
49 backend net.Conn
50 backendRef domain.TrafficBackend
51 lastSeen atomic.Int64
52 done chan struct{}
53 once sync.Once
54}
55
56func newUDPEntryPointRuntime(parentCtx context.Context, manager *Manager, entryPoint domain.EntryPoint, packetConn net.PacketConn, trusted []*net.IPNet) *udpEntryPointRuntime {
57 ctx, cancel := context.WithCancel(context.WithoutCancel(parentCtx))
58 return &udpEntryPointRuntime{
59 manager: manager,
60 entryPoint: entryPoint,
61 packetConn: packetConn,
62 trusted: trusted,
63 ctx: ctx,
64 cancel: cancel,
65 done: make(chan struct{}),
66 datagrams: make(chan udpDatagram, udpDatagramBacklog),
67 sessions: map[string]*udpSession{},
68 }
69}
70
71func (r *udpEntryPointRuntime) start() {
72 if !r.started.CompareAndSwap(false, true) {
73 return
74 }
75 entryPoint := r.entryPointSnapshot()
76 trafficInfo(r.ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Str("protocol", string(entryPoint.Protocol)).Msg("started udp traffic entrypoint")
77 for range udpWorkerCount {
78 go r.datagramWorker()
79 }
80 go r.readLoop()
81 go r.expireLoop()
82}
83
84func (r *udpEntryPointRuntime) readLoop() {
85 defer r.closeDone()
86 buf := make([]byte, udpBufferSize)
87 for {
88 n, clientAddr, err := r.packetConn.ReadFrom(buf)
89 if err != nil {
90 if r.isClosed() || errors.Is(err, net.ErrClosed) {
91 return
92 }
93 r.counters.totalErrors.Add(1)
94 continue
95 }
96 if !trustedRemoteAddr(r.trustedSnapshot(), clientAddr) {
97 r.counters.totalRefused.Add(1)
98 continue
99 }
100 select {
101 case <-r.ctx.Done():
102 return
103 default:
104 }
105 if len(r.datagrams) >= cap(r.datagrams) {
106 r.counters.totalRefused.Add(1)
107 continue
108 }
109 packet := append([]byte(nil), buf[:n]...)
110 select {
111 case r.datagrams <- udpDatagram{clientAddr: clientAddr, packet: packet}:
112 case <-r.ctx.Done():
113 return
114 default:
115 r.counters.totalRefused.Add(1)
116 }
117 }
118}
119
120func (r *udpEntryPointRuntime) datagramWorker() {
121 for {
122 select {
123 case <-r.ctx.Done():
124 return
125 case datagram := <-r.datagrams:
126 select {
127 case <-r.ctx.Done():
128 return
129 default:
130 }
131 r.handleDatagram(datagram.clientAddr, datagram.packet)
132 }
133 }
134}
135
136func (r *udpEntryPointRuntime) handleDatagram(clientAddr net.Addr, packet []byte) {
137 options := effectiveUDPOptions(snapshotUDPOptions(r.manager.snapshot.Load()))
138 session, ok := r.session(clientAddr.String(), clientAddr, options)
139 if !ok {
140 return
141 }
142 n, err := session.backend.Write(packet)
143 r.counters.bytesIn.Add(int64(n))
144 if err != nil {
145 r.counters.totalErrors.Add(1)
146 r.removeSession(clientAddr.String())
147 }
148}
149
150func (r *udpEntryPointRuntime) session(key string, clientAddr net.Addr, options domain.UDPOptions) (*udpSession, bool) {
151 r.mu.Lock()
152 if session := r.sessions[key]; session != nil {
153 session.touch()
154 r.mu.Unlock()
155 return session, true
156 }
157 r.mu.Unlock()
158
159 backend, ok := r.resolveUDPBackend()
160 if !ok {
161 r.counters.totalRefused.Add(1)
162 return nil, false
163 }
164
165 r.mu.Lock()
166 if options.MaxSessions > 0 && len(r.sessions) >= options.MaxSessions {
167 r.mu.Unlock()
168 r.counters.totalRefused.Add(1)
169 return nil, false
170 }
171 r.mu.Unlock()
172 dialCtx, cancel := context.WithTimeout(r.ctx, udpDialTimeout(options))
173 backendConn, err := (&net.Dialer{}).DialContext(dialCtx, "udp", net.JoinHostPort(backend.Host, strconv.Itoa(backend.Port)))
174 cancel()
175 if err != nil {
176 r.counters.totalErrors.Add(1)
177 return nil, false
178 }
179 session := &udpSession{clientAddr: clientAddr, backend: backendConn, backendRef: backend, done: make(chan struct{})}
180 session.touch()
181
182 r.mu.Lock()
183 if existing := r.sessions[key]; existing != nil {
184 r.mu.Unlock()
185 _ = backendConn.Close()
186 return existing, true
187 }
188 if options.MaxSessions > 0 && len(r.sessions) >= options.MaxSessions {
189 r.mu.Unlock()
190 _ = backendConn.Close()
191 r.counters.totalRefused.Add(1)
192 return nil, false
193 }
194 r.sessions[key] = session
195 r.mu.Unlock()
196
197 r.counters.activeUDPSessions.Add(1)
198 r.counters.totalAccepted.Add(1)
199 go r.backendLoop(key, session)
200 return session, true
201}
202
203func trafficBackendEqual(left domain.TrafficBackend, right domain.TrafficBackend) bool {
204 return left.Name == right.Name && left.Host == right.Host && left.Port == right.Port && left.Protocol == right.Protocol
205}
206
207func udpDialTimeout(options domain.UDPOptions) time.Duration {
208 if options.IdleTimeout > 0 && options.IdleTimeout < 5*time.Second {
209 return options.IdleTimeout
210 }
211 return 5 * time.Second
212}
213
214func (r *udpEntryPointRuntime) backendLoop(key string, session *udpSession) {
215 defer r.removeSession(key)
216 buf := make([]byte, udpBufferSize)
217 for {
218 n, err := session.backend.Read(buf)
219 if err != nil {
220 if !errors.Is(err, net.ErrClosed) && !r.isClosed() {
221 r.counters.totalErrors.Add(1)
222 }
223 return
224 }
225 written, err := r.packetConn.WriteTo(buf[:n], session.clientAddr)
226 r.counters.bytesOut.Add(int64(written))
227 if err != nil {
228 if !r.isClosed() {
229 r.counters.totalErrors.Add(1)
230 }
231 return
232 }
233 }
234}
235
236func (r *udpEntryPointRuntime) expireLoop() {
237 ticker := time.NewTicker(25 * time.Millisecond)
238 defer ticker.Stop()
239 for {
240 select {
241 case <-r.ctx.Done():
242 return
243 case <-ticker.C:
244 r.expireIdleSessions(effectiveUDPOptions(snapshotUDPOptions(r.manager.snapshot.Load())).IdleTimeout)
245 }
246 }
247}
248
249func (r *udpEntryPointRuntime) expireIdleSessions(idleTimeout time.Duration) {
250 if idleTimeout <= 0 {
251 idleTimeout = defaultUDPOptions().IdleTimeout
252 }
253 cutoff := time.Now().Add(-idleTimeout).UnixNano()
254 expired := []*udpSession{}
255 r.mu.Lock()
256 for key, session := range r.sessions {
257 if session.lastSeen.Load() <= cutoff {
258 delete(r.sessions, key)
259 expired = append(expired, session)
260 }
261 }
262 r.mu.Unlock()
263 for _, session := range expired {
264 session.close()
265 r.counters.activeUDPSessions.Add(-1)
266 }
267}
268
269func (r *udpEntryPointRuntime) removeSession(key string) {
270 r.mu.Lock()
271 session := r.sessions[key]
272 if session != nil {
273 delete(r.sessions, key)
274 }
275 r.mu.Unlock()
276 if session == nil {
277 return
278 }
279 session.close()
280 r.counters.activeUDPSessions.Add(-1)
281}
282
283func (r *udpEntryPointRuntime) resolveUDPBackend() (domain.TrafficBackend, bool) {
284 graph := r.manager.snapshot.Load()
285 if graph == nil {
286 return domain.TrafficBackend{}, false
287 }
288 entryPoint := r.entryPointSnapshot()
289 var router domain.TrafficRouter
290 for _, candidate := range graph.Routers {
291 if candidate.EntryPoint == entryPoint.Name && candidate.Protocol == domain.RouterProtocolUDP {
292 if router.Name != "" {
293 return domain.TrafficBackend{}, false
294 }
295 router = candidate
296 }
297 }
298 if router.Name == "" {
299 return domain.TrafficBackend{}, false
300 }
301 for _, service := range graph.Services {
302 if service.Name == router.Service && len(service.Backends) == 1 && service.Backends[0].Protocol == domain.NetworkProtocolUDP {
303 return service.Backends[0], true
304 }
305 }
306 return domain.TrafficBackend{}, false
307}
308
309func (r *udpEntryPointRuntime) stop(ctx context.Context, drainTimeout time.Duration) {
310 if r.closed.CompareAndSwap(false, true) {
311 entryPoint := r.entryPointSnapshot()
312 trafficInfo(ctx).Str("entrypoint", entryPoint.Name).Str("address", entryPoint.Address).Msg("stopping udp traffic entrypoint")
313 r.cancel()
314 _ = r.packetConn.Close()
315 if !r.started.Load() {
316 r.closeDone()
317 }
318 }
319 select {
320 case <-r.done:
321 case <-ctx.Done():
322 r.closeSessions()
323 return
324 }
325 if r.waitSessions(ctx, drainTimeout) {
326 trafficInfo(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Msg("stopped udp traffic entrypoint")
327 return
328 }
329 trafficDebug(ctx).Str("entrypoint", r.entryPointSnapshot().Name).Dur("drain_timeout", drainTimeout).Msg("forcing udp traffic entrypoint drain")
330 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():
335 }
336}
337
338func (r *udpEntryPointRuntime) drainSessionsAfter(drainTimeout time.Duration) {
339 trafficDebug(r.ctx).Str("entrypoint", r.entryPointSnapshot().Name).Dur("drain_timeout", drainTimeout).Msg("scheduled udp session drain after router removal")
340 r.drainSessionsMatchingAfter(func(*udpSession) bool { return true }, drainTimeout)
341}
342
343func (r *udpEntryPointRuntime) drainSessionsNotMatchingAfter(backend domain.TrafficBackend, drainTimeout time.Duration) {
344 trafficDebug(r.ctx).Str("entrypoint", r.entryPointSnapshot().Name).Dur("drain_timeout", drainTimeout).Msg("scheduled stale udp session drain after backend update")
345 r.drainSessionsMatchingAfter(func(session *udpSession) bool { return !trafficBackendEqual(session.backendRef, backend) }, drainTimeout)
346}
347
348func (r *udpEntryPointRuntime) drainSessionsMatchingAfter(match func(*udpSession) bool, drainTimeout time.Duration) {
349 go func() {
350 if drainTimeout <= 0 {
351 drainTimeout = defaultUDPOptions().DrainTimeout
352 }
353 select {
354 case <-time.After(drainTimeout):
355 r.closeSessionsMatching(match)
356 case <-r.ctx.Done():
357 }
358 }()
359}
360
361func (r *udpEntryPointRuntime) waitSessions(ctx context.Context, drainTimeout time.Duration) bool {
362 if drainTimeout <= 0 {
363 drainTimeout = defaultUDPOptions().DrainTimeout
364 }
365 drainCtx, cancel := context.WithTimeout(ctx, drainTimeout)
366 defer cancel()
367 select {
368 case <-r.sessionsDone():
369 return true
370 case <-drainCtx.Done():
371 return false
372 }
373}
374
375func (r *udpEntryPointRuntime) closeDone() {
376 r.doneOnce.Do(func() { close(r.done) })
377}
378
379func (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
396func (r *udpEntryPointRuntime) closeSessions() {
397 r.closeSessionsMatching(func(*udpSession) bool { return true })
398}
399
400func (r *udpEntryPointRuntime) closeSessionsMatching(match func(*udpSession) bool) {
401 r.mu.Lock()
402 keys := make([]string, 0, len(r.sessions))
403 for key, session := range r.sessions {
404 if match(session) {
405 keys = append(keys, key)
406 }
407 }
408 r.mu.Unlock()
409 for _, key := range keys {
410 r.removeSession(key)
411 }
412}
413
414func (r *udpEntryPointRuntime) matches(entryPoint domain.EntryPoint) bool {
415 current := r.entryPointSnapshot()
416 return current.Name == entryPoint.Name && current.Address == entryPoint.Address && current.Protocol == entryPoint.Protocol && trustedCIDRsEqual(current.TrustedCIDRs, entryPoint.TrustedCIDRs)
417}
418
419func (r *udpEntryPointRuntime) sameAddress(entryPoint domain.EntryPoint) bool {
420 return r.entryPointSnapshot().Address == entryPoint.Address
421}
422
423func (r *udpEntryPointRuntime) updateEntryPoint(entryPoint domain.EntryPoint, trusted []*net.IPNet) {
424 r.mu.Lock()
425 r.entryPoint = entryPoint
426 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)
431 }
432 }
433 r.mu.Unlock()
434 for _, key := range stale {
435 r.removeSession(key)
436 }
437}
438
439func (r *udpEntryPointRuntime) entryPointSnapshot() domain.EntryPoint {
440 r.mu.Lock()
441 defer r.mu.Unlock()
442 return r.entryPoint
443}
444
445func (r *udpEntryPointRuntime) trustedSnapshot() []*net.IPNet {
446 r.mu.Lock()
447 defer r.mu.Unlock()
448 return r.trusted
449}
450
451func (r *udpEntryPointRuntime) isClosed() bool { return r.closed.Load() }
452
453func (s *udpSession) touch() { s.lastSeen.Store(time.Now().UnixNano()) }
454
455func (s *udpSession) close() {
456 s.once.Do(func() {
457 _ = s.backend.Close()
458 close(s.done)
459 })
460}