diff --git a/internal/netconn/pipe.go b/internal/netconn/pipe.go new file mode 100644 index 0000000..511df03 --- /dev/null +++ b/internal/netconn/pipe.go @@ -0,0 +1,54 @@ +package netconn + +import ( + "context" + "io" + "net" + "sync" + "time" +) + +// BiCopy runs io.Copy in both directions between c1 and c2, returning +// (bytesC1FromC2, bytesC2FromC1). When either copy exits — by EOF, error, or +// ctx cancellation — both sides get their deadlines pinned to time.Now() so the +// peer goroutine unblocks. Errors are surfaced via errf when non-nil. +func BiCopy(ctx context.Context, c1, c2 net.Conn, errf func(format string, v ...any)) (int64, int64) { + ctx2, cancel := context.WithCancel(ctx) + context.AfterFunc(ctx2, func() { + now := time.Now() + if err := c1.SetDeadline(now); err != nil && errf != nil { + errf("BiCopy: c1 SetDeadline: %v", err) + } + if err := c2.SetDeadline(now); err != nil && errf != nil { + errf("BiCopy: c2 SetDeadline: %v", err) + } + }) + + var wg sync.WaitGroup + var c1FromC2, c2FromC1 int64 + wg.Go(func() { + defer cancel() + n, err := io.Copy(c1, c2) + c1FromC2 = n + if err != nil && errf != nil { + errf("BiCopy: c1<-c2: %v", err) + } + }) + wg.Go(func() { + defer cancel() + n, err := io.Copy(c2, c1) + c2FromC1 = n + if err != nil && errf != nil { + errf("BiCopy: c2<-c1: %v", err) + } + }) + wg.Wait() + + if err := c1.SetDeadline(time.Time{}); err != nil && errf != nil { + errf("BiCopy: c1 clear deadline: %v", err) + } + if err := c2.SetDeadline(time.Time{}); err != nil && errf != nil { + errf("BiCopy: c2 clear deadline: %v", err) + } + return c1FromC2, c2FromC1 +} diff --git a/internal/proxy/tcpfwd/tcpfwd.go b/internal/proxy/tcpfwd/tcpfwd.go index e701ae7..e44e31e 100644 --- a/internal/proxy/tcpfwd/tcpfwd.go +++ b/internal/proxy/tcpfwd/tcpfwd.go @@ -3,13 +3,13 @@ package tcpfwd import ( "context" "fmt" - "io" "net" "sync" "time" "github.com/cacggghp/vk-turn-proxy/internal/client/ish" "github.com/cacggghp/vk-turn-proxy/internal/logx" + "github.com/cacggghp/vk-turn-proxy/internal/netconn" "github.com/cacggghp/vk-turn-proxy/internal/proxy/common" "github.com/cacggghp/vk-turn-proxy/internal/stats" "github.com/cacggghp/vk-turn-proxy/internal/transport/dtlsdial" @@ -162,7 +162,12 @@ func Run(ctx context.Context, deps *Deps, params *Params, peer *net.UDPAddr, lis return } defer func() { _ = stream.Close() }() - fromSession, toSession := pipe(deps, ctx, tc, stream) + errf := func(format string, v ...any) { + if deps.log().DebugEnabled() { + deps.log().Errorf(format, v...) + } + } + fromSession, toSession := netconn.BiCopy(ctx, tc, stream, errf) sessRef.FromSession.Add(uint64(fromSession)) sessRef.ToSession.Add(uint64(toSession)) deps.log().Debugf("[session %d] TCP done #%d local<-session=%s local->session=%s", @@ -275,44 +280,3 @@ func createSmuxSession(ctx context.Context, deps *Deps, params *Params, peer *ne return smuxSess, cleanup, nil } -// pipe copies data bidirectionally between two connections, cancelling both -// sides as soon as either copy finishes. Returns (c1<-c2, c2<-c1) bytes. -func pipe(deps *Deps, ctx context.Context, c1, c2 net.Conn) (int64, int64) { - ctx2, cancel := context.WithCancel(ctx) - context.AfterFunc(ctx2, func() { - if err := c1.SetDeadline(time.Now()); err != nil { - deps.log().Errorf("pipe: failed to set deadline c1: %v", err) - } - if err := c2.SetDeadline(time.Now()); err != nil { - deps.log().Errorf("pipe: failed to set deadline c2: %v", err) - } - }) - - var wg sync.WaitGroup - var c1FromC2 int64 - var c2FromC1 int64 - wg.Go(func() { - defer cancel() - n, err := io.Copy(c1, c2) - c1FromC2 = n - if err != nil && deps.log().DebugEnabled() { - deps.log().Errorf("pipe: c1<-c2 copy error: %v", err) - } - }) - wg.Go(func() { - defer cancel() - n, err := io.Copy(c2, c1) - c2FromC1 = n - if err != nil && deps.log().DebugEnabled() { - deps.log().Errorf("pipe: c2<-c1 copy error: %v", err) - } - }) - wg.Wait() - if err := c1.SetDeadline(time.Time{}); err != nil && deps.log().DebugEnabled() { - deps.log().Errorf("pipe: failed to reset deadline c1: %v", err) - } - if err := c2.SetDeadline(time.Time{}); err != nil && deps.log().DebugEnabled() { - deps.log().Errorf("pipe: failed to reset deadline c2: %v", err) - } - return c1FromC2, c2FromC1 -} diff --git a/internal/proxy/tcpfwdserver/tcpfwdserver.go b/internal/proxy/tcpfwdserver/tcpfwdserver.go index 0eed048..8de770e 100644 --- a/internal/proxy/tcpfwdserver/tcpfwdserver.go +++ b/internal/proxy/tcpfwdserver/tcpfwdserver.go @@ -12,6 +12,7 @@ import ( "time" "github.com/cacggghp/vk-turn-proxy/internal/logx" + "github.com/cacggghp/vk-turn-proxy/internal/netconn" "github.com/cacggghp/vk-turn-proxy/internal/proxy/bondserver" "github.com/cacggghp/vk-turn-proxy/internal/stats" "github.com/cacggghp/vk-turn-proxy/internal/transport/kcptun" @@ -110,7 +111,7 @@ func handleStream(ctx context.Context, logger logx.Logger, registry *bondserver. } }() - pipeConn(ctx, logger, &prefixedConn{Conn: s, prefix: prefix[:]}, backendConn) + netconn.BiCopy(ctx, &prefixedConn{Conn: s, prefix: prefix[:]}, backendConn, logger.Debugf) } // prefixedConn re-injects the magic-peek prefix on the first reads so the @@ -129,37 +130,3 @@ func (c *prefixedConn) Read(p []byte) (int, error) { return c.Conn.Read(p) } -// pipeConn copies data bidirectionally between two connections. -func pipeConn(ctx context.Context, logger logx.Logger, c1, c2 net.Conn) { - ctx2, cancel := context.WithCancel(ctx) - defer cancel() - - context.AfterFunc(ctx2, func() { - if err := c1.SetDeadline(time.Now()); err != nil { - logger.Debugf("pipeConn: failed to set deadline c1: %v", err) - } - if err := c2.SetDeadline(time.Now()); err != nil { - logger.Debugf("pipeConn: failed to set deadline c2: %v", err) - } - }) - - var wg sync.WaitGroup - wg.Go(func() { - if _, err := io.Copy(c1, c2); err != nil { - logger.Debugf("pipeConn: c1<-c2 copy error: %v", err) - } - }) - wg.Go(func() { - if _, err := io.Copy(c2, c1); err != nil { - logger.Debugf("pipeConn: c2<-c1 copy error: %v", err) - } - }) - wg.Wait() - - if err := c1.SetDeadline(time.Time{}); err != nil { - logger.Debugf("pipeConn: failed to reset deadline c1: %v", err) - } - if err := c2.SetDeadline(time.Time{}); err != nil { - logger.Debugf("pipeConn: failed to reset deadline c2: %v", err) - } -}