[Connection] Add an optional Dialer to ClientOptions

pulsar/client.go

@@ -18,7 +18,9 @@
 package pulsar
 
 import (
+	"context"
 	"crypto/tls"
+	"net"
 	"time"
 
 	"github.com/apache/pulsar-client-go/pulsar/auth"
@@ -103,6 +105,11 @@ type ClientOptions struct {
 	// If your application is sensitive to service disruption, set this explicitly (e.g., 10s or 15s).
 	ConnectionTimeout time.Duration
 
+	// Dialer is an optional custom dialer used to establish TCP connections to
+	// brokers. When nil, a net.Dialer honouring ConnectionTimeout is used.
+	// The address passed to the dialer is the broker's host:port
+	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
+
 	// Set the operation timeout (default: 30 seconds)
 	// Producer-create, subscribe and unsubscribe operations will be retried until this interval, after which the
 	// operation will be marked as failed

pulsar/client_impl.go

@@ -166,7 +166,7 @@ func newClient(options ClientOptions) (Client, error) {
 
 	c := &client{
 		cnxPool: internal.NewConnectionPool(tlsConfig, authProvider, connectionTimeout, keepAliveInterval,
-			maxConnectionsPerHost, logger, metrics, options.Description, connectionMaxIdleTime),
+			maxConnectionsPerHost, logger, metrics, options.Description, connectionMaxIdleTime, options.Dialer),
 		log:              logger,
 		metrics:          metrics,
 		memLimit:         internal.NewMemoryLimitController(memLimitBytes, defaultMemoryLimitTriggerThreshold),

pulsar/internal/connection.go

@@ -138,6 +138,7 @@ type dataRequest struct {
 type connection struct {
 	started           int32
 	connectionTimeout time.Duration
+	dialer            func(ctx context.Context, network, addr string) (net.Conn, error)
 	closeOnce         sync.Once
 
 	// mu protects the fields below against concurrency accesses.
@@ -182,6 +183,7 @@ type connection struct {
 
 // connectionOptions defines configurations for creating connection.
 type connectionOptions struct {
+	dialer            func(ctx context.Context, network, addr string) (net.Conn, error)
 	logicalAddr       *url.URL
 	physicalAddr      *url.URL
 	tls               *TLSOptions
@@ -195,6 +197,7 @@ type connectionOptions struct {
 
 func newConnection(opts connectionOptions) *connection {
 	cnx := &connection{
+		dialer:               opts.dialer,
 		connectionTimeout:    opts.connectionTimeout,
 		keepAliveInterval:    opts.keepAliveInterval,
 		logicalAddr:          opts.logicalAddr,
@@ -258,13 +261,31 @@ func (c *connection) connect() bool {
 		tlsConfig *tls.Config
 	)
 
+	// time.Duration is initialized to 0 by default, net.Dialer's default timeout is no timeout
+	// therefore if c.connectionTimeout is 0, it means no timeout.
+	// As in tls.DialWithDialer, this deadline spans the dial and the TLS
+	// handshake together, not just the dial.
+	ctx := context.Background()
+	if c.connectionTimeout.Nanoseconds() > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, c.connectionTimeout)
+		defer cancel()
+	}
+
+	// kept separate from the TLS handshake below so that a custom dialer returns
+	// a plain net.Conn and the handshake is performed here, not inside
+	// tls.DialWithDialer.
+	dial := func() (net.Conn, error) {
+		if c.dialer != nil {
+			return c.dialer(ctx, "tcp", c.physicalAddr.Host)
+		}
+		var d net.Dialer
+		return d.DialContext(ctx, "tcp", c.physicalAddr.Host)
+	}
+
 	if c.tlsOptions == nil {
 		// Clear text connection
-		if c.connectionTimeout.Nanoseconds() > 0 {
-			cnx, err = net.DialTimeout("tcp", c.physicalAddr.Host, c.connectionTimeout)
-		} else {
-			cnx, err = net.Dial("tcp", c.physicalAddr.Host)
-		}
+		cnx, err = dial()
 	} else {
 		// TLS connection
 		tlsConfig, err = c.getTLSConfig()
@@ -273,10 +294,28 @@ func (c *connection) connect() bool {
 			return false
 		}
 
-		// time.Duration is initialized to 0 by default, net.Dialer's default timeout is no timeout
-		// therefore if c.connectionTimeout is 0, it means no timeout
-		d := &net.Dialer{Timeout: c.connectionTimeout}
-		cnx, err = tls.DialWithDialer(d, "tcp", c.physicalAddr.Host, tlsConfig)
+		// tls.DialWithDialer infers ServerName from the address being dialed when
+		// the config does not set one; tls.Client does not, and would verify
+		// against an empty name. Preserve that so behaviour is unchanged
+		// if ServerName is left empty.
+		if tlsConfig.ServerName == "" {
+			host := c.physicalAddr.Hostname()
+			if host != "" {
+				tlsConfig = tlsConfig.Clone()
+				tlsConfig.ServerName = host
+			}
+		}
+
+		var rawConn net.Conn
+		rawConn, err = dial()
+		if err == nil {
+			tlsConn := tls.Client(rawConn, tlsConfig)
+			if err = tlsConn.HandshakeContext(ctx); err != nil {
+				_ = rawConn.Close()
+			} else {
+				cnx = tlsConn
+			}
+		}
 	}
 
 	if err != nil {

pulsar/internal/connection_pool.go

@@ -18,7 +18,9 @@
 package internal
 
 import (
+	"context"
 	"fmt"
+	"net"
 	"net/url"
 	"sync"
 	"sync/atomic"
@@ -58,6 +60,7 @@ type connectionPool struct {
 	metrics     *Metrics
 	log         log.Logger
 	description string
+	dialer      func(ctx context.Context, network, addr string) (net.Conn, error)
 }
 
 // NewConnectionPool init connection pool.
@@ -70,7 +73,8 @@ func NewConnectionPool(
 	logger log.Logger,
 	metrics *Metrics,
 	description string,
-	connectionMaxIdleTime time.Duration) ConnectionPool {
+	connectionMaxIdleTime time.Duration,
+	dialer func(ctx context.Context, network, addr string) (net.Conn, error)) ConnectionPool {
 	p := &connectionPool{
 		connections:           make(map[string]*connection),
 		tlsOptions:            tlsOptions,
@@ -82,6 +86,7 @@ func NewConnectionPool(
 		metrics:               metrics,
 		closeCh:               make(chan struct{}),
 		description:           description,
+		dialer:                dialer,
 	}
 	go p.checkAndCleanIdleConnections(connectionMaxIdleTime)
 	return p
@@ -114,6 +119,7 @@ func (p *connectionPool) GetConnection(logicalAddr *url.URL, physicalAddr *url.U
 
 	if conn == nil {
 		conn = newConnection(connectionOptions{
+			dialer:            p.dialer,
 			logicalAddr:       logicalAddr,
 			physicalAddr:      physicalAddr,
 			tls:               p.tlsOptions,