Fix: 🎯 Prevent Connection Pollution by Closing/Resetting Connection on Context Cancellation During BeginTx

connection.go

@@ -0,0 +1,223 @@
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"net"
+	"time"
+)
+
+type mysqlConn struct {
+	netConn           net.Conn
+	affectedRows      uint64
+	insertId          uint64
+	cfg               *Config
+	maxAllowedPacket  int
+	maxWriteSize      int
+	writeTimeout      time.Duration
+	flags             clientFlag
+	status            statusFlag
+	sequence          uint8
+	parseTime         bool
+	strict            bool
+	buf               buffer
+	closed            bool
+}
+
+func (mc *mysqlConn) Begin() (driver.Tx, error) {
+	return mc.begin(false)
+}
+
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+	defer mc.finish()
+
+	var level string
+	switch sql.IsolationLevel(opts.Isolation) {
+	case sql.LevelDefault:
+		// No isolation level specified, use server default
+		level = ""
+	case sql.LevelReadUncommitted:
+		level = "SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED"
+	case sql.LevelReadCommitted:
+		level = "SET TRANSACTION ISOLATION LEVEL READ COMMITTED"
+	case sql.LevelRepeatableRead:
+		level = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
+	case sql.LevelSerializable:
+		level = "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE"
+	default:
+		return nil, driver.ErrBadConn
+	}
+
+	if level != "" {
+		if err := mc.exec(level); err != nil {
+			// Check if context was canceled during isolation level setting
+			if ctx.Err() != nil {
+				// Connection state is ambiguous, mark as bad
+				mc.closed = true
+				return nil, driver.ErrBadConn
+			}
+			return nil, err
+		}
+	}
+
+	var readOnly string
+	if opts.ReadOnly {
+		readOnly = "START TRANSACTION READ ONLY"
+	} else {
+		readOnly = "START TRANSACTION"
+	}
+
+	err := mc.exec(readOnly)
+	if err != nil {
+		// Check if context was canceled during START TRANSACTION
+		if ctx.Err() != nil {
+			// Context was canceled - connection state is ambiguous
+			// The server may have started the transaction before we aborted
+			// Mark connection as bad to prevent pool pollution
+			mc.closed = true
+			return nil, driver.ErrBadConn
+		}
+		return nil, err
+	}
+
+	return &mysqlTx{mc}, nil
+}
+
+func (mc *mysqlConn) begin(readOnly bool) (driver.Tx, error) {
+	if mc.closed {
+		return nil, driver.ErrBadConn
+	}
+
+	var query string
+	if readOnly {
+		query = "START TRANSACTION READ ONLY"
+	} else {
+		query = "START TRANSACTION"
+	}
+
+	err := mc.exec(query)
+	if err != nil {
+		return nil, err
+	}
+	return &mysqlTx{mc}, nil
+}
+
+func (mc *mysqlConn) watchCancel(ctx context.Context) error {
+	if mc.closed {
+		return driver.ErrBadConn
+	}
+
+	if ctx.Done() == nil {
+		return nil
+	}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
+	if mc.cfg.InterpolateParams {
+		return nil
+	}
+
+	mc.startWatcher(ctx)
+	return nil
+}
+
+func (mc *mysqlConn) startWatcher(ctx context.Context) {
+	go func() {
+		<-ctx.Done()
+		if mc.netConn != nil {
+			mc.netConn.Close()
+		}
+	}()
+}
+
+func (mc *mysqlConn) finish() {
+	// Placeholder for cleanup after context watching
+}
+
+func (mc *mysqlConn) exec(query string) error {
+	// Execute query implementation
+	return mc.writeCommandPacketStr(comQuery, query)
+}
+
+func (mc *mysqlConn) writeCommandPacketStr(command byte, arg string) error {
+	if mc.closed {
+		return driver.ErrBadConn
+	}
+
+	// Reset packet sequence
+	mc.sequence = 0
+
+	// Send command packet
+	data := make([]byte, 1+len(arg))
+	data[0] = command
+	copy(data[1:], arg)
+
+	return mc.writePacket(data)
+}
+
+func (mc *mysqlConn) writePacket(data []byte) error {
+	if mc.closed {
+		return driver.ErrBadConn
+	}
+
+	pktLen := len(data)
+	if pktLen == 0 {
+		return nil
+	}
+
+	// Write packet header and data
+	header := make([]byte, 4)
+	header[0] = byte(pktLen)
+	header[1] = byte(pktLen >> 8)
+	header[2] = byte(pktLen >> 16)
+	header[3] = mc.sequence
+
+	mc.sequence++
+
+	if mc.writeTimeout > 0 {
+		mc.netConn.SetWriteDeadline(time.Now().Add(mc.writeTimeout))
+	}
+
+	if _, err := mc.netConn.Write(header); err != nil {
+		return err
+	}
+
+	if _, err := mc.netConn.Write(data); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+type mysqlTx struct {
+	mc *mysqlConn
+}
+
+func (tx *mysqlTx) Commit() error {
+	if tx.mc.closed {
+		return driver.ErrBadConn
+	}
+	return tx.mc.exec("COMMIT")
+}
+
+func (tx *mysqlTx) Rollback() error {
+	if tx.mc.closed {
+		return driver.ErrBadConn
+	}
+	return tx.mc.exec("ROLLBACK")
+}

connection_test.go

@@ -0,0 +1,172 @@
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"net"
+	"testing"
+	"time"
+)
+
+// mockConn simulates a network connection with controllable delay
+type mockConn struct {
+	net.Conn
+	writeDelay time.Duration
+	closed     bool
+}
+
+func (m *mockConn) Write(b []byte) (n int, err error) {
+	if m.closed {
+		return 0, errors.New("connection closed")
+	}
+	if m.writeDelay > 0 {
+		time.Sleep(m.writeDelay)
+	}
+	return len(b), nil
+}
+
+func (m *mockConn) Close() error {
+	m.closed = true
+	return nil
+}
+
+func (m *mockConn) Read(b []byte) (n int, err error) {
+	return 0, errors.New("not implemented")
+}
+
+func (m *mockConn) LocalAddr() net.Addr                { return nil }
+func (m *mockConn) RemoteAddr() net.Addr               { return nil }
+func (m *mockConn) SetDeadline(t time.Time) error      { return nil }
+func (m *mockConn) SetReadDeadline(t time.Time) error  { return nil }
+func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil }
+
+// TestBeginTxContextCancellation verifies that when BeginTx is canceled,
+// the connection is marked as bad and not returned to the pool
+func TestBeginTxContextCancellation(t *testing.T) {
+	// Create a mock connection with 100ms write delay
+	mockNet := &mockConn{
+		writeDelay: 100 * time.Millisecond,
+	}
+
+	mc := &mysqlConn{
+		netConn:      mockNet,
+		cfg:          &Config{},
+		writeTimeout: 0,
+		closed:       false,
+	}
+
+	// Create a context that times out before the write completes
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	// Attempt to begin transaction
+	_, err := mc.BeginTx(ctx, driver.TxOptions{})
+
+	// Should return an error
+	if err == nil {
+		t.Fatal("Expected error from BeginTx with canceled context")
+	}
+
+	// Should return driver.ErrBadConn to signal pool to discard connection
+	if !errors.Is(err, driver.ErrBadConn) {
+		t.Errorf("Expected driver.ErrBadConn, got %v", err)
+	}
+
+	// Connection should be marked as closed
+	if !mc.closed {
+		t.Error("Connection should be marked as closed after context cancellation")
+	}
+}
+
+// TestBeginTxContextCancellationDuringIsolationLevel verifies that
+// cancellation during isolation level setting also marks connection as bad
+func TestBeginTxContextCancellationDuringIsolationLevel(t *testing.T) {
+	mockNet := &mockConn{
+		writeDelay: 100 * time.Millisecond,
+	}
+
+	mc := &mysqlConn{
+		netConn:      mockNet,
+		cfg:          &Config{},
+		writeTimeout: 0,
+		closed:       false,
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	// Attempt to begin transaction with isolation level
+	_, err := mc.BeginTx(ctx, driver.TxOptions{
+		Isolation: driver.IsolationLevel(sql.LevelSerializable),
+	})
+
+	if err == nil {
+		t.Fatal("Expected error from BeginTx with canceled context")
+	}
+
+	if !errors.Is(err, driver.ErrBadConn) {
+		t.Errorf("Expected driver.ErrBadConn, got %v", err)
+	}
+
+	if !mc.closed {
+		t.Error("Connection should be marked as closed after context cancellation")
+	}
+}
+
+// TestBeginTxSuccessfulWithoutCancellation verifies normal operation
+func TestBeginTxSuccessfulWithoutCancellation(t *testing.T) {
+	mockNet := &mockConn{
+		writeDelay: 10 * time.Millisecond,
+	}
+
+	mc := &mysqlConn{
+		netConn:      mockNet,
+		cfg:          &Config{},
+		writeTimeout: 0,
+		closed:       false,
+	}
+
+	// Context with sufficient timeout
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	tx, err := mc.BeginTx(ctx, driver.TxOptions{})
+
+	if err != nil {
+		t.Fatalf("Expected successful BeginTx, got error: %v", err)
+	}
+
+	if tx == nil {
+		t.Fatal("Expected non-nil transaction")
+	}
+
+	if mc.closed {
+		t.Error("Connection should not be closed after successful BeginTx")
+	}
+}
+
+// TestBeginTxAlreadyClosedConnection verifies behavior with closed connection
+func TestBeginTxAlreadyClosedConnection(t *testing.T) {
+	mc := &mysqlConn{
+		netConn: &mockConn{},
+		cfg:     &Config{},
+		closed:  true,
+	}
+
+	ctx := context.Background()
+	_, err := mc.BeginTx(ctx, driver.TxOptions{})
+
+	if !errors.Is(err, driver.ErrBadConn) {
+		t.Errorf("Expected driver.ErrBadConn for closed connection, got %v", err)
+	}
+}