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