1package traffic
2
3import (
4 "context"
5 "net"
6 "testing"
7 "time"
8
9 "github.com/stretchr/testify/assert"
10 "github.com/stretchr/testify/require"
11
12 "github.com/bnema/gordon/internal/domain"
13)
14
15func TestUDPPassthroughEcho(t *testing.T) {
16 backend := startUDPEchoServer(t)
17 graph := udpGraph(t, freeUDPAddress(t), backend.address)
18 manager := NewManager()
19 require.NoError(t, manager.Apply(context.Background(), &graph))
20 defer shutdownManager(t, manager)
21
22 conn := dialUDP(t, graph.EntryPoints[0].Address)
23 defer conn.Close()
24 assertUDPRoundTrip(t, conn, "hello udp")
25}
26
27func TestUDPRejectsUntrustedCIDR(t *testing.T) {
28 backend := startUDPEchoServer(t)
29 graph := udpGraph(t, freeUDPAddress(t), backend.address)
30 graph.EntryPoints[0].TrustedCIDRs = []string{"192.0.2.0/24"}
31 manager := NewManager()
32 require.NoError(t, manager.Apply(context.Background(), &graph))
33 defer shutdownManager(t, manager)
34
35 assertUDPRejected(t, manager, graph.EntryPoints[0].Address, 1)
36}
37
38func TestUDPReloadAppliesTrustedCIDRChange(t *testing.T) {
39 backend := startUDPEchoServer(t)
40 graph := udpGraph(t, freeUDPAddress(t), backend.address)
41 graph.EntryPoints[0].TrustedCIDRs = []string{"127.0.0.0/8"}
42 manager := NewManager()
43 require.NoError(t, manager.Apply(context.Background(), &graph))
44 defer shutdownManager(t, manager)
45
46 conn := dialUDP(t, graph.EntryPoints[0].Address)
47 assertUDPRoundTrip(t, conn, "allowed")
48 require.NoError(t, conn.Close())
49
50 updated := graph
51 updated.EntryPoints[0].TrustedCIDRs = []string{"192.0.2.0/24"}
52 require.NoError(t, manager.Apply(context.Background(), &updated))
53 assertUDPRejected(t, manager, graph.EntryPoints[0].Address, 1)
54}
55
56func TestUDPReloadClosesActiveUntrustedSessions(t *testing.T) {
57 backend := startUDPEchoServer(t)
58 graph := udpGraph(t, freeUDPAddress(t), backend.address)
59 graph.EntryPoints[0].TrustedCIDRs = []string{"127.0.0.0/8"}
60 manager := NewManager()
61 require.NoError(t, manager.Apply(context.Background(), &graph))
62 defer shutdownManager(t, manager)
63
64 conn := dialUDP(t, graph.EntryPoints[0].Address)
65 defer conn.Close()
66 assertUDPRoundTrip(t, conn, "before")
67 require.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 1 }, time.Second, 10*time.Millisecond)
68
69 updated := graph
70 updated.EntryPoints[0].TrustedCIDRs = []string{"192.0.2.0/24"}
71 require.NoError(t, manager.Apply(context.Background(), &updated))
72
73 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 0 }, time.Second, 10*time.Millisecond)
74 _, err := conn.Write([]byte("after"))
75 require.NoError(t, err)
76 _ = conn.SetReadDeadline(time.Now().Add(80 * time.Millisecond))
77 buf := make([]byte, 32)
78 _, err = conn.Read(buf)
79 require.Error(t, err)
80}
81
82func TestUDPReloadReplacesSessionWhenBackendChanges(t *testing.T) {
83 backendA := startUDPEchoServerWithPrefix(t, "a:")
84 backendB := startUDPEchoServerWithPrefix(t, "b:")
85 graph := udpGraph(t, freeUDPAddress(t), backendA.address)
86 manager := NewManager()
87 require.NoError(t, manager.Apply(context.Background(), &graph))
88 defer shutdownManager(t, manager)
89
90 conn := dialUDP(t, graph.EntryPoints[0].Address)
91 defer conn.Close()
92 assertUDPRoundTripWant(t, conn, "first", "a:first")
93
94 updated := graph
95 backend, err := backendFromAddress("echo:udp", backendB.address)
96 require.NoError(t, err)
97 backend.Protocol = domain.NetworkProtocolUDP
98 updated.Services[0].Backends = []domain.TrafficBackend{backend}
99 require.NoError(t, manager.Apply(context.Background(), &updated))
100
101 assertUDPRoundTripWant(t, conn, "second", "a:second")
102 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 0 }, time.Second, 10*time.Millisecond)
103
104 newConn := dialUDP(t, graph.EntryPoints[0].Address)
105 defer newConn.Close()
106 assertUDPRoundTripWant(t, newConn, "third", "b:third")
107}
108
109func assertUDPRejected(t *testing.T, manager *Manager, address string, refused int64) {
110 t.Helper()
111 conn := dialUDP(t, address)
112 defer conn.Close()
113 _, err := conn.Write([]byte("blocked"))
114 require.NoError(t, err)
115 _ = conn.SetReadDeadline(time.Now().Add(80 * time.Millisecond))
116 buf := make([]byte, 32)
117 _, err = conn.Read(buf)
118 require.Error(t, err)
119 assert.Eventually(t, func() bool { return manager.Status().Counters.TotalRefused == refused }, time.Second, 10*time.Millisecond)
120}
121
122func TestUDPTwoClientsGetIsolatedSessions(t *testing.T) {
123 backend := startUDPEchoServer(t)
124 graph := udpGraph(t, freeUDPAddress(t), backend.address)
125 manager := NewManager()
126 require.NoError(t, manager.Apply(context.Background(), &graph))
127 defer shutdownManager(t, manager)
128
129 first := dialUDP(t, graph.EntryPoints[0].Address)
130 defer first.Close()
131 second := dialUDP(t, graph.EntryPoints[0].Address)
132 defer second.Close()
133 assertUDPRoundTrip(t, first, "first")
134 assertUDPRoundTrip(t, second, "second")
135 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 2 }, time.Second, 10*time.Millisecond)
136}
137
138func TestUDPIdleSessionExpires(t *testing.T) {
139 backend := startUDPEchoServer(t)
140 graph := udpGraph(t, freeUDPAddress(t), backend.address)
141 graph.Options.UDP.IdleTimeout = 40 * time.Millisecond
142 manager := NewManager()
143 require.NoError(t, manager.Apply(context.Background(), &graph))
144 defer shutdownManager(t, manager)
145
146 conn := dialUDP(t, graph.EntryPoints[0].Address)
147 defer conn.Close()
148 assertUDPRoundTrip(t, conn, "expire")
149 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 0 }, time.Second, 10*time.Millisecond)
150}
151
152func TestUDPMaxSessionsRejectsOverflow(t *testing.T) {
153 backend := startUDPEchoServer(t)
154 graph := udpGraph(t, freeUDPAddress(t), backend.address)
155 graph.Options.UDP.MaxSessions = 1
156 manager := NewManager()
157 require.NoError(t, manager.Apply(context.Background(), &graph))
158 defer shutdownManager(t, manager)
159
160 first := dialUDP(t, graph.EntryPoints[0].Address)
161 defer first.Close()
162 assertUDPRoundTrip(t, first, "first")
163
164 second := dialUDP(t, graph.EntryPoints[0].Address)
165 defer second.Close()
166 _, err := second.Write([]byte("second"))
167 require.NoError(t, err)
168 _ = second.SetReadDeadline(time.Now().Add(80 * time.Millisecond))
169 buf := make([]byte, 32)
170 _, err = second.Read(buf)
171 require.Error(t, err)
172 assert.Eventually(t, func() bool { return manager.Status().Counters.TotalRefused == 1 }, time.Second, 10*time.Millisecond)
173}
174
175func TestUDPRemovedRouterWithRetainedEntryPointDrainsSession(t *testing.T) {
176 backend := startUDPEchoServer(t)
177 graph := udpGraph(t, freeUDPAddress(t), backend.address)
178 graph.Options.UDP.DrainTimeout = 50 * time.Millisecond
179 manager := NewManager()
180 require.NoError(t, manager.Apply(context.Background(), &graph))
181 defer shutdownManager(t, manager)
182
183 conn := dialUDP(t, graph.EntryPoints[0].Address)
184 defer conn.Close()
185 assertUDPRoundTrip(t, conn, "held")
186
187 next := graph
188 next.Routers = nil
189 next.Services = nil
190 require.NoError(t, manager.Apply(context.Background(), &next))
191 assertUDPRoundTrip(t, conn, "during-drain")
192 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 0 }, time.Second, 10*time.Millisecond)
193
194 _, err := conn.Write([]byte("after"))
195 require.NoError(t, err)
196 _ = conn.SetReadDeadline(time.Now().Add(80 * time.Millisecond))
197 buf := make([]byte, 32)
198 _, err = conn.Read(buf)
199 require.Error(t, err)
200}
201
202func TestUDPBackendChangeDrainsExistingSession(t *testing.T) {
203 oldBackend := startUDPEchoServerWithPrefix(t, "old:")
204 newBackend := startUDPEchoServerWithPrefix(t, "new:")
205 graph := udpGraph(t, freeUDPAddress(t), oldBackend.address)
206 graph.Options.UDP.DrainTimeout = 50 * time.Millisecond
207 manager := NewManager()
208 require.NoError(t, manager.Apply(context.Background(), &graph))
209 defer shutdownManager(t, manager)
210
211 oldConn := dialUDP(t, graph.EntryPoints[0].Address)
212 defer oldConn.Close()
213 assertUDPRoundTripWant(t, oldConn, "held", "old:held")
214
215 next := udpGraph(t, graph.EntryPoints[0].Address, newBackend.address)
216 next.Options = graph.Options
217 require.NoError(t, manager.Apply(context.Background(), &next))
218 assertUDPRoundTripWant(t, oldConn, "during-drain", "old:during-drain")
219 assert.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 0 }, time.Second, 10*time.Millisecond)
220
221 newConn := dialUDP(t, graph.EntryPoints[0].Address)
222 defer newConn.Close()
223 assertUDPRoundTripWant(t, newConn, "after", "new:after")
224}
225
226func TestUDPRemovedRouterDrainsThenClosesSessions(t *testing.T) {
227 backend := startUDPEchoServer(t)
228 graph := udpGraph(t, freeUDPAddress(t), backend.address)
229 graph.Options.UDP.DrainTimeout = 50 * time.Millisecond
230 manager := NewManager()
231 require.NoError(t, manager.Apply(context.Background(), &graph))
232 defer shutdownManager(t, manager)
233
234 conn := dialUDP(t, graph.EntryPoints[0].Address)
235 defer conn.Close()
236 assertUDPRoundTrip(t, conn, "held")
237 require.Eventually(t, func() bool { return manager.Status().Counters.ActiveUDPSessions == 1 }, time.Second, 10*time.Millisecond)
238
239 empty := domain.TrafficGraph{Options: graph.Options}
240 started := time.Now()
241 require.NoError(t, manager.Apply(context.Background(), &empty))
242 assert.GreaterOrEqual(t, time.Since(started), 40*time.Millisecond)
243 assert.Equal(t, int64(0), manager.Status().Counters.ActiveUDPSessions)
244}
245
246func TestUDPShutdownClosesSocket(t *testing.T) {
247 backend := startUDPEchoServer(t)
248 graph := udpGraph(t, freeUDPAddress(t), backend.address)
249 manager := NewManager()
250 require.NoError(t, manager.Apply(context.Background(), &graph))
251 require.NoError(t, manager.Shutdown(context.Background()))
252
253 conn := dialUDP(t, graph.EntryPoints[0].Address)
254 defer conn.Close()
255 _, err := conn.Write([]byte("closed"))
256 require.NoError(t, err)
257 _ = conn.SetReadDeadline(time.Now().Add(80 * time.Millisecond))
258 buf := make([]byte, 32)
259 _, err = conn.Read(buf)
260 require.Error(t, err)
261}
262
263func TestUDPStatusCountersTrackDatagrams(t *testing.T) {
264 backend := startUDPEchoServer(t)
265 graph := udpGraph(t, freeUDPAddress(t), backend.address)
266 manager := NewManager()
267 require.NoError(t, manager.Apply(context.Background(), &graph))
268 defer shutdownManager(t, manager)
269
270 conn := dialUDP(t, graph.EntryPoints[0].Address)
271 defer conn.Close()
272 assertUDPRoundTrip(t, conn, "bytes")
273
274 assert.Eventually(t, func() bool {
275 status := manager.Status()
276 return status.Counters.TotalAccepted == 1 &&
277 status.Counters.ActiveUDPSessions == 1 &&
278 status.Counters.BytesIn >= int64(len("bytes")) &&
279 status.Counters.BytesOut >= int64(len("bytes"))
280 }, time.Second, 10*time.Millisecond)
281}
282
283func udpGraph(t *testing.T, listenAddress string, backendAddress string) domain.TrafficGraph {
284 t.Helper()
285 backend, err := backendFromAddress("echo:udp", backendAddress)
286 require.NoError(t, err)
287 backend.Protocol = domain.NetworkProtocolUDP
288 ref := serviceRef("echo", "udp")
289 graph := domain.TrafficGraph{
290 Options: domain.TrafficOptions{UDP: domain.UDPOptions{IdleTimeout: time.Minute, DrainTimeout: 50 * time.Millisecond}},
291 EntryPoints: []domain.EntryPoint{{Name: "udp", Address: listenAddress, Protocol: domain.EntryPointProtocolUDP}},
292 Routers: []domain.TrafficRouter{{Name: "echo", EntryPoint: "udp", Protocol: domain.RouterProtocolUDP, Service: ref}},
293 Services: []domain.TrafficService{{Name: ref, Backends: []domain.TrafficBackend{backend}}},
294 }
295 require.NoError(t, graph.Validate())
296 return graph
297}
298
299func freeUDPAddress(t *testing.T) string {
300 t.Helper()
301 conn, err := net.ListenPacket("udp", "127.0.0.1:0")
302 require.NoError(t, err)
303 address := conn.LocalAddr().String()
304 require.NoError(t, conn.Close())
305 return address
306}
307
308func dialUDP(t *testing.T, address string) *net.UDPConn {
309 t.Helper()
310 udpAddr, err := net.ResolveUDPAddr("udp", address)
311 require.NoError(t, err)
312 conn, err := net.DialUDP("udp", nil, udpAddr)
313 require.NoError(t, err)
314 return conn
315}
316
317func assertUDPRoundTrip(t *testing.T, conn *net.UDPConn, message string) {
318 t.Helper()
319 assertUDPRoundTripWant(t, conn, message, message)
320}
321
322func assertUDPRoundTripWant(t *testing.T, conn *net.UDPConn, message string, want string) {
323 t.Helper()
324 _, err := conn.Write([]byte(message))
325 require.NoError(t, err)
326 _ = conn.SetReadDeadline(time.Now().Add(time.Second))
327 buf := make([]byte, 128)
328 n, err := conn.Read(buf)
329 require.NoError(t, err)
330 assert.Equal(t, want, string(buf[:n]))
331}
332
333type udpEchoServer struct{ address string }
334
335func startUDPEchoServer(t *testing.T) udpEchoServer {
336 t.Helper()
337 return startUDPEchoServerWithPrefix(t, "")
338}
339
340func startUDPEchoServerWithPrefix(t *testing.T, prefix string) udpEchoServer {
341 t.Helper()
342 conn, err := net.ListenPacket("udp", "127.0.0.1:0")
343 require.NoError(t, err)
344 done := make(chan struct{})
345 go func() {
346 defer close(done)
347 buf := make([]byte, udpBufferSize)
348 for {
349 n, addr, err := conn.ReadFrom(buf)
350 if err != nil {
351 return
352 }
353 _, _ = conn.WriteTo(append([]byte(prefix), buf[:n]...), addr)
354 }
355 }()
356 t.Cleanup(func() {
357 _ = conn.Close()
358 <-done
359 })
360 return udpEchoServer{address: conn.LocalAddr().String()}
361}