GODRIVER-1710 Ensure TLS ServerName is not overriden if already set (#486)

Author: Divjot Arora

x/mongo/driver/topology/connection.go

@@ -149,7 +149,7 @@ func (c *connection) connect(ctx context.Context) {
 			Cache:                   c.config.ocspCache,
 			DisableEndpointChecking: c.config.disableOCSPEndpointCheck,
 		}
-		tlsNc, err := configureTLS(ctx, c.nc, c.addr, tlsConfig, ocspOpts)
+		tlsNc, err := configureTLS(ctx, c.config.tlsConnectionSource, c.nc, c.addr, tlsConfig, ocspOpts)
 		if err != nil {
 			c.processInitializationError(err)
 			return
@@ -604,19 +604,27 @@ func (c *Connection) LocalAddress() address.Address {
 var notMasterCodes = []int32{10107, 13435}
 var recoveringCodes = []int32{11600, 11602, 13436, 189, 91}
 
-func configureTLS(ctx context.Context, nc net.Conn, addr address.Address, config *tls.Config, ocspOpts *ocsp.VerifyOptions) (net.Conn, error) {
-	// Ensure config.ServerName is always set for SNI.
-	hostname := addr.String()
-	colonPos := strings.LastIndex(hostname, ":")
-	if colonPos == -1 {
-		colonPos = len(hostname)
-	}
+func configureTLS(ctx context.Context,
+	tlsConnSource tlsConnectionSource,
+	nc net.Conn,
+	addr address.Address,
+	config *tls.Config,
+	ocspOpts *ocsp.VerifyOptions,
+) (net.Conn, error) {
 
-	hostname = hostname[:colonPos]
-	config.ServerName = hostname
+	// Ensure config.ServerName is always set for SNI.
+	if config.ServerName == "" {
+		hostname := addr.String()
+		colonPos := strings.LastIndex(hostname, ":")
+		if colonPos == -1 {
+			colonPos = len(hostname)
+		}
 
-	client := tls.Client(nc, config)
+		hostname = hostname[:colonPos]
+		config.ServerName = hostname
+	}
 
+	client := tlsConnSource.Client(nc, config)
 	errChan := make(chan error, 1)
 	go func() {
 		errChan <- client.Handshake()

x/mongo/driver/topology/connection_options.go

@@ -52,13 +52,15 @@ type connectionConfig struct {
 	ocspCache                ocsp.Cache
 	disableOCSPEndpointCheck bool
 	errorHandlingCallback    func(error, uint64)
+	tlsConnectionSource      tlsConnectionSource
 }
 
 func newConnectionConfig(opts ...ConnectionOption) (*connectionConfig, error) {
 	cfg := &connectionConfig{
-		connectTimeout: 30 * time.Second,
-		dialer:         nil,
-		lifeTimeout:    30 * time.Minute,
+		connectTimeout:      30 * time.Second,
+		dialer:              nil,
+		lifeTimeout:         30 * time.Minute,
+		tlsConnectionSource: defaultTLSConnectionSource,
 	}
 
 	for _, opt := range opts {
@@ -78,6 +80,13 @@ func newConnectionConfig(opts ...ConnectionOption) (*connectionConfig, error) {
 // ConnectionOption is used to configure a connection.
 type ConnectionOption func(*connectionConfig) error
 
+func withTLSConnectionSource(fn func(tlsConnectionSource) tlsConnectionSource) ConnectionOption {
+	return func(c *connectionConfig) error {
+		c.tlsConnectionSource = fn(c.tlsConnectionSource)
+		return nil
+	}
+}
+
 func withErrorHandlingCallback(fn func(error, uint64)) ConnectionOption {
 	return func(c *connectionConfig) error {
 		c.errorHandlingCallback = fn

x/mongo/driver/topology/connection_test.go

@@ -8,6 +8,7 @@ package topology
 
 import (
 	"context"
+	"crypto/tls"
 	"errors"
 	"net"
 	"sync"
@@ -175,6 +176,58 @@ func TestConnection(t *testing.T) {
 					wg.Wait()
 				})
 			})
+			t.Run("tls", func(t *testing.T) {
+				t.Run("connection source is set to default if unspecified", func(t *testing.T) {
+					conn, err := newConnection(address.Address(""))
+					assert.Nil(t, err, "newConnection error: %v", err)
+					assert.NotNil(t, conn.config.tlsConnectionSource, "expected tlsConnectionSource to be set but was not")
+				})
+				t.Run("server name", func(t *testing.T) {
+					testCases := []struct {
+						name               string
+						addr               address.Address
+						cfg                *tls.Config
+						expectedServerName string
+					}{
+						{"set to connection address if empty", "localhost:27017", &tls.Config{}, "localhost"},
+						{"left alone if non-empty", "localhost:27017", &tls.Config{ServerName: "other"}, "other"},
+					}
+					for _, tc := range testCases {
+						t.Run(tc.name, func(t *testing.T) {
+							var sentCfg *tls.Config
+							var testTLSConnectionSource tlsConnectionSourceFn = func(nc net.Conn, cfg *tls.Config) *tls.Conn {
+								sentCfg = cfg
+								return tls.Client(nc, cfg)
+							}
+
+							connOpts := []ConnectionOption{
+								WithDialer(func(Dialer) Dialer {
+									return DialerFunc(func(context.Context, string, string) (net.Conn, error) {
+										return &net.TCPConn{}, nil
+									})
+								}),
+								WithHandshaker(func(Handshaker) Handshaker {
+									return &testHandshaker{}
+								}),
+								WithTLSConfig(func(*tls.Config) *tls.Config {
+									return tc.cfg
+								}),
+								withTLSConnectionSource(func(tlsConnectionSource) tlsConnectionSource {
+									return testTLSConnectionSource
+								}),
+							}
+							conn, err := newConnection(tc.addr, connOpts...)
+							assert.Nil(t, err, "newConnection error: %v", err)
+
+							conn.connect(context.Background())
+							err = conn.wait()
+							assert.NotNil(t, sentCfg, "expected TLS config to be set, but was not")
+							assert.Equal(t, tc.expectedServerName, sentCfg.ServerName, "expected ServerName %s, got %s",
+								tc.expectedServerName, sentCfg.ServerName)
+						})
+					}
+				})
+			})
 		})
 		t.Run("writeWireMessage", func(t *testing.T) {
 			t.Run("closed connection", func(t *testing.T) {

x/mongo/driver/topology/tls_connection_source.go

@@ -0,0 +1,26 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package topology
+
+import (
+	"crypto/tls"
+	"net"
+)
+
+type tlsConnectionSource interface {
+	Client(net.Conn, *tls.Config) *tls.Conn
+}
+
+type tlsConnectionSourceFn func(net.Conn, *tls.Config) *tls.Conn
+
+func (t tlsConnectionSourceFn) Client(nc net.Conn, cfg *tls.Config) *tls.Conn {
+	return t(nc, cfg)
+}
+
+var defaultTLSConnectionSource tlsConnectionSourceFn = func(nc net.Conn, cfg *tls.Config) *tls.Conn {
+	return tls.Client(nc, cfg)
+}