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

connection.go

@@ -0,0 +1,203 @@
+// 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 main
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"net"
+	"sync"
+	"time"
+)
+
+// Config holds the configuration for the driver
+type Config struct {
+	InterpolateParams bool
+}
+
+type clientFlag uint32
+type statusFlag uint16
+
+type mysqlConn struct {
+	netConn net.Conn
+	closed  bool
+	mu      sync.Mutex
+	cfg     *Config
+}
+
+func (mc *mysqlConn) Begin() (driver.Tx, error) {
+	return mc.begin(false)
+}
+
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	mc.mu.Lock()
+	if mc.closed {
+		mc.mu.Unlock()
+		return nil, driver.ErrBadConn
+	}
+	mc.mu.Unlock()
+
+	// Check if context is already canceled
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	var level string
+	switch sql.IsolationLevel(opts.Isolation) {
+	case sql.LevelDefault:
+		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, errors.New("invalid isolation level")
+	}
+
+	if level != "" {
+		err := mc.exec(ctx, level)
+		if ctx.Err() != nil {
+			mc.mu.Lock()
+			mc.closed = true
+			if mc.netConn != nil {
+				mc.netConn.Close()
+			}
+			mc.mu.Unlock()
+			return nil, driver.ErrBadConn
+		}
+		if err != nil {
+			var netErr net.Error
+			isTimeout := errors.As(err, &netErr) && netErr.Timeout()
+			if isTimeout || err.Error() == "write timeout" || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+				mc.mu.Lock()
+				mc.closed = true
+				if mc.netConn != nil {
+					mc.netConn.Close()
+				}
+				mc.mu.Unlock()
+				return nil, driver.ErrBadConn
+			}
+			return nil, err
+		}
+	}
+
+	var startTxQuery string
+	if opts.ReadOnly {
+		startTxQuery = "START TRANSACTION READ ONLY"
+	} else {
+		startTxQuery = "START TRANSACTION"
+	}
+
+	err := mc.exec(ctx, startTxQuery)
+	if ctx.Err() != nil {
+		mc.mu.Lock()
+		mc.closed = true
+		if mc.netConn != nil {
+			mc.netConn.Close()
+		}
+		mc.mu.Unlock()
+		return nil, driver.ErrBadConn
+	}
+	if err != nil {
+		var netErr net.Error
+		isTimeout := errors.As(err, &netErr) && netErr.Timeout()
+		if isTimeout || err.Error() == "write timeout" || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+			mc.mu.Lock()
+			mc.closed = true
+			if mc.netConn != nil {
+				mc.netConn.Close()
+			}
+			mc.mu.Unlock()
+			return nil, driver.ErrBadConn
+		}
+		return nil, err
+	}
+
+	return &mysqlTx{mc: mc}, nil
+}
+
+func (mc *mysqlConn) begin(readOnly bool) (driver.Tx, error) {
+	mc.mu.Lock()
+	if mc.closed {
+		mc.mu.Unlock()
+		return nil, driver.ErrBadConn
+	}
+	mc.mu.Unlock()
+
+	var query string
+	if readOnly {
+		query = "START TRANSACTION READ ONLY"
+	} else {
+		query = "START TRANSACTION"
+	}
+
+	err := mc.exec(context.Background(), query)
+	if err != nil {
+		return nil, err
+	}
+	return &mysqlTx{mc: mc}, nil
+}
+
+func (mc *mysqlConn) exec(ctx context.Context, query string) error {
+	mc.mu.Lock()
+	closed := mc.closed
+	netConn := mc.netConn
+	mc.mu.Unlock()
+
+	if closed || netConn == nil {
+		return driver.ErrBadConn
+	}
+
+	// Make the net write operation context-aware using SetWriteDeadline
+	if deadline, ok := ctx.Deadline(); ok {
+		netConn.SetWriteDeadline(deadline)
+	} else if ctx.Done() != nil {
+		// Context has cancellation but no deadline (cancel-only).
+		// Spawn a watcher goroutine to close the net connection if context is canceled during Write.
+		done := make(chan struct{})
+		defer close(done)
+		go func() {
+			select {
+			case <-ctx.Done():
+				mc.mu.Lock()
+				if mc.netConn != nil {
+					mc.netConn.Close()
+				}
+				mc.mu.Unlock()
+			case <-done:
+			}
+		}()
+	} else {
+		netConn.SetWriteDeadline(time.Time{})
+	}
+
+	_, err := netConn.Write([]byte(query))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+type mysqlTx struct {
+	mc *mysqlConn
+}
+
+func (tx *mysqlTx) Commit() error {
+	return nil
+}
+
+func (tx *mysqlTx) Rollback() error {
+	return nil
+}