1package quicnettest
2
3import (
4 "container/heap"
5 "context"
6 "errors"
7 "fmt"
8 "net"
9 "sync"
10 "sync/atomic"
11 "time"
12)
13
14// ErrClosed is returned by operations on a closed proxy.
15var ErrClosed = errors.New("quicnettest: proxy is closed")
16
17// Proxy is a bounded, deterministic UDP relay with per-direction impairment. It
18// listens on an ephemeral loopback port for clients and forwards to
19// Config.ServerAddr from a separate upstream socket that can be rebound to a
20// new source address at any time.
21type Proxy struct {
22 cfg Config
23 clock Clock
24 server *net.UDPAddr
25
26 clientConn *net.UDPConn
27 upstream atomic.Pointer[net.UDPConn]
28
29 done chan struct{}
30 wg sync.WaitGroup
31 once sync.Once
32 closeErr error
33
34 // slots bounds queued plus deadline-pending packets; queue only hands
35 // packets to the scheduler and is drained eagerly.
36 slots chan struct{}
37 queue chan packet
38
39 start time.Time
40 seq atomic.Uint64
41
42 toServer impairer
43 toClient impairer
44
45 toServerCounters linkCounters
46 toClientCounters linkCounters
47
48 strayPackets atomic.Uint64
49 rebinds atomic.Uint64
50 rebindErrors atomic.Uint64
51 sentToServer atomic.Uint64
52
53 rebindMu sync.Mutex
54 clientAddr atomic.Pointer[net.UDPAddr]
55 manualBlackout atomic.Bool
56
57 onPacket func(Event)
58}
59
60type linkCounters struct {
61 received atomic.Uint64
62 sent atomic.Uint64
63 lossDrops atomic.Uint64
64 blackoutDrops atomic.Uint64
65 oversizeDrops atomic.Uint64
66 overflowDrops atomic.Uint64
67 duplicates atomic.Uint64
68 reordered atomic.Uint64
69 deliveryErrors atomic.Uint64
70}
71
72func (c *linkCounters) snapshot() LinkStats {
73 return LinkStats{
74 Received: c.received.Load(),
75 Sent: c.sent.Load(),
76 LossDrops: c.lossDrops.Load(),
77 BlackoutDrops: c.blackoutDrops.Load(),
78 OversizeDrops: c.oversizeDrops.Load(),
79 OverflowDrops: c.overflowDrops.Load(),
80 Duplicates: c.duplicates.Load(),
81 Reordered: c.reordered.Load(),
82 DeliveryErrors: c.deliveryErrors.Load(),
83 }
84}
85
86type packet struct {
87 data []byte
88 dst *net.UDPAddr
89 toServer bool
90 due time.Time
91 seq uint64
92}
93
94type packetHeap []packet
95
96func (h packetHeap) Len() int { return len(h) }
97
98func (h packetHeap) Less(i, j int) bool {
99 if h[i].due.Equal(h[j].due) {
100 return h[i].seq < h[j].seq
101 }
102 return h[i].due.Before(h[j].due)
103}
104
105func (h packetHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
106
107func (h *packetHeap) Push(value any) {
108 pkt, ok := value.(packet)
109 if !ok {
110 panic(fmt.Sprintf("quicnettest: packet heap received %T", value))
111 }
112 *h = append(*h, pkt)
113}
114
115func (h *packetHeap) Pop() any {
116 old := *h
117 n := len(old)
118 pkt := old[n-1]
119 old[n-1] = packet{}
120 *h = old[:n-1]
121 return pkt
122}
123
124// New starts a proxy in front of cfg.ServerAddr and returns it. Callers must
125// Close the proxy to release its three goroutines and two sockets.
126func New(cfg Config) (*Proxy, error) {
127 cfg = cfg.withDefaults()
128 if err := cfg.validate(); err != nil {
129 return nil, err
130 }
131 clientConn, err := listenLoopback()
132 if err != nil {
133 return nil, fmt.Errorf("quicnettest: listen client socket: %w", err)
134 }
135 upstream, err := listenLoopback()
136 if err != nil {
137 _ = clientConn.Close()
138 return nil, fmt.Errorf("quicnettest: listen upstream socket: %w", err)
139 }
140 start := cfg.Clock.Now()
141 var shared RNG
142 if cfg.RNG != nil {
143 shared = &lockedRNG{rng: cfg.RNG}
144 }
145 rngFor := func(direction Direction) RNG {
146 if shared != nil {
147 return shared
148 }
149 return newRNG(cfg.Seed, direction)
150 }
151 p := &Proxy{
152 cfg: cfg,
153 clock: cfg.Clock,
154 server: cfg.ServerAddr,
155 clientConn: clientConn,
156 done: make(chan struct{}),
157 slots: make(chan struct{}, cfg.QueueCapacity),
158 queue: make(chan packet, cfg.QueueCapacity),
159 start: start,
160 toServer: impairer{link: cfg.ToServer, rng: rngFor(ToServer), start: start},
161 toClient: impairer{link: cfg.ToClient, rng: rngFor(ToClient), start: start},
162 onPacket: cfg.OnPacket,
163 }
164 p.upstream.Store(upstream)
165 p.wg.Add(3)
166 go p.runClientReader()
167 go p.runServerReader()
168 go p.runScheduler()
169 return p, nil
170}
171
172// Addr returns the client-facing address. Sending here is equivalent to sending
173// to the impaired server.
174func (p *Proxy) Addr() *net.UDPAddr {
175 addr, ok := p.clientConn.LocalAddr().(*net.UDPAddr)
176 if !ok {
177 return nil
178 }
179 return &net.UDPAddr{IP: append(net.IP(nil), addr.IP...), Port: addr.Port}
180}
181
182// UpstreamAddr returns the current server-facing local address, or nil once the
183// proxy is closed. It changes on every successful Rebind.
184func (p *Proxy) UpstreamAddr() *net.UDPAddr {
185 conn := p.upstream.Load()
186 if conn == nil {
187 return nil
188 }
189 addr, ok := conn.LocalAddr().(*net.UDPAddr)
190 if !ok {
191 return nil
192 }
193 return &net.UDPAddr{IP: append(net.IP(nil), addr.IP...), Port: addr.Port}
194}
195
196// SetBlackout toggles a manual blackout applied on top of the scheduled
197// windows. It is safe to call from any goroutine.
198func (p *Proxy) SetBlackout(active bool) { p.manualBlackout.Store(active) }
199
200// Rebind replaces the upstream socket, changing the source address the server
201// observes. In-flight packets that the server sends to the previous address are
202// lost, which mirrors a NAT rebinding.
203func (p *Proxy) Rebind() error {
204 p.rebindMu.Lock()
205 defer p.rebindMu.Unlock()
206 if p.stopping() {
207 return ErrClosed
208 }
209 conn, err := listenLoopback()
210 if err != nil {
211 return fmt.Errorf("quicnettest: rebind: %w", err)
212 }
213 if p.stopping() {
214 _ = conn.Close()
215 return ErrClosed
216 }
217 old := p.upstream.Swap(conn)
218 if old != nil {
219 _ = old.Close()
220 }
221 p.rebinds.Add(1)
222 return nil
223}
224
225// Stats returns a snapshot of the proxy counters.
226func (p *Proxy) Stats() Stats {
227 return Stats{
228 ToServer: p.toServerCounters.snapshot(),
229 ToClient: p.toClientCounters.snapshot(),
230 StrayPackets: p.strayPackets.Load(),
231 Rebinds: p.rebinds.Load(),
232 RebindErrors: p.rebindErrors.Load(),
233 }
234}
235
236// Close stops the proxy and waits for its goroutines to exit. It is idempotent
237// and returns the same error on every call.
238func (p *Proxy) Close() error {
239 p.once.Do(func() {
240 close(p.done)
241 var errs []error
242 if err := p.clientConn.Close(); err != nil {
243 errs = append(errs, err)
244 }
245 p.rebindMu.Lock()
246 if upstream := p.upstream.Swap(nil); upstream != nil {
247 if err := upstream.Close(); err != nil {
248 errs = append(errs, err)
249 }
250 }
251 p.rebindMu.Unlock()
252 p.wg.Wait()
253 p.closeErr = errors.Join(errs...)
254 })
255 return p.closeErr
256}
257
258func (p *Proxy) stopping() bool {
259 select {
260 case <-p.done:
261 return true
262 default:
263 return false
264 }
265}
266
267func (p *Proxy) counters(toServer bool) *linkCounters {
268 if toServer {
269 return &p.toServerCounters
270 }
271 return &p.toClientCounters
272}
273
274func (p *Proxy) linkFor(toServer bool) LinkConfig {
275 if toServer {
276 return p.cfg.ToServer
277 }
278 return p.cfg.ToClient
279}
280
281func (p *Proxy) impairerFor(toServer bool) *impairer {
282 if toServer {
283 return &p.toServer
284 }
285 return &p.toClient
286}
287
288func (p *Proxy) runClientReader() {
289 defer p.wg.Done()
290 buf := make([]byte, p.cfg.MaxPacketSize+1)
291 for {
292 n, from, err := p.clientConn.ReadFromUDP(buf)
293 if err != nil {
294 return
295 }
296 p.clientAddr.Store(from)
297 p.ingest(true, buf[:n], p.server)
298 }
299}
300
301func (p *Proxy) runServerReader() {
302 defer p.wg.Done()
303 buf := make([]byte, p.cfg.MaxPacketSize+1)
304 for {
305 conn := p.upstream.Load()
306 if conn == nil {
307 return
308 }
309 n, from, err := conn.ReadFromUDP(buf)
310 if err != nil {
311 if p.stopping() {
312 return
313 }
314 if conn != p.upstream.Load() {
315 // The socket was rebound; continue on the new one.
316 continue
317 }
318 return
319 }
320 if !sameUDPAddr(from, p.server) {
321 p.strayPackets.Add(1)
322 continue
323 }
324 client := p.clientAddr.Load()
325 if client == nil {
326 p.strayPackets.Add(1)
327 continue
328 }
329 p.ingest(false, buf[:n], client)
330 }
331}
332
333func (p *Proxy) runScheduler() {
334 defer p.wg.Done()
335 var pending packetHeap
336 heap.Init(&pending)
337 timer := p.clock.NewTimer(time.Hour)
338 if !timer.Stop() {
339 <-timer.C()
340 }
341 defer timer.Stop()
342 for {
343 // Drain everything already queued before evaluating deadlines. A test
344 // clock jump must observe all queued packets, otherwise a later packet
345 // with an earlier deadline could be delivered behind a later one.
346 p.drainQueue(&pending)
347 if len(pending) == 0 {
348 select {
349 case <-p.done:
350 return
351 case pkt := <-p.queue:
352 heap.Push(&pending, pkt)
353 }
354 continue
355 }
356 now := p.clock.Now()
357 if !pending[0].due.After(now) {
358 p.deliverDue(&pending, now)
359 continue
360 }
361 if !timer.Stop() {
362 select {
363 case <-timer.C():
364 default:
365 }
366 }
367 timer.ResetAt(pending[0].due)
368 select {
369 case <-p.done:
370 return
371 case pkt := <-p.queue:
372 heap.Push(&pending, pkt)
373 case <-timer.C():
374 }
375 }
376}
377
378func (p *Proxy) drainQueue(pending *packetHeap) {
379 for {
380 select {
381 case pkt := <-p.queue:
382 heap.Push(pending, pkt)
383 default:
384 return
385 }
386 }
387}
388
389func (p *Proxy) deliverDue(pending *packetHeap, now time.Time) {
390 for pending.Len() > 0 && !(*pending)[0].due.After(now) {
391 pkt, ok := heap.Pop(pending).(packet)
392 if !ok {
393 return
394 }
395 p.deliver(pkt)
396 <-p.slots
397 }
398}
399
400// ingest applies impairments to one freshly read datagram and schedules its
401// copies. It never blocks: a full budget drops the packet instead.
402func (p *Proxy) ingest(toServer bool, raw []byte, dst *net.UDPAddr) {
403 counters := p.counters(toServer)
404 seq := p.seq.Add(1)
405 now := p.clock.Now()
406 if len(raw) > p.cfg.MaxPacketSize {
407 counters.oversizeDrops.Add(1)
408 p.emit(Event{
409 Direction: directionOf(toServer),
410 Sequence: seq,
411 Bytes: len(raw),
412 At: now,
413 Decision: Decision{Drop: true, Oversize: true},
414 })
415 return
416 }
417 counters.received.Add(1)
418
419 var decision Decision
420 if p.manualBlackout.Load() {
421 decision = Decision{Drop: true, Blackout: true}
422 } else {
423 decision = p.impairerFor(toServer).decide(now)
424 }
425 if decision.Drop {
426 if decision.Blackout {
427 counters.blackoutDrops.Add(1)
428 } else {
429 counters.lossDrops.Add(1)
430 }
431 p.emit(Event{
432 Direction: directionOf(toServer),
433 Sequence: seq,
434 Bytes: len(raw),
435 At: now,
436 Decision: decision,
437 })
438 return
439 }
440
441 link := p.linkFor(toServer)
442 if decision.Duplicates > 0 {
443 counters.duplicates.Add(uint64(decision.Duplicates))
444 }
445 if decision.Reordered {
446 counters.reordered.Add(1)
447 }
448 due := now.Add(link.BaseLatency + decision.ExtraDelay)
449 for copyIndex := 0; copyIndex <= decision.Duplicates; copyIndex++ {
450 copyDue := due.Add(time.Duration(copyIndex) * link.DuplicateDelay)
451 if !p.enqueue(packet{
452 data: append([]byte(nil), raw...),
453 dst: dst,
454 toServer: toServer,
455 due: copyDue,
456 seq: seq,
457 }) {
458 counters.overflowDrops.Add(1)
459 }
460 }
461 p.emit(Event{
462 Direction: directionOf(toServer),
463 Sequence: seq,
464 Bytes: len(raw),
465 At: now,
466 Due: due,
467 Decision: decision,
468 })
469}
470
471func (p *Proxy) enqueue(pkt packet) bool {
472 select {
473 case p.slots <- struct{}{}:
474 case <-p.done:
475 return false
476 default:
477 return false
478 }
479 select {
480 case p.queue <- pkt:
481 return true
482 case <-p.done:
483 <-p.slots
484 return false
485 }
486}
487
488func (p *Proxy) deliver(pkt packet) {
489 counters := p.counters(pkt.toServer)
490 conn := p.clientConn
491 if pkt.toServer {
492 conn = p.upstream.Load()
493 }
494 if conn == nil {
495 if !p.stopping() {
496 counters.deliveryErrors.Add(1)
497 }
498 return
499 }
500 // Count before writing so that a peer observing the packet implies the
501 // counter has been updated.
502 counters.sent.Add(1)
503 if _, err := conn.WriteToUDP(pkt.data, pkt.dst); err != nil {
504 if !p.stopping() {
505 counters.deliveryErrors.Add(1)
506 }
507 return
508 }
509 if pkt.toServer && p.cfg.NATRebindAfterPackets > 0 {
510 every := uint64(p.cfg.NATRebindAfterPackets)
511 if sent := p.sentToServer.Add(1); sent%every == 0 {
512 if err := p.Rebind(); err != nil {
513 p.rebindErrors.Add(1)
514 }
515 }
516 }
517}
518
519func (p *Proxy) emit(event Event) {
520 if p.onPacket == nil {
521 return
522 }
523 p.onPacket(event)
524}
525
526func directionOf(toServer bool) Direction {
527 if toServer {
528 return ToServer
529 }
530 return ToClient
531}
532
533func listenLoopback() (*net.UDPConn, error) {
534 var lc net.ListenConfig
535 conn, err := lc.ListenPacket(context.Background(), "udp", "127.0.0.1:0")
536 if err != nil {
537 return nil, err
538 }
539 udp, ok := conn.(*net.UDPConn)
540 if !ok {
541 _ = conn.Close()
542 return nil, fmt.Errorf("quicnettest: udp listener returned %T", conn)
543 }
544 return udp, nil
545}
546
547func sameUDPAddr(a, b *net.UDPAddr) bool {
548 return a != nil && b != nil && a.Port == b.Port && a.IP.Equal(b.IP)
549}