github.com/gocql/gocql@v1.6.0/dial.go (about)

     1  package gocql
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"fmt"
     7  	"net"
     8  	"strings"
     9  )
    10  
    11  // HostDialer allows customizing connection to cluster nodes.
    12  type HostDialer interface {
    13  	// DialHost establishes a connection to the host.
    14  	// The returned connection must be directly usable for CQL protocol,
    15  	// specifically DialHost is responsible also for setting up the TLS session if needed.
    16  	// DialHost should disable write coalescing if the returned net.Conn does not support writev.
    17  	// As of Go 1.18, only plain TCP connections support writev, TLS sessions should disable coalescing.
    18  	// You can use WrapTLS helper function if you don't need to override the TLS setup.
    19  	DialHost(ctx context.Context, host *HostInfo) (*DialedHost, error)
    20  }
    21  
    22  // DialedHost contains information about established connection to a host.
    23  type DialedHost struct {
    24  	// Conn used to communicate with the server.
    25  	Conn net.Conn
    26  
    27  	// DisableCoalesce disables write coalescing for the Conn.
    28  	// If true, the effect is the same as if WriteCoalesceWaitTime was configured to 0.
    29  	DisableCoalesce bool
    30  }
    31  
    32  // defaultHostDialer dials host in a default way.
    33  type defaultHostDialer struct {
    34  	dialer    Dialer
    35  	tlsConfig *tls.Config
    36  }
    37  
    38  func (hd *defaultHostDialer) DialHost(ctx context.Context, host *HostInfo) (*DialedHost, error) {
    39  	ip := host.ConnectAddress()
    40  	port := host.Port()
    41  
    42  	if !validIpAddr(ip) {
    43  		return nil, fmt.Errorf("host missing connect ip address: %v", ip)
    44  	} else if port == 0 {
    45  		return nil, fmt.Errorf("host missing port: %v", port)
    46  	}
    47  
    48  	connAddr := host.ConnectAddressAndPort()
    49  	conn, err := hd.dialer.DialContext(ctx, "tcp", connAddr)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	addr := host.HostnameAndPort()
    54  	return WrapTLS(ctx, conn, addr, hd.tlsConfig)
    55  }
    56  
    57  func tlsConfigForAddr(tlsConfig *tls.Config, addr string) *tls.Config {
    58  	// the TLS config is safe to be reused by connections but it must not
    59  	// be modified after being used.
    60  	if !tlsConfig.InsecureSkipVerify && tlsConfig.ServerName == "" {
    61  		colonPos := strings.LastIndex(addr, ":")
    62  		if colonPos == -1 {
    63  			colonPos = len(addr)
    64  		}
    65  		hostname := addr[:colonPos]
    66  		// clone config to avoid modifying the shared one.
    67  		tlsConfig = tlsConfig.Clone()
    68  		tlsConfig.ServerName = hostname
    69  	}
    70  	return tlsConfig
    71  }
    72  
    73  // WrapTLS optionally wraps a net.Conn connected to addr with the given tlsConfig.
    74  // If the tlsConfig is nil, conn is not wrapped into a TLS session, so is insecure.
    75  // If the tlsConfig does not have server name set, it is updated based on the default gocql rules.
    76  func WrapTLS(ctx context.Context, conn net.Conn, addr string, tlsConfig *tls.Config) (*DialedHost, error) {
    77  	if tlsConfig != nil {
    78  		tlsConfig := tlsConfigForAddr(tlsConfig, addr)
    79  		tconn := tls.Client(conn, tlsConfig)
    80  		if err := tconn.HandshakeContext(ctx); err != nil {
    81  			conn.Close()
    82  			return nil, err
    83  		}
    84  		conn = tconn
    85  	}
    86  
    87  	return &DialedHost{
    88  		Conn:            conn,
    89  		DisableCoalesce: tlsConfig != nil, // write coalescing can't use writev when the connection is wrapped.
    90  	}, nil
    91  }