Skip to content

Commit

Permalink
Fix context cancelling without sending COMM_QUIT for authenticated co…
Browse files Browse the repository at this point in the history
…nnections
  • Loading branch information
minhquang4334 committed Dec 15, 2024
1 parent c9f41c0 commit 55ec77f
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 2 deletions.
13 changes: 11 additions & 2 deletions connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions connector.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 55ec77f

Please sign in to comment.