diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go new file mode 100644 index 00000000..9221ba41 --- /dev/null +++ b/internal/proxy/proxy.go @@ -0,0 +1,131 @@ +package proxy + +import ( + "context" + "fmt" + "io" + "net" + "os" + "sync" + "time" +) + +// Proxy proxies local connections to a remote TCP address optionally using a custom dialer. +type Proxy struct { + Listener net.Listener + RemoteAddr string + DialContext func(ctx context.Context, network, address string) (net.Conn, error) + OnError func(error) + activeConns sync.WaitGroup +} + +// deadliner is an interface for listeners that support setting deadlines. +type deadliner interface { + SetDeadline(t time.Time) error +} + +// halfCloser is an interface for connections that support half-close. +type halfCloser interface { + CloseWrite() error +} + +// Run starts the proxy and runs until the context is canceled. +func (p *Proxy) Run(ctx context.Context) { + if p.DialContext == nil { + p.DialContext = (&net.Dialer{}).DialContext + } + + defer p.Listener.Close() + + // Handle incoming connections until context is canceled. +Loop: + for { + select { + case <-ctx.Done(): + break Loop + default: + } + + // Set a deadline on the listener if supported to check context periodically. + if dl, ok := p.Listener.(deadliner); ok { + dl.SetDeadline(time.Now().Add(1 * time.Second)) + } + + conn, err := p.Listener.Accept() + if err != nil { + if os.IsTimeout(err) { + // Just a timeout, continue to check context and accept again. + continue + } + + select { + case <-ctx.Done(): + break Loop + default: + if p.OnError != nil { + p.OnError(fmt.Errorf("accept local connection: %w", err)) + } + continue + } + } + + p.activeConns.Add(1) + go p.handleConnection(ctx, conn) + } + + // Wait for all connections to finish. + p.activeConns.Wait() +} + +func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) { + defer p.activeConns.Done() + defer localConn.Close() + + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + remoteConn, err := p.DialContext(ctx, "tcp", p.RemoteAddr) + if err != nil { + if p.OnError != nil { + p.OnError(fmt.Errorf("connect remote address '%s': %w", p.RemoteAddr, err)) + } + return + } + defer remoteConn.Close() + + // Bidirectional copy with proper half-close handling. + done := make(chan error, 2) + + go func() { + _, err := io.Copy(remoteConn, localConn) + // Close write half of remote connection if supported. + if hc, ok := remoteConn.(halfCloser); ok { + hc.CloseWrite() + } + done <- err + }() + + go func() { + _, err := io.Copy(localConn, remoteConn) + // Close write half of local connection if supported. + if hc, ok := localConn.(halfCloser); ok { + hc.CloseWrite() + } + done <- err + }() + + // Wait for both copies to complete or context cancel. + for i := 0; i < 2; i++ { + select { + case <-ctx.Done(): + // Close connections to abort ongoing copies. + localConn.Close() + remoteConn.Close() + return + case err = <-done: + if err != nil && p.OnError != nil { + p.OnError(fmt.Errorf("data copy: %w", err)) + } + } + } +}