From d82258f844ef65ba06ba05aa7de4130c43f2c56e Mon Sep 17 00:00:00 2001 From: samosvalishe Date: Thu, 21 May 2026 01:25:59 +0300 Subject: [PATCH] =?UTF-8?q?refactor(udprelay):=20=D1=80=D0=B0=D0=B7=D0=B1?= =?UTF-8?q?=D0=B8=D1=82=D1=8C=20=D0=BD=D0=B0=20run/loop/listener?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Разбивает монолитный udprelay.go на 530 строк на три файла: - run.go: Run, Deps, ErrFatal, Params, AuthHandler, GetCredsFunc - loop.go: DTLSLoop, TURNLoop, oneDTLS, oneTURN - listener.go: runListener, Packet, packetPool, inboundQueueCap Поведение не изменено; чисто структурный рефакторинг. --- internal/proxy/udprelay/listener.go | 65 ++++++ .../proxy/udprelay/{udprelay.go => loop.go} | 199 +----------------- internal/proxy/udprelay/run.go | 153 ++++++++++++++ 3 files changed, 219 insertions(+), 198 deletions(-) create mode 100644 internal/proxy/udprelay/listener.go rename internal/proxy/udprelay/{udprelay.go => loop.go} (58%) create mode 100644 internal/proxy/udprelay/run.go diff --git a/internal/proxy/udprelay/listener.go b/internal/proxy/udprelay/listener.go new file mode 100644 index 0000000..927ac3d --- /dev/null +++ b/internal/proxy/udprelay/listener.go @@ -0,0 +1,65 @@ +package udprelay + +import ( + "context" + "net" + "sync" + "sync/atomic" +) + +const inboundQueueCap = 2000 + +// Packet is a pooled UDP datagram carried from the listener to the per-stream +// DTLS worker. N is the populated prefix of Data. +type Packet struct { + Data []byte + N int +} + +// packetPool reuses Packet buffers across the inbound hot path. Buffer size +// matches the 2048-byte default the listener loop expects. +var packetPool = sync.Pool{ + New: func() any { return &Packet{Data: make([]byte, 2048)} }, +} + +// runListener reads packets from listenConn, refreshes the active-peer cache, +// and posts each packet to inboundChan. Packets are dropped when the channel +// is full to keep the read loop wait-free. +func runListener(ctx context.Context, listenConn net.PacketConn, activeLocalPeer *atomic.Value, inboundChan chan<- *Packet) { + // Pointer-cache for the last seen local peer addr. Avoids the + // per-packet addr.String() allocation pair on the hot WG ingest path: + // most packets come from the same UDPAddr instance, so a pointer + // equality check covers the fast path. The slow path (new instance + // from ReadFrom for the same ip:port) does one String compare and + // then refreshes the cache. + var lastAddr net.Addr + var lastAddrStr string + for { + if ctx.Err() != nil { + return + } + pktIface := packetPool.Get() + pkt := pktIface.(*Packet) //nolint:errcheck // pool New always returns *Packet + nRead, addr, err := listenConn.ReadFrom(pkt.Data) + if err != nil { + return + } + + if addr != lastAddr { + s := addr.String() + if s != lastAddrStr { + activeLocalPeer.Store(addr) + lastAddrStr = s + } + lastAddr = addr + } + + pkt.N = nRead + + select { + case inboundChan <- pkt: + default: + packetPool.Put(pkt) + } + } +} diff --git a/internal/proxy/udprelay/udprelay.go b/internal/proxy/udprelay/loop.go similarity index 58% rename from internal/proxy/udprelay/udprelay.go rename to internal/proxy/udprelay/loop.go index fe7f57e..b43fe6a 100644 --- a/internal/proxy/udprelay/udprelay.go +++ b/internal/proxy/udprelay/loop.go @@ -1,7 +1,3 @@ -// Package udprelay implements the UDP-mode proxy loop: it terminates DTLS from a -// local peer (WireGuard) and relays its packets through a per-stream TURN -// allocation back to a remote peer. Run is the entrypoint; it owns the local -// listener, the inbound dispatch fan-in, and the per-stream DTLS/TURN loops. package udprelay import ( @@ -14,206 +10,13 @@ import ( "sync/atomic" "time" - "github.com/samosvalishe/btp/internal/client/vkauth" - "github.com/samosvalishe/btp/internal/logx" "github.com/samosvalishe/btp/internal/proxy/common" "github.com/samosvalishe/btp/internal/stats" - "github.com/samosvalishe/btp/internal/transport/dtlsdial" "github.com/samosvalishe/btp/internal/wire/srtpmimicry" + "github.com/samosvalishe/btp/internal/client/vkauth" "github.com/cbeuw/connutil" ) -// Packet is a pooled UDP datagram carried from the listener to the per-stream -// DTLS worker. N is the populated prefix of Data. -type Packet struct { - Data []byte - N int -} - -// packetPool reuses Packet buffers across the inbound hot path. Buffer size -// matches the 2048-byte default the listener loop expects. -var packetPool = sync.Pool{ - New: func() any { return &Packet{Data: make([]byte, 2048)} }, -} - -// GetCredsFunc is re-exported from common so callers can keep their imports -// scoped to this package. -type GetCredsFunc = common.GetCredsFunc - -// AuthHandler is the subset of vkauth.Client this package needs. Defined as -// an interface so tests can inject fakes; the production wiring still imports -// vkauth for its sentinel errors (ErrFatalCaptchaNoStreams, etc.). -type AuthHandler interface { - IsAuthError(err error) bool - HandleAuthError(streamID int) bool - ResetErrors(streamID int) - LockoutUntilUnix() int64 -} - -// Params is the per-stream TURN/wrap configuration shared by the DTLS and TURN loops. -type Params struct { - Host string - Port string - Link string - UDP bool - WrapKey []byte - GetCreds GetCredsFunc -} - -// ErrFatal is returned by Run when a stream encounters a condition that -// requires the entire application to exit (e.g. manual captcha solver failed -// with no connected streams). Callers should check with errors.Is and call -// os.Exit themselves — udprelay does not reach into the host process. -var ErrFatal = errors.New("udprelay: fatal error") - -// Deps groups everything the loops need from the host process. The atomics -// are owned by Run and exposed here so DTLSLoop/TURNLoop can share them when -// called directly (Run wires them automatically). -type Deps struct { - DTLSDialer *dtlsdial.Dialer - Auth AuthHandler - Log logx.Logger - ActiveLocalPeer *atomic.Value - ConnectedStreams *atomic.Int32 - // fatalCh is an internal signalling channel; set by Run, written by - // TURNLoop, and drained by Run to propagate the fatal error up. - fatalCh chan error -} - -func (d *Deps) log() logx.Logger { - if d.Log == nil { - return logx.Nop() - } - return d.Log -} - -// Run is the UDP-mode entrypoint. It binds listenAddr, fans inbound packets -// into a shared queue, and spawns numStreams pairs of (DTLSLoop, TURNLoop). -// connectedStreams is owned by the caller (vkauth reads it via StreamsAlive) -// and incremented/decremented by oneTURN. -// Returns after all stream loops exit (i.e. when ctx is cancelled). -// If a fatal captcha condition is encountered, Run returns ErrFatal so the -// caller can perform os.Exit without udprelay reaching into the host process. -func Run(ctx context.Context, dtlsDialer *dtlsdial.Dialer, auth AuthHandler, logger logx.Logger, connectedStreams *atomic.Int32, params *Params, peer *net.UDPAddr, listenAddr string, numStreams int) error { - listenConn, err := net.ListenPacket("udp", listenAddr) - if err != nil { - return fmt.Errorf("udprelay listen %s: %w", listenAddr, err) - } - context.AfterFunc(ctx, func() { - if closeErr := listenConn.Close(); closeErr != nil { - logger.Errorf("udprelay: close local connection: %s", closeErr) - } - }) - - if numStreams <= 0 { - numStreams = 1 - } - - fatalCh := make(chan error, 1) - var activeLocalPeer atomic.Value - deps := &Deps{ - DTLSDialer: dtlsDialer, - Auth: auth, - Log: logger, - ActiveLocalPeer: &activeLocalPeer, - ConnectedStreams: connectedStreams, - fatalCh: fatalCh, - } - - // runCtx is cancelled when a fatal error is detected (via fatalCh), which - // propagates cancellation into all stream loops without requiring them to - // hold a reference to the host-process cancel function. - runCtx, runCancel := context.WithCancel(ctx) - defer runCancel() - - inboundChan := make(chan *Packet, inboundQueueCap) - wg := sync.WaitGroup{} - wg.Go(func() { - runListener(runCtx, listenConn, &activeLocalPeer, inboundChan) - }) - t := time.Tick(200 * time.Millisecond) - - // Stream 1 gets okchan so it can signal the first successful handshake to - // the log. All streams start concurrently — no gate between stream 1 and - // the rest, so a slow DTLS handshake on stream 1 never delays streams 2..N. - okchan := make(chan struct{}, 1) - for i := 0; i < numStreams; i++ { - cchan := make(chan net.PacketConn) - var ok chan<- struct{} - if i == 0 { - ok = okchan - } - streamID := i + 1 - wg.Go(func() { - DTLSLoop(runCtx, deps, peer, listenConn, inboundChan, cchan, ok, streamID) - }) - wg.Go(func() { - TURNLoop(runCtx, deps, params, peer, cchan, t, streamID) - }) - } - - // If a fatal error was sent, cancel remaining goroutines and propagate up. - var fatalErr error - go func() { - select { - case err := <-fatalCh: - fatalErr = err - runCancel() - case <-runCtx.Done(): - } - }() - - wg.Wait() - if fatalErr != nil { - return fatalErr - } - return nil -} - -const inboundQueueCap = 2000 - -// runListener reads packets from listenConn, refreshes the active-peer cache, -// and posts each packet to inboundChan. Packets are dropped when the channel -// is full to keep the read loop wait-free. -func runListener(ctx context.Context, listenConn net.PacketConn, activeLocalPeer *atomic.Value, inboundChan chan<- *Packet) { - // Pointer-cache for the last seen local peer addr. Avoids the - // per-packet addr.String() allocation pair on the hot WG ingest path: - // most packets come from the same UDPAddr instance, so a pointer - // equality check covers the fast path. The slow path (new instance - // from ReadFrom for the same ip:port) does one String compare and - // then refreshes the cache. - var lastAddr net.Addr - var lastAddrStr string - for { - if ctx.Err() != nil { - return - } - pktIface := packetPool.Get() - pkt := pktIface.(*Packet) //nolint:errcheck // pool New always returns *Packet - nRead, addr, err := listenConn.ReadFrom(pkt.Data) - if err != nil { - return - } - - if addr != lastAddr { - s := addr.String() - if s != lastAddrStr { - activeLocalPeer.Store(addr) - lastAddrStr = s - } - lastAddr = addr - } - - pkt.N = nRead - - select { - case inboundChan <- pkt: - default: - packetPool.Put(pkt) - } - } -} - // DTLSLoop keeps a single DTLS termination alive for streamID, restarting it // on failure with a 10-30s backoff (skipped while a captcha lockout is active // and the prior error was a deadline). connchan is fed a fresh AsyncPacketPipe diff --git a/internal/proxy/udprelay/run.go b/internal/proxy/udprelay/run.go new file mode 100644 index 0000000..5b53047 --- /dev/null +++ b/internal/proxy/udprelay/run.go @@ -0,0 +1,153 @@ +// Package udprelay implements the UDP-mode proxy loop: it terminates DTLS from a +// local peer (WireGuard) and relays its packets through a per-stream TURN +// allocation back to a remote peer. Run is the entrypoint; it owns the local +// listener, the inbound dispatch fan-in, and the per-stream DTLS/TURN loops. +package udprelay + +import ( + "context" + "errors" + "fmt" + "net" + "sync" + "sync/atomic" + "time" + + "github.com/samosvalishe/btp/internal/logx" + "github.com/samosvalishe/btp/internal/proxy/common" + "github.com/samosvalishe/btp/internal/transport/dtlsdial" +) + +// GetCredsFunc is re-exported from common so callers can keep their imports +// scoped to this package. +type GetCredsFunc = common.GetCredsFunc + +// AuthHandler is the subset of vkauth.Client this package needs. Defined as +// an interface so tests can inject fakes; the production wiring still imports +// vkauth for its sentinel errors (ErrFatalCaptchaNoStreams, etc.). +type AuthHandler interface { + IsAuthError(err error) bool + HandleAuthError(streamID int) bool + ResetErrors(streamID int) + LockoutUntilUnix() int64 +} + +// Params is the per-stream TURN/wrap configuration shared by the DTLS and TURN loops. +type Params struct { + Host string + Port string + Link string + UDP bool + WrapKey []byte + GetCreds GetCredsFunc +} + +// ErrFatal is returned by Run when a stream encounters a condition that +// requires the entire application to exit (e.g. manual captcha solver failed +// with no connected streams). Callers should check with errors.Is and call +// os.Exit themselves — udprelay does not reach into the host process. +var ErrFatal = errors.New("udprelay: fatal error") + +// Deps groups everything the loops need from the host process. The atomics +// are owned by Run and exposed here so DTLSLoop/TURNLoop can share them when +// called directly (Run wires them automatically). +type Deps struct { + DTLSDialer *dtlsdial.Dialer + Auth AuthHandler + Log logx.Logger + ActiveLocalPeer *atomic.Value + ConnectedStreams *atomic.Int32 + // fatalCh is an internal signalling channel; set by Run, written by + // TURNLoop, and drained by Run to propagate the fatal error up. + fatalCh chan error +} + +func (d *Deps) log() logx.Logger { + if d.Log == nil { + return logx.Nop() + } + return d.Log +} + +// Run is the UDP-mode entrypoint. It binds listenAddr, fans inbound packets +// into a shared queue, and spawns numStreams pairs of (DTLSLoop, TURNLoop). +// connectedStreams is owned by the caller (vkauth reads it via StreamsAlive) +// and incremented/decremented by oneTURN. +// Returns after all stream loops exit (i.e. when ctx is cancelled). +// If a fatal captcha condition is encountered, Run returns ErrFatal so the +// caller can perform os.Exit without udprelay reaching into the host process. +func Run(ctx context.Context, dtlsDialer *dtlsdial.Dialer, auth AuthHandler, logger logx.Logger, connectedStreams *atomic.Int32, params *Params, peer *net.UDPAddr, listenAddr string, numStreams int) error { + listenConn, err := net.ListenPacket("udp", listenAddr) + if err != nil { + return fmt.Errorf("udprelay listen %s: %w", listenAddr, err) + } + context.AfterFunc(ctx, func() { + if closeErr := listenConn.Close(); closeErr != nil { + logger.Errorf("udprelay: close local connection: %s", closeErr) + } + }) + + if numStreams <= 0 { + numStreams = 1 + } + + fatalCh := make(chan error, 1) + var activeLocalPeer atomic.Value + deps := &Deps{ + DTLSDialer: dtlsDialer, + Auth: auth, + Log: logger, + ActiveLocalPeer: &activeLocalPeer, + ConnectedStreams: connectedStreams, + fatalCh: fatalCh, + } + + // runCtx is cancelled when a fatal error is detected (via fatalCh), which + // propagates cancellation into all stream loops without requiring them to + // hold a reference to the host-process cancel function. + runCtx, runCancel := context.WithCancel(ctx) + defer runCancel() + + inboundChan := make(chan *Packet, inboundQueueCap) + wg := sync.WaitGroup{} + wg.Go(func() { + runListener(runCtx, listenConn, &activeLocalPeer, inboundChan) + }) + t := time.Tick(200 * time.Millisecond) + + // Stream 1 gets okchan so it can signal the first successful handshake to + // the log. All streams start concurrently — no gate between stream 1 and + // the rest, so a slow DTLS handshake on stream 1 never delays streams 2..N. + okchan := make(chan struct{}, 1) + for i := 0; i < numStreams; i++ { + cchan := make(chan net.PacketConn) + var ok chan<- struct{} + if i == 0 { + ok = okchan + } + streamID := i + 1 + wg.Go(func() { + DTLSLoop(runCtx, deps, peer, listenConn, inboundChan, cchan, ok, streamID) + }) + wg.Go(func() { + TURNLoop(runCtx, deps, params, peer, cchan, t, streamID) + }) + } + + // If a fatal error was sent, cancel remaining goroutines and propagate up. + var fatalErr error + go func() { + select { + case err := <-fatalCh: + fatalErr = err + runCancel() + case <-runCtx.Done(): + } + }() + + wg.Wait() + if fatalErr != nil { + return fatalErr + } + return nil +}