observer, net: read packet from SQL port in health check (#578)

Author: djshow832

pkg/balance/observer/health_check.go

@@ -86,6 +86,9 @@ func (dhc *DefaultHealthCheck) checkSqlPort(ctx context.Context, addr string, bh
 		if err = conn.SetReadDeadline(time.Now().Add(dhc.cfg.DialTimeout)); err != nil {
 			return err
 		}
+		if err = pnet.CheckSqlPort(conn); err != nil {
+			return err
+		}
 		if ignoredErr := conn.Close(); ignoredErr != nil && !pnet.IsDisconnectError(ignoredErr) {
 			dhc.logger.Warn("close connection in health check failed", zap.Error(ignoredErr))
 		}

pkg/balance/observer/health_check_test.go

@@ -9,9 +9,11 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/go-mysql-org/go-mysql/packet"
 	"github.com/pingcap/tiproxy/lib/util/logger"
 	"github.com/pingcap/tiproxy/lib/util/waitgroup"
 	"github.com/pingcap/tiproxy/pkg/testkit"
@@ -75,6 +77,12 @@ func TestHealthCheck(t *testing.T) {
 	health = hc.Check(context.Background(), backend.sqlAddr, info)
 	require.True(t, health.Healthy)
 
+	backend.setSqlResp(false)
+	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	require.False(t, health.Healthy)
+	backend.setSqlResp(true)
+	health = hc.Check(context.Background(), backend.sqlAddr, info)
+	require.True(t, health.Healthy)
 	backend.close()
 }
 
@@ -88,6 +96,7 @@ type backendServer struct {
 	wg         waitgroup.WaitGroup
 	ip         string
 	statusPort uint
+	sqlResp    atomic.Bool
 }
 
 func newBackendServer(t *testing.T) (*backendServer, *BackendInfo) {
@@ -97,6 +106,7 @@ func newBackendServer(t *testing.T) (*backendServer, *BackendInfo) {
 	backend.startHTTPServer()
 	backend.setHTTPResp(true)
 	backend.setHTTPRespBody("")
+	backend.setSqlResp(true)
 	backend.startSQLServer()
 	return backend, &BackendInfo{
 		IP:         backend.ip,
@@ -133,6 +143,10 @@ func (srv *backendServer) stopHTTPServer() {
 	require.NoError(srv.t, err)
 }
 
+func (srv *backendServer) setSqlResp(sqlResp bool) {
+	srv.sqlResp.Store(sqlResp)
+}
+
 func (srv *backendServer) startSQLServer() {
 	srv.sqlListener, srv.sqlAddr = testkit.StartListener(srv.t, srv.sqlAddr)
 	srv.wg.Run(func() {
@@ -142,6 +156,11 @@ func (srv *backendServer) startSQLServer() {
 				// listener is closed
 				break
 			}
+			if srv.sqlResp.Load() {
+				data := []byte{0, 0, 0, 0, 0}
+				c := packet.NewConn(conn)
+				require.NoError(srv.t, c.WritePacket(data))
+			}
 			_ = conn.Close()
 		}
 	})

pkg/proxy/net/mysql.go

@@ -358,32 +358,17 @@ func ParseChangeUser(data []byte, capability Capability) (*ChangeUserReq, error)
 	return req, err
 }
 
-// ReadServerVersion only reads server version.
-func ReadServerVersion(conn net.Conn) (string, error) {
+// CheckSqlPort checks whether the SQL port is available.
+func CheckSqlPort(conn net.Conn) error {
 	c := packet.NewConn(conn)
 	data, err := c.ReadPacket()
 	if err != nil {
-		return "", err
+		return err
 	}
 	if data[0] == ErrHeader.Byte() {
-		return "", errors.New("read initial handshake error")
+		return errors.New("read initial handshake error")
 	}
-	pos := 1
-	version := data[pos : pos+bytes.IndexByte(data[pos:], 0x00)]
-	return string(version), nil
-}
-
-// WriteServerVersion only writes server version. It's only used for testing.
-func WriteServerVersion(conn net.Conn, serverVersion string) error {
-	data := make([]byte, 0, 128)
-	data = append(data, []byte{0, 0, 0, 0}...)
-	// min version 10
-	data = append(data, 10)
-	// server version[NUL]
-	data = append(data, serverVersion...)
-	data = append(data, 0)
-	c := packet.NewConn(conn)
-	return c.WritePacket(data)
+	return nil
 }
 
 // ParseOKPacket parses an OK packet and only returns server status.

pkg/proxy/net/mysql_test.go

@@ -4,11 +4,14 @@
 package net
 
 import (
+	"net"
 	"testing"
 
 	"github.com/go-mysql-org/go-mysql/mysql"
+	"github.com/go-mysql-org/go-mysql/packet"
 	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/lib/util/logger"
+	"github.com/pingcap/tiproxy/pkg/testkit"
 	"github.com/stretchr/testify/require"
 )
 
@@ -72,3 +75,38 @@ func TestMySQLError(t *testing.T) {
 	require.True(t, errors.Is(errors.Wrap(ErrHandshakeTLS, myerr), ErrHandshakeTLS))
 	require.True(t, errors.Is(errors.Wrap(myerr, ErrHandshakeTLS), ErrHandshakeTLS))
 }
+
+func TestCheckSqlPort(t *testing.T) {
+	// normal
+	testkit.TestTCPConn(t,
+		func(t *testing.T, c net.Conn) {
+			err := CheckSqlPort(c)
+			require.NoError(t, err)
+		},
+		func(t *testing.T, c net.Conn) {
+			data := []byte{0, 0, 0, 0, 0}
+			conn := packet.NewConn(c)
+			require.NoError(t, conn.WritePacket(data))
+		}, 1)
+
+	// no write
+	testkit.TestTCPConn(t,
+		func(t *testing.T, c net.Conn) {
+			err := CheckSqlPort(c)
+			require.Error(t, err)
+		},
+		func(t *testing.T, c net.Conn) {
+		}, 1)
+
+	// write error code
+	testkit.TestTCPConn(t,
+		func(t *testing.T, c net.Conn) {
+			err := CheckSqlPort(c)
+			require.Error(t, err)
+		},
+		func(t *testing.T, c net.Conn) {
+			data := []byte{0, 0, 0, 0, 0xff}
+			conn := packet.NewConn(c)
+			require.NoError(t, conn.WritePacket(data))
+		}, 1)
+}