From 55ec77fabb0a3a78fa46a5f82169522c820c7216 Mon Sep 17 00:00:00 2001
From: minhquang4334 <minhquang4334@gmail.com>
Date: Sun, 15 Dec 2024 12:03:15 +0700
Subject: [PATCH] Fix context cancelling without sending COMM_QUIT for
 authenticated connections

---
 connection.go | 13 +++++++++++--
 connector.go  |  1 +
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/connection.go b/connection.go
index acc62708..f2899567 100644
--- a/connection.go
+++ b/connection.go
@@ -40,6 +40,7 @@ type mysqlConn struct {
 
 	// for context support (Go 1.8+)
 	watching bool
+	authed   bool
 	watcher  chan<- context.Context
 	closech  chan struct{}
 	finished chan<- struct{}
@@ -436,9 +437,15 @@ func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 }
 
 // cancel is called when the query has canceled.
-func (mc *mysqlConn) cancel(err error) {
+func (mc *mysqlConn) cancel(err error) error {
 	mc.canceled.Set(err)
+	if mc.authed {
+		return mc.Close()
+	}
+
 	mc.cleanup()
+
+	return nil
 }
 
 // finish is called when the query has succeeded.
@@ -624,7 +631,9 @@ func (mc *mysqlConn) startWatcher() {
 
 			select {
 			case <-ctx.Done():
-				mc.cancel(ctx.Err())
+				if err := mc.cancel(ctx.Err()); err != nil {
+					mc.log(err)
+				}
 			case <-finished:
 			case <-mc.closech:
 				return
diff --git a/connector.go b/connector.go
index 769b3adc..08789c2c 100644
--- a/connector.go
+++ b/connector.go
@@ -169,6 +169,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		mc.cleanup()
 		return nil, err
 	}
+	mc.authed = true
 
 	if mc.cfg.MaxAllowedPacket > 0 {
 		mc.maxAllowedPacket = mc.cfg.MaxAllowedPacket