From acbfe7855976afab773d5950169e3626c2a86691 Mon Sep 17 00:00:00 2001
From: Jakub Nyckowski <jakub.nyckowski@goteleport.com>
Date: Wed, 9 Mar 2022 14:22:55 -0500
Subject: [PATCH] Update go-mysql package (#10997)

Send COM_QUIT message when closing MySQL connection.

Backport #10984 to branch/v9
---
 go.mod                         |  2 +-
 go.sum                         |  4 ++--
 lib/srv/db/access_test.go      | 32 +++++++++++++++++++++++++++++++-
 lib/srv/db/mysql/test.go       | 25 +++++++++++++++++++++++++
 lib/srv/db/proxy_test.go       |  4 +++-
 lib/srv/db/proxyserver_test.go |  8 ++++----
 6 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 37bebeac23f7a..16ebbd81fb1a2 100644
--- a/go.mod
+++ b/go.mod
@@ -271,7 +271,7 @@ replace (
 	github.com/go-redis/redis/v8 => github.com/gravitational/redis/v8 v8.11.5-0.20220211010318-7af711b76a91
 	github.com/gogo/protobuf => github.com/gravitational/protobuf v1.3.2-0.20201123192827-2b9fcfaffcbf
 	github.com/gravitational/teleport/api => ./api
-	github.com/siddontang/go-mysql v1.1.0 => github.com/gravitational/go-mysql v1.1.1-teleport.1
+	github.com/siddontang/go-mysql v1.1.0 => github.com/gravitational/go-mysql v1.1.1-teleport.2
 	github.com/sirupsen/logrus => github.com/gravitational/logrus v1.4.4-0.20210817004754-047e20245621
 	github.com/vulcand/predicate => github.com/gravitational/predicate v1.2.1
 )
diff --git a/go.sum b/go.sum
index 27758b1196c28..468cf93640ef3 100644
--- a/go.sum
+++ b/go.sum
@@ -454,8 +454,8 @@ github.com/gravitational/form v0.0.0-20151109031454-c4048f792f70 h1:To76nCJtM3DI
 github.com/gravitational/form v0.0.0-20151109031454-c4048f792f70/go.mod h1:88hFR45MpUd23d2vNWE/dYtesU50jKsbz0I9kH7UaBY=
 github.com/gravitational/go-mssqldb v0.11.1-0.20220202000043-bec708e9bfd0 h1:DC+S+j/tBs/0MnQCC5j7GWWbMGcR3ca5v75ksAU1LJM=
 github.com/gravitational/go-mssqldb v0.11.1-0.20220202000043-bec708e9bfd0/go.mod h1:iiK0YP1ZeepvmBQk/QpLEhhTNJgfzrpArPY/aFvc9yU=
-github.com/gravitational/go-mysql v1.1.1-teleport.1 h1:062V8u0juCyUvpYMdkYch8JDDw7wf5rdhKaIfhnojDg=
-github.com/gravitational/go-mysql v1.1.1-teleport.1/go.mod h1:re0JQZ1Cy5dVlIDGq0YksfDIla/GRZlxqOoC0XPSSGE=
+github.com/gravitational/go-mysql v1.1.1-teleport.2 h1:XZ36BZ7BgslA5ZCyCHjpc1wilFITThIH7cLcbLWKWzM=
+github.com/gravitational/go-mysql v1.1.1-teleport.2/go.mod h1:re0JQZ1Cy5dVlIDGq0YksfDIla/GRZlxqOoC0XPSSGE=
 github.com/gravitational/go-oidc v0.0.5 h1:kxsCknoOZ+KqIAoYLLdHuQcvcc+SrQlnT7xxIM8oo6o=
 github.com/gravitational/go-oidc v0.0.5/go.mod h1:SevmOUNdOB0aD9BAIgjptZ6oHkKxMZZgA70nwPfgU/w=
 github.com/gravitational/kingpin v2.1.11-0.20190130013101-742f2714c145+incompatible h1:CfyZl3nyo9K5lLqOmqvl9/IElY1UCnOWKZiQxJ8HKdA=
diff --git a/lib/srv/db/access_test.go b/lib/srv/db/access_test.go
index 321c04f1b182c..e859ba60e47dc 100644
--- a/lib/srv/db/access_test.go
+++ b/lib/srv/db/access_test.go
@@ -428,6 +428,31 @@ func TestAccessMySQLChangeUser(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestMySQLCloseConnection(t *testing.T) {
+	ctx := context.Background()
+	testCtx := setupTestContext(ctx, t, withSelfHostedMySQL("mysql"))
+	go testCtx.startHandlingConnections()
+
+	// Create user/role with the requested permissions.
+	testCtx.createUserAndRole(ctx, t, "alice", "admin", []string{"alice"}, []string{types.Wildcard})
+
+	// Connect to the database as this user.
+	mysqlConn, err := testCtx.mysqlClient("alice", "mysql", "alice")
+	require.NoError(t, err)
+
+	_, err = mysqlConn.Execute("select 1")
+	require.NoError(t, err)
+
+	// Close connection to DB proxy
+	err = mysqlConn.Close()
+	require.NoError(t, err)
+
+	// DB proxy should close the DB connection and send COM_QUIT message.
+	require.Eventually(t, func() bool {
+		return testCtx.mysql["mysql"].db.ConnsClosed()
+	}, 2*time.Second, 100*time.Millisecond)
+}
+
 // TestAccessRedisAUTHDefaultCmd checks if empty user can log in to Redis as default.
 func TestAccessRedisAUTHDefaultCmd(t *testing.T) {
 	ctx := context.Background()
@@ -492,6 +517,9 @@ func TestAccessMySQLServerPacket(t *testing.T) {
 	// in a way that previously would cause our packet parsing logic to fail.
 	_, err = mysqlConn.Execute("show tables")
 	require.NoError(t, err)
+
+	err = mysqlConn.Close()
+	require.NoError(t, err)
 }
 
 // TestGCPRequireSSL tests connecting to GCP Cloud SQL Postgres and MySQL
@@ -1929,7 +1957,9 @@ func withSelfHostedMySQL(name string) withDatabaseOption {
 		})
 		require.NoError(t, err)
 		go mysqlServer.Serve()
-		t.Cleanup(func() { mysqlServer.Close() })
+		t.Cleanup(func() {
+			require.NoError(t, mysqlServer.Close())
+		})
 		database, err := types.NewDatabaseV3(types.Metadata{
 			Name: name,
 		}, types.DatabaseSpecV3{
diff --git a/lib/srv/db/mysql/test.go b/lib/srv/db/mysql/test.go
index 18b4ae22db6c7..520bdc8dc31de 100644
--- a/lib/srv/db/mysql/test.go
+++ b/lib/srv/db/mysql/test.go
@@ -19,6 +19,7 @@ package mysql
 import (
 	"crypto/tls"
 	"net"
+	"sync"
 	"sync/atomic"
 
 	"github.com/gravitational/teleport/lib/defaults"
@@ -62,6 +63,11 @@ type TestServer struct {
 	tlsConfig *tls.Config
 	log       logrus.FieldLogger
 	handler   *testHandler
+
+	// serverConnsMtx is a mutex that guards serverConns.
+	serverConnsMtx sync.Mutex
+	// serverConns holds all connections created by the server.
+	serverConns []*server.Conn
 }
 
 // NewTestServer returns a new instance of a test MySQL server.
@@ -152,6 +158,11 @@ func (s *TestServer) handleConnection(conn net.Conn) error {
 	if err != nil {
 		return trace.Wrap(err)
 	}
+
+	s.serverConnsMtx.Lock()
+	s.serverConns = append(s.serverConns, serverConn)
+	s.serverConnsMtx.Unlock()
+
 	for {
 		if serverConn.Closed() {
 			return nil
@@ -196,6 +207,20 @@ func (s *TestServer) Close() error {
 	return s.listener.Close()
 }
 
+// ConnsClosed returns true if all connections has been correctly closed (message COM_QUIT), false otherwise.
+func (s *TestServer) ConnsClosed() bool {
+	s.serverConnsMtx.Lock()
+	defer s.serverConnsMtx.Unlock()
+
+	for _, conn := range s.serverConns {
+		if !conn.Closed() {
+			return false
+		}
+	}
+
+	return true
+}
+
 type testHandler struct {
 	server.EmptyHandler
 	log logrus.FieldLogger
diff --git a/lib/srv/db/proxy_test.go b/lib/srv/db/proxy_test.go
index c8baf6d720c60..277fe77009850 100644
--- a/lib/srv/db/proxy_test.go
+++ b/lib/srv/db/proxy_test.go
@@ -197,7 +197,9 @@ func TestProxyClientDisconnectDueToLockInForce(t *testing.T) {
 		Target: types.LockTarget{User: "alice"},
 	})
 	require.NoError(t, err)
-	testCtx.authServer.UpsertLock(ctx, lock)
+
+	err = testCtx.authServer.UpsertLock(ctx, lock)
+	require.NoError(t, err)
 
 	waitForEvent(t, testCtx, events.ClientDisconnectCode)
 	err = mysql.Ping()
diff --git a/lib/srv/db/proxyserver_test.go b/lib/srv/db/proxyserver_test.go
index 33fe5a69bce4e..df6500b49641a 100644
--- a/lib/srv/db/proxyserver_test.go
+++ b/lib/srv/db/proxyserver_test.go
@@ -91,10 +91,10 @@ func TestProxyConnectionLimiting(t *testing.T) {
 			t.Run("limit can be hit", func(t *testing.T) {
 				for i := 0; i < connLimitNumber; i++ {
 					// Try to connect to the database.
-					pgConn, err := tt.connect()
+					dbConn, err := tt.connect()
 					require.NoError(t, err)
 
-					connsClosers = append(connsClosers, pgConn)
+					connsClosers = append(connsClosers, dbConn)
 				}
 
 				// This connection should go over the limit.
@@ -114,9 +114,9 @@ func TestProxyConnectionLimiting(t *testing.T) {
 				require.NoError(t, err)
 
 				// Create a new connection. We do not expect an error here as we have just closed one.
-				pgConn, err := tt.connect()
+				dbConn, err := tt.connect()
 				require.NoError(t, err)
-				connsClosers = append(connsClosers, pgConn)
+				connsClosers = append(connsClosers, dbConn)
 
 				// Here the limit should be reached again.
 				_, err = tt.connect()