feat: expose DialContext in Dialer and propagate context (#165)

kruskall · web-flow · commit 473983911d7c · 2024-01-22T22:35:42.000+01:00
* feat: expose DialContext in Dialer and propagate context

A dialer only exposed the Dial method, making it impossible to correctly
propagate context along with timeouts and cancellations

* fix: update windows specific file to comply with interface

* build: fix windows compile errors

* docs: add comment about contextDialer

* feat: add ConnectContext to transport client

* lint: fix spelling mistake
diff --git a/transport/client.go b/transport/client.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -88,6 +89,10 @@ func NewClientWithDialer(d Dialer, c Config, network, host string, defaultPort i
 }
 
 func (c *Client) Connect() error {
+	return c.ConnectContext(context.Background())
+}
+
+func (c *Client) ConnectContext(ctx context.Context) error {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
@@ -96,7 +101,7 @@ func (c *Client) Connect() error {
 		c.conn = nil
 	}
 
-	conn, err := c.dialer.Dial(c.network, c.host)
+	conn, err := c.dialer.DialContext(ctx, c.network, c.host)
 	if err != nil {
 		return err
 	}
@@ -217,7 +222,7 @@ func (c *Client) Test(d testing.Driver) {
 	d.Run("logstash: "+c.host, func(d testing.Driver) {
 		d.Run("connection", func(d testing.Driver) {
 			netDialer := TestNetDialer(d, c.config.Timeout)
-			_, err := netDialer.Dial("tcp", c.host)
+			_, err := netDialer.DialContext(context.Background(), "tcp", c.host)
 			d.Fatal("dial up", err)
 		})
 
@@ -227,7 +232,7 @@ func (c *Client) Test(d testing.Driver) {
 			d.Run("TLS", func(d testing.Driver) {
 				netDialer := NetDialer(c.config.Timeout)
 				tlsDialer := TestTLSDialer(d, netDialer, c.config.TLS, c.config.Timeout)
-				_, err := tlsDialer.Dial("tcp", c.host)
+				_, err := tlsDialer.DialContext(context.Background(), "tcp", c.host)
 				d.Fatal("dial up", err)
 			})
 		}
diff --git a/transport/dialer/dialer_windows.go b/transport/dialer/dialer_windows.go
@@ -20,6 +20,7 @@
 package dialer
 
 import (
+	"context"
 	"errors"
 	"net"
 	"strings"
@@ -60,10 +61,12 @@ func (t *NpipeDialerBuilder) String() string {
 func (t *NpipeDialerBuilder) Make(timeout time.Duration) (transport.Dialer, error) {
 	to := timeout
 	return transport.DialerFunc(
-		func(_, _ string) (net.Conn, error) {
-			return winio.DialPipe(
+		func(ctx context.Context, _, _ string) (net.Conn, error) {
+			ctx, cancel := context.WithTimeout(ctx, to)
+			defer cancel()
+			return winio.DialPipeContext(
+				ctx,
 				strings.TrimSuffix(npipe.TransformString(t.Path), "/"),
-				&to,
 			)
 		},
 	), nil
diff --git a/transport/httpcommon/httpcommon.go b/transport/httpcommon/httpcommon.go
@@ -257,10 +257,8 @@ func (settings *HTTPTransportSettings) httpRoundTripper(
 	opts ...TransportOption,
 ) *http.Transport {
 	t := http.DefaultTransport.(*http.Transport).Clone()
-	t.DialContext = nil
-	t.DialTLSContext = nil
-	t.Dial = dialer.Dial       //nolint:staticcheck // use deprecated function to preserve functionality
-	t.DialTLS = tlsDialer.Dial //nolint:staticcheck // use deprecated function to preserve functionality
+	t.DialContext = dialer.DialContext
+	t.DialTLSContext = tlsDialer.DialContext
 	t.TLSClientConfig = tls.ToConfig()
 	t.ForceAttemptHTTP2 = false
 	t.Proxy = settings.Proxy.ProxyFunc()
diff --git a/transport/logging.go b/transport/logging.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"errors"
 	"io"
 	"net"
@@ -31,9 +32,9 @@ type loggingConn struct {
 }
 
 func LoggingDialer(d Dialer, logger *logp.Logger) Dialer {
-	return DialerFunc(func(network, addr string) (net.Conn, error) {
+	return DialerFunc(func(ctx context.Context, network, addr string) (net.Conn, error) {
 		logger := logger.With("network", network, "address", addr)
-		c, err := d.Dial(network, addr)
+		c, err := d.DialContext(ctx, network, addr)
 		if err != nil {
 			logger.Errorf("Error dialing %v", err)
 			return nil, err
diff --git a/transport/proxy.go b/transport/proxy.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"net"
 	"net/url"
 
@@ -68,7 +69,7 @@ func ProxyDialer(log *logp.Logger, config *ProxyConfig, forward Dialer) (Dialer,
 	}
 
 	log.Infof("proxy host: '%s'", url.Host)
-	return DialerFunc(func(network, address string) (net.Conn, error) {
+	return DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
 		var err error
 		var addresses []string
 
@@ -94,6 +95,17 @@ func ProxyDialer(log *logp.Logger, config *ProxyConfig, forward Dialer) (Dialer,
 		if err != nil {
 			return nil, err
 		}
-		return DialWith(dialer, network, host, addresses, port)
+
+		contextDialer, ok := dialer.(Dialer)
+		// This will never be executed because the proxy package always returns
+		// a ContextDialer but they didn't break the interface for backward compatibility.
+		// See golang/go#58376
+		if !ok {
+			contextDialer = DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
+				return dialer.Dial(network, address)
+			})
+		}
+
+		return DialWith(ctx, contextDialer, network, host, addresses, port)
 	}), nil
 }
diff --git a/transport/tcp.go b/transport/tcp.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"strings"
@@ -32,7 +33,7 @@ func NetDialer(timeout time.Duration) Dialer {
 }
 
 func TestNetDialer(d testing.Driver, timeout time.Duration) Dialer {
-	return DialerFunc(func(network, address string) (net.Conn, error) {
+	return DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
 		switch network {
 		case "tcp", "tcp4", "tcp6", "udp", "udp4", "udp6":
 		default:
@@ -55,7 +56,7 @@ func TestNetDialer(d testing.Driver, timeout time.Duration) Dialer {
 
 		// dial via host IP by randomized iteration of known IPs
 		dialer := &net.Dialer{Timeout: timeout}
-		return DialWith(dialer, network, host, addresses, port)
+		return DialWith(ctx, dialer, network, host, addresses, port)
 	})
 }
 
@@ -66,7 +67,7 @@ func UnixDialer(timeout time.Duration, sockFile string) Dialer {
 
 // TestUnixDialer creates a Test Unix Dialer when using domain socket.
 func TestUnixDialer(d testing.Driver, timeout time.Duration, sockFile string) Dialer {
-	return DialerFunc(func(network, address string) (net.Conn, error) {
+	return DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
 		d.Info("connecting using unix domain socket", sockFile)
 		return net.DialTimeout("unix", sockFile, timeout)
 	})
diff --git a/transport/tls.go b/transport/tls.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -44,7 +45,7 @@ func TestTLSDialer(
 	var lastAddress string
 	var m sync.Mutex
 
-	return DialerFunc(func(network, address string) (net.Conn, error) {
+	return DialerFunc(func(ctx context.Context, network, address string) (net.Conn, error) {
 		switch network {
 		case "tcp", "tcp4", "tcp6":
 		default:
@@ -69,18 +70,18 @@ func TestTLSDialer(
 		}
 		m.Unlock()
 
-		return tlsDialWith(d, forward, network, address, timeout, tlsConfig, config)
+		return tlsDialWith(ctx, d, forward, network, address, timeout, tlsConfig, config)
 	})
 }
 
 type DialerH2 interface {
-	Dial(network, address string, cfg *tls.Config) (net.Conn, error)
+	DialContext(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error)
 }
 
-type DialerFuncH2 func(network, address string, cfg *tls.Config) (net.Conn, error)
+type DialerFuncH2 func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error)
 
-func (d DialerFuncH2) Dial(network, address string, cfg *tls.Config) (net.Conn, error) {
-	return d(network, address, cfg)
+func (d DialerFuncH2) DialContext(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) {
+	return d(ctx, network, address, cfg)
 }
 
 func TLSDialerH2(forward Dialer, config *tlscommon.TLSConfig, timeout time.Duration) (DialerH2, error) {
@@ -98,7 +99,7 @@ func TestTLSDialerH2(
 	var lastAddress string
 	var m sync.Mutex
 
-	return DialerFuncH2(func(network, address string, cfg *tls.Config) (net.Conn, error) {
+	return DialerFuncH2(func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) {
 		switch network {
 		case "tcp", "tcp4", "tcp6":
 		default:
@@ -126,19 +127,20 @@ func TestTLSDialerH2(
 		// NextProtos must be set from the passed h2 connection or it will fail
 		tlsConfig.NextProtos = cfg.NextProtos
 
-		return tlsDialWith(d, forward, network, address, timeout, tlsConfig, config)
+		return tlsDialWith(ctx, d, forward, network, address, timeout, tlsConfig, config)
 	}), nil
 }
 
 func tlsDialWith(
+	ctx context.Context,
 	d testing.Driver,
 	dialer Dialer,
 	network, address string,
 	timeout time.Duration,
 	tlsConfig *tls.Config,
 	config *tlscommon.TLSConfig,
 ) (net.Conn, error) {
-	socket, err := dialer.Dial(network, address)
+	socket, err := dialer.DialContext(ctx, network, address)
 	if err != nil {
 		return nil, err
 	}
diff --git a/transport/transport.go b/transport/transport.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"errors"
 	"net"
 
@@ -26,24 +27,29 @@ import (
 
 type Dialer interface {
 	Dial(network, address string) (net.Conn, error)
+	DialContext(ctx context.Context, network, address string) (net.Conn, error)
 }
 
-type DialerFunc func(network, address string) (net.Conn, error)
+type DialerFunc func(ctx context.Context, network, address string) (net.Conn, error)
 
 var (
 	ErrNotConnected = errors.New("client is not connected")
 )
 
 func (d DialerFunc) Dial(network, address string) (net.Conn, error) {
-	return d(network, address)
+	return d(context.Background(), network, address)
 }
 
-func Dial(c Config, network, address string) (net.Conn, error) {
+func (d DialerFunc) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	return d(ctx, network, address)
+}
+
+func DialContext(ctx context.Context, c Config, network, address string) (net.Conn, error) {
 	d, err := MakeDialer(c)
 	if err != nil {
 		return nil, err
 	}
-	return d.Dial(network, address)
+	return d.DialContext(ctx, network, address)
 }
 
 func MakeDialer(c Config) (Dialer, error) {
diff --git a/transport/util.go b/transport/util.go
@@ -18,6 +18,7 @@
 package transport
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"net"
@@ -43,6 +44,7 @@ func fullAddress(host string, defaultPort int) string {
 //
 // Use this to select and dial one IP being known for one host name.
 func DialWith(
+	ctx context.Context,
 	dialer Dialer,
 	network, host string,
 	addresses []string,
@@ -52,7 +54,7 @@ func DialWith(
 	case 0:
 		return nil, fmt.Errorf("no route to host %v", host)
 	case 1:
-		return dialer.Dial(network, net.JoinHostPort(addresses[0], port))
+		return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0], port))
 	}
 
 	// Use randomization on DNS reported addresses combined with timeout and ACKs
@@ -69,7 +71,7 @@ func DialWith(
 	// > "Clients, of course, may reorder this information" - with respect to
 	// > handling order of dns records in a response.forwarded. Really required?
 	for _, i := range rand.Perm(len(addresses)) {
-		c, err = dialer.Dial(network, net.JoinHostPort(addresses[i], port))
+		c, err = dialer.DialContext(ctx, network, net.JoinHostPort(addresses[i], port))
 		if err == nil && c != nil {
 			return c, err
 		}
diff --git a/transport/wrap.go b/transport/wrap.go
@@ -18,12 +18,13 @@
 package transport
 
 import (
+	"context"
 	"net"
 )
 
 func ConnWrapper(d Dialer, w func(net.Conn) net.Conn) Dialer {
-	return DialerFunc(func(network, addr string) (net.Conn, error) {
-		c, err := d.Dial(network, addr)
+	return DialerFunc(func(ctx context.Context, network, addr string) (net.Conn, error) {
+		c, err := d.DialContext(ctx, network, addr)
 		if err != nil {
 			return nil, err
 		}

Original file line number	Diff line number	Diff line change
`@@ -18,12 +18,13 @@`
`18`	`18`	`package transport`
`19`	`19`
`20`	`20`	`import (`
	`21`	`+ "context"`
`21`	`22`	`"net"`
`22`	`23`	`)`
`23`	`24`
`24`	`25`	`func ConnWrapper(d Dialer, w func(net.Conn) net.Conn) Dialer {`
`25`		`- return DialerFunc(func(network, addr string) (net.Conn, error) {`
`26`		`- c, err := d.Dial(network, addr)`
	`26`	`+ return DialerFunc(func(ctx context.Context, network, addr string) (net.Conn, error) {`
	`27`	`+ c, err := d.DialContext(ctx, network, addr)`
`27`	`28`	`if err != nil {`
`28`	`29`	`return nil, err`
`29`	`30`	`}`