Improve Config handling

connection.go

@@ -30,7 +30,6 @@ type mysqlConn struct {
 	result           mysqlResult // managed by clearResult() and handleOkPacket().
 	compIO           *compIO
 	cfg              *Config
-	connector        *connector
 	maxAllowedPacket int
 	maxWriteSize     int
 	capabilities     capabilityFlag

connector.go

@@ -19,8 +19,7 @@ import (
 )
 
 type connector struct {
-	cfg               *Config // immutable private copy.
-	encodedAttributes string  // Encoded connection attributes.
+	cfg *Config // immutable private copy.
 }
 
 func encodeConnectionAttributes(cfg *Config) string {
@@ -55,10 +54,8 @@ func encodeConnectionAttributes(cfg *Config) string {
 }
 
 func newConnector(cfg *Config) *connector {
-	encodedAttributes := encodeConnectionAttributes(cfg)
 	return &connector{
-		cfg:               cfg,
-		encodedAttributes: encodedAttributes,
+		cfg: cfg,
 	}
 }
 
@@ -75,6 +72,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		if err != nil {
 			return nil, err
 		}
+		cfg.encodedAttributes = encodeConnectionAttributes(cfg)
 	}
 
 	// New mysqlConn
@@ -83,20 +81,19 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		maxWriteSize:     maxPacketSize - 1,
 		closech:          make(chan struct{}),
 		cfg:              cfg,
-		connector:        c,
 	}
 	mc.parseTime = mc.cfg.ParseTime
 
 	// Connect to Server
 	dctx := ctx
 	if mc.cfg.Timeout > 0 {
 		var cancel context.CancelFunc
-		dctx, cancel = context.WithTimeout(ctx, c.cfg.Timeout)
+		dctx, cancel = context.WithTimeout(ctx, mc.cfg.Timeout)
 		defer cancel()
 	}
 
-	if c.cfg.DialFunc != nil {
-		mc.netConn, err = c.cfg.DialFunc(dctx, mc.cfg.Net, mc.cfg.Addr)
+	if mc.cfg.DialFunc != nil {
+		mc.netConn, err = mc.cfg.DialFunc(dctx, mc.cfg.Net, mc.cfg.Addr)
 	} else {
 		dialsLock.RLock()
 		dial, ok := dials[mc.cfg.Net]
@@ -116,7 +113,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	// Enable TCP Keepalives on TCP connections
 	if tc, ok := mc.netConn.(*net.TCPConn); ok {
 		if err := tc.SetKeepAlive(true); err != nil {
-			c.cfg.Logger.Print(err)
+			mc.cfg.Logger.Print(err)
 		}
 	}
 
@@ -145,15 +142,15 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	authResp, err := mc.auth(authData, plugin)
 	if err != nil {
 		// try the default auth plugin, if using the requested plugin failed
-		c.cfg.Logger.Print("could not use requested auth plugin '"+plugin+"': ", err.Error())
+		mc.cfg.Logger.Print("could not use requested auth plugin '"+plugin+"': ", err.Error())
 		plugin = defaultAuthPlugin
 		authResp, err = mc.auth(authData, plugin)
 		if err != nil {
 			mc.cleanup()
 			return nil, err
 		}
 	}
-	mc.initCapabilities(serverCapabilities, serverExtCapabilities, mc.cfg)
+	mc.initCapabilities(serverCapabilities, serverExtCapabilities)
 	if err = mc.writeHandshakeResponsePacket(authResp, plugin); err != nil {
 		mc.cleanup()
 		return nil, err

connector_test.go

@@ -1,7 +1,9 @@
 package mysql
 
 import (
+	"bytes"
 	"context"
+	"errors"
 	"net"
 	"testing"
 	"time"
@@ -28,3 +30,117 @@ func TestConnectorReturnsTimeout(t *testing.T) {
 		t.Fatalf("expected %T, got %T", nerr, err)
 	}
 }
+
+func TestBeforeConnectUsesEffectiveTimeout(t *testing.T) {
+	dialErr := errors.New("stop after observing dial context")
+	var remaining time.Duration
+
+	cfg := NewConfig()
+	cfg.Timeout = 2 * time.Hour
+	cfg.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
+		deadline, ok := ctx.Deadline()
+		if !ok {
+			t.Fatal("dial context has no deadline")
+		}
+		remaining = time.Until(deadline)
+		return nil, dialErr
+	}
+	if err := cfg.Apply(BeforeConnect(func(_ context.Context, cfg *Config) error {
+		cfg.Timeout = time.Hour
+		return nil
+	})); err != nil {
+		t.Fatal(err)
+	}
+
+	connector, err := NewConnector(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := connector.Connect(context.Background()); !errors.Is(err, dialErr) {
+		t.Fatalf("Connect() error = %v, want %v", err, dialErr)
+	}
+	if remaining < 59*time.Minute || remaining > 61*time.Minute {
+		t.Fatalf("dial timeout = %v, want about 1h", remaining)
+	}
+}
+
+func TestBeforeConnectUsesEffectiveDialerAndAttributes(t *testing.T) {
+	serverHandshake := []byte(
+		"\x48\x00\x00\x00" + // Packet header: 72-byte payload, sequence 0.
+			"\x0a" + // Protocol version 10.
+			"5.5.8\x00" + // NUL-terminated server version.
+			"\xa5\x00\x00\x00" + // Connection ID 165.
+			"<F?:Dh\"a" + // First 8 bytes of the authentication scramble.
+			"\x00" + // Filler.
+			"\xdf\xf7" + // Lower 2 bytes of the server capability flags.
+			"\x21" + // utf8_general_ci character set.
+			"\x02\x00" + // SERVER_STATUS_AUTOCOMMIT.
+			"\x1f\x80" + // Upper 2 bytes of the server capability flags.
+			"\x15" + // Authentication plugin data length: 21 bytes.
+			"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + // Reserved.
+			"bxr/UKmc3M2@\x00" + // Remaining authentication scramble.
+			"mysql_native_password", // Authentication plugin name.
+	)
+	okPacket := []byte(
+		"\x07\x00\x00\x02" + // Packet header: 7-byte payload, sequence 2.
+			"\x00" + // OK packet header.
+			"\x00\x00" + // Zero affected rows and last insert ID.
+			"\x02\x00" + // SERVER_STATUS_AUTOCOMMIT.
+			"\x00\x00", // Zero warnings.
+	)
+	mock := &mockConn{
+		data:          serverHandshake,
+		queuedReplies: [][]byte{okPacket},
+	}
+
+	var (
+		initialDialCalled bool
+		dialNetwork       string
+		dialAddress       string
+	)
+	cfg := NewConfig()
+	cfg.Addr = "initial.example:3306"
+	cfg.ConnectionAttributes = "phase:initial"
+	cfg.DialFunc = func(context.Context, string, string) (net.Conn, error) {
+		initialDialCalled = true
+		return nil, errors.New("initial dialer must not be called")
+	}
+	if err := cfg.Apply(BeforeConnect(func(_ context.Context, cfg *Config) error {
+		cfg.Addr = "callback.example:3306"
+		cfg.ConnectionAttributes = "phase:callback"
+		cfg.DialFunc = func(_ context.Context, network, address string) (net.Conn, error) {
+			dialNetwork = network
+			dialAddress = address
+			return mock, nil
+		}
+		return nil
+	})); err != nil {
+		t.Fatal(err)
+	}
+
+	connector, err := NewConnector(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	conn, err := connector.Connect(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	if initialDialCalled {
+		t.Fatal("Connect() used the pre-callback DialFunc")
+	}
+	if dialNetwork != "tcp" || dialAddress != "callback.example:3306" {
+		t.Fatalf("dialed %s(%s), want tcp(callback.example:3306)", dialNetwork, dialAddress)
+	}
+	if !bytes.Contains(mock.written, []byte("phase\bcallback")) {
+		t.Fatalf("handshake response does not contain callback attributes: %q", mock.written)
+	}
+	if bytes.Contains(mock.written, []byte("phase\x07initial")) {
+		t.Fatalf("handshake response contains stale attributes: %q", mock.written)
+	}
+	if !bytes.Contains(mock.written, []byte("callback.example")) {
+		t.Fatalf("handshake response does not contain callback server host: %q", mock.written)
+	}
+}

dsn.go

@@ -80,11 +80,12 @@ type Config struct {
 	compress       bool // Enable zlib compression
 	tinyInt1IsBool bool // Treat signed TINYINT(1) as boolean
 
-	beforeConnect func(context.Context, *Config) error // Invoked before a connection is established
-	paramOrder    []string                             // Order of connection parameters parsed from the DSN
-	pubKey        *rsa.PublicKey                       // Server public key
-	timeTruncate  time.Duration                        // Truncate time.Time values to the specified duration
-	charsets      []string                             // Connection charset. When set, this will be set in SET NAMES <charset> query
+	beforeConnect     func(context.Context, *Config) error // Invoked before a connection is established
+	encodedAttributes string                               // Encoded connection attributes
+	paramOrder        []string                             // Order of connection parameters parsed from the DSN
+	pubKey            *rsa.PublicKey                       // Server public key
+	timeTruncate      time.Duration                        // Truncate time.Time values to the specified duration
+	charsets          []string                             // Connection charset. When set, this will be set in SET NAMES <charset> query
 }
 
 // Functional Options Pattern
@@ -146,6 +147,9 @@ func TimeTruncate(d time.Duration) Option {
 }
 
 // BeforeConnect sets the function to be invoked before a connection is established.
+// If the function changes [Config.Addr] while [Config.TLS] is non-nil and its
+// InsecureSkipVerify field is false, it must also update ServerName to match
+// the hostname in the new address.
 func BeforeConnect(fn func(context.Context, *Config) error) Option {
 	return func(cfg *Config) error {
 		cfg.beforeConnect = fn
@@ -263,6 +267,7 @@ func (cfg *Config) normalize() error {
 	if cfg.Logger == nil {
 		cfg.Logger = defaultLogger
 	}
+	cfg.encodedAttributes = encodeConnectionAttributes(cfg)
 
 	return nil
 }

dsn_test.go

@@ -33,6 +33,7 @@ func newTestConfig(update func(*Config)) *Config {
 	if update != nil {
 		update(cfg)
 	}
+	cfg.encodedAttributes = encodeConnectionAttributes(cfg)
 	return cfg
 }
 

packets.go

@@ -277,7 +277,7 @@ func (mc *mysqlConn) readHandshakePacket() (data []byte, capabilities capability
 }
 
 // initCapabilities initializes the capabilities based on server support and configuration
-func (mc *mysqlConn) initCapabilities(serverCapabilities capabilityFlag, serverExtCapabilities extendedCapabilityFlag, cfg *Config) {
+func (mc *mysqlConn) initCapabilities(serverCapabilities capabilityFlag, serverExtCapabilities extendedCapabilityFlag) {
 	clientCapabilities :=
 		clientMySQL |
 			clientLongFlag |
@@ -291,10 +291,10 @@ func (mc *mysqlConn) initCapabilities(serverCapabilities capabilityFlag, serverE
 			clientConnectAttrs |
 			clientDeprecateEOF
 
-	if cfg.ClientFoundRows {
+	if mc.cfg.ClientFoundRows {
 		clientCapabilities |= clientFoundRows
 	}
-	if cfg.compress {
+	if mc.cfg.compress {
 		clientCapabilities |= clientCompress
 	}
 	// To enable TLS / SSL
@@ -305,7 +305,7 @@ func (mc *mysqlConn) initCapabilities(serverCapabilities capabilityFlag, serverE
 	if mc.cfg.MultiStatements {
 		clientCapabilities |= clientMultiStatements
 	}
-	if n := len(cfg.DBName); n > 0 {
+	if n := len(mc.cfg.DBName); n > 0 {
 		clientCapabilities |= clientConnectWithDB
 	}
 
@@ -405,9 +405,9 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 
 	// Connection Attributes
 	if mc.capabilities&clientConnectAttrs != 0 {
-		connAttrsLen := len(mc.connector.encodedAttributes)
+		connAttrsLen := len(mc.cfg.encodedAttributes)
 		data = appendLengthEncodedInteger(data, uint64(connAttrsLen))
-		data = append(data, mc.connector.encodedAttributes...)
+		data = append(data, mc.cfg.encodedAttributes...)
 	}
 
 	// Send Auth packet

packets_test.go

@@ -96,11 +96,13 @@ var _ net.Conn = new(mockConn)
 
 func newRWMockConn(sequence uint8) (*mockConn, *mysqlConn) {
 	conn := new(mockConn)
-	connector := newConnector(NewConfig())
+	cfg := NewConfig()
+	if err := cfg.normalize(); err != nil {
+		panic(err)
+	}
 	mc := &mysqlConn{
 		buf:              newBuffer(),
-		cfg:              connector.cfg,
-		connector:        connector,
+		cfg:              cfg,
 		netConn:          conn,
 		closech:          make(chan struct{}),
 		maxAllowedPacket: defaultMaxAllowedPacket,