backend: use BackendConnMgr as ConnContext

pkg/manager/router/backend_selector.go

@@ -26,10 +26,10 @@ func (bs *BackendSelector) Reset() {
 }
 
 func (bs *BackendSelector) Next() string {
-	if len(bs.cur) > 0 {
+	bs.cur = bs.routeOnce(bs.excluded)
+	if bs.cur != "" {
 		bs.excluded = append(bs.excluded, bs.cur)
 	}
-	bs.cur = bs.routeOnce(bs.excluded)
 	return bs.cur
 }
 

pkg/manager/router/router_static.go

@@ -17,18 +17,18 @@ package router
 var _ Router = &StaticRouter{}
 
 type StaticRouter struct {
-	addr []string
-	cnt  int
+	addrs []string
+	cnt   int
 }
 
 func NewStaticRouter(addr []string) *StaticRouter {
-	return &StaticRouter{addr: addr}
+	return &StaticRouter{addrs: addr}
 }
 
 func (r *StaticRouter) GetBackendSelector() BackendSelector {
 	return BackendSelector{
 		routeOnce: func(excluded []string) string {
-			for _, addr := range r.addr {
+			for _, addr := range r.addrs {
 				found := false
 				for _, e := range excluded {
 					if e == addr {

pkg/proxy/backend/authenticator.go

@@ -19,7 +19,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
-	"sync"
 	"time"
 
 	"github.com/pingcap/TiProxy/lib/util/errors"
@@ -47,19 +46,14 @@ const SupportedServerCapabilities = pnet.ClientLongPassword | pnet.ClientFoundRo
 
 // Authenticator handshakes with the client and the backend.
 type Authenticator struct {
-	backendTLSConfig            *tls.Config
-	ctxmap                      sync.Map
-	supportedServerCapabilities pnet.Capability
-	dbname                      string // default database name
-	serverAddr                  string
-	clientAddr                  string
-	user                        string
-	attrs                       map[string]string
-	salt                        []byte
-	capability                  uint32 // client capability
-	collation                   uint8
-	proxyProtocol               bool
-	requireBackendTLS           bool
+	dbname            string // default database name
+	user              string
+	attrs             map[string]string
+	salt              []byte
+	capability        uint32 // client capability
+	collation         uint8
+	proxyProtocol     bool
+	requireBackendTLS bool
 }
 
 func (auth *Authenticator) String() string {
@@ -104,17 +98,16 @@ func (auth *Authenticator) verifyBackendCaps(logger *zap.Logger, backendCapabili
 
 type backendIOGetter func(ctx ConnContext, auth *Authenticator, resp *pnet.HandshakeResp, timeout time.Duration) (*pnet.PacketIO, error)
 
-func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, clientIO *pnet.PacketIO, handshakeHandler HandshakeHandler,
+func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, cctx ConnContext, clientIO *pnet.PacketIO, handshakeHandler HandshakeHandler,
 	getBackendIO backendIOGetter, frontendTLSConfig, backendTLSConfig *tls.Config) error {
 	clientIO.ResetSequence()
-	auth.clientAddr = clientIO.SourceAddr().String()
 
-	proxyCapability := auth.supportedServerCapabilities
+	proxyCapability := handshakeHandler.GetCapability()
 	if frontendTLSConfig == nil {
 		proxyCapability ^= pnet.ClientSSL
 	}
 
-	if err := clientIO.WriteInitialHandshake(proxyCapability.Uint32(), auth.salt, mysql.AuthNativePassword); err != nil {
+	if err := clientIO.WriteInitialHandshake(proxyCapability, auth.salt, mysql.AuthNativePassword); err != nil {
 		return err
 	}
 	pkt, isSSL, err := clientIO.ReadSSLRequestOrHandshakeResp()
@@ -151,17 +144,17 @@ func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, clientIO *pnet
 	}
 	auth.capability = commonCaps.Uint32()
 
-	resp := pnet.ParseHandshakeResponse(pkt)
-	if err = handshakeHandler.HandleHandshakeResp(auth, resp); err != nil {
+	clientResp := pnet.ParseHandshakeResponse(pkt)
+	if err = handshakeHandler.HandleHandshakeResp(cctx, clientResp); err != nil {
 		return err
 	}
-	auth.user = resp.User
-	auth.dbname = resp.DB
-	auth.collation = resp.Collation
-	auth.attrs = resp.Attrs
+	auth.user = clientResp.User
+	auth.dbname = clientResp.DB
+	auth.collation = clientResp.Collation
+	auth.attrs = clientResp.Attrs
 
 	// In case of testing, backendIO is passed manually that we don't want to bother with the routing logic.
-	backendIO, err := getBackendIO(auth, auth, resp, 5*time.Second)
+	backendIO, err := getBackendIO(cctx, auth, clientResp, 5*time.Second)
 	if err != nil {
 		return err
 	}
@@ -173,11 +166,10 @@ func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, clientIO *pnet
 	}
 
 	// read backend initial handshake
-	_, backendCapabilityU, err := auth.readInitialHandshake(backendIO)
+	_, backendCapability, err := auth.readInitialHandshake(backendIO)
 	if err != nil {
 		return err
 	}
-	backendCapability := pnet.Capability(backendCapabilityU)
 
 	if err := auth.verifyBackendCaps(logger, backendCapability); err != nil {
 		return err
@@ -195,38 +187,12 @@ func (auth *Authenticator) handshakeFirstTime(logger *zap.Logger, clientIO *pnet
 		logger.Info("backend does not support capabilities from proxy", zap.Stringer("common", common), zap.Stringer("proxy", proxyCapability^common), zap.Stringer("backend", backendCapability^common))
 	}
 
-	// Send an unknown auth plugin so that the backend will request the auth data again.
-	resp.AuthPlugin = unknownAuthPlugin
-	resp.Capability = auth.capability
-
-	if backendCapability&pnet.ClientSSL != 0 && backendTLSConfig != nil {
-		resp.Capability |= mysql.ClientSSL
-		pkt = pnet.MakeHandshakeResponse(resp)
-		// write SSL Packet
-		if err := backendIO.WritePacket(pkt[:32], true); err != nil {
-			return err
-		}
-		auth.backendTLSConfig = backendTLSConfig.Clone()
-		addr := backendIO.RemoteAddr().String()
-		if auth.serverAddr != "" {
-			// NOTE: should use DNS name as much as possible
-			// Usually certs are signed with domain instead of IP addrs
-			// And `RemoteAddr()` will return IP addr
-			addr = auth.serverAddr
-		}
-		host, _, err := net.SplitHostPort(addr)
-		if err == nil {
-			auth.backendTLSConfig.ServerName = host
-		}
-		if err = backendIO.ClientTLSHandshake(auth.backendTLSConfig); err != nil {
-			return err
-		}
-	} else {
-		pkt = pnet.MakeHandshakeResponse(resp)
-	}
-
 	// forward client handshake resp
-	if err := backendIO.WritePacket(pkt, true); err != nil {
+	if err := auth.writeAuthHandshake(
+		backendIO, backendTLSConfig, backendCapability,
+		// send an unknown auth plugin so that the backend will request the auth data again.
+		unknownAuthPlugin, nil, 0,
+	); err != nil {
 		return err
 	}
 
@@ -258,7 +224,7 @@ func forwardMsg(srcIO, destIO *pnet.PacketIO) (data []byte, err error) {
 	return
 }
 
-func (auth *Authenticator) handshakeSecondTime(logger *zap.Logger, clientIO, backendIO *pnet.PacketIO, sessionToken string) error {
+func (auth *Authenticator) handshakeSecondTime(logger *zap.Logger, clientIO, backendIO *pnet.PacketIO, backendTLSConfig *tls.Config, sessionToken string) error {
 	if len(sessionToken) == 0 {
 		return errors.New("session token is empty")
 	}
@@ -268,24 +234,26 @@ func (auth *Authenticator) handshakeSecondTime(logger *zap.Logger, clientIO, bac
 		return err
 	}
 
-	_, serverCapability, err := auth.readInitialHandshake(backendIO)
+	_, backendCapability, err := auth.readInitialHandshake(backendIO)
 	if err != nil {
 		return err
 	}
 
-	if err := auth.verifyBackendCaps(logger, pnet.Capability(serverCapability)); err != nil {
+	if err := auth.verifyBackendCaps(logger, pnet.Capability(backendCapability)); err != nil {
 		return err
 	}
 
-	tokenBytes := hack.Slice(sessionToken)
-	if err = auth.writeAuthHandshake(backendIO, tokenBytes, serverCapability&mysql.ClientSSL != 0); err != nil {
+	if err = auth.writeAuthHandshake(
+		backendIO, backendTLSConfig, backendCapability,
+		mysql.AuthTiDBSessionToken, hack.Slice(sessionToken), mysql.ClientPluginAuth,
+	); err != nil {
 		return err
 	}
 
 	return auth.handleSecondAuthResult(backendIO)
 }
 
-func (auth *Authenticator) readInitialHandshake(backendIO *pnet.PacketIO) (serverPkt []byte, capability uint32, err error) {
+func (auth *Authenticator) readInitialHandshake(backendIO *pnet.PacketIO) (serverPkt []byte, capability pnet.Capability, err error) {
 	if serverPkt, err = backendIO.ReadPacket(); err != nil {
 		return
 	}
@@ -297,32 +265,49 @@ func (auth *Authenticator) readInitialHandshake(backendIO *pnet.PacketIO) (serve
 	return
 }
 
-func (auth *Authenticator) writeAuthHandshake(backendIO *pnet.PacketIO, authData []byte, tls bool) error {
+func (auth *Authenticator) writeAuthHandshake(
+	backendIO *pnet.PacketIO,
+	backendTLSConfig *tls.Config,
+	backendCapability pnet.Capability,
+	authPlugin string,
+	authData []byte,
+	authCap uint32,
+) error {
 	// Always handshake with SSL enabled and enable auth_plugin.
 	resp := &pnet.HandshakeResp{
 		User:       auth.user,
 		DB:         auth.dbname,
-		AuthPlugin: mysql.AuthTiDBSessionToken,
 		Attrs:      auth.attrs,
-		AuthData:   authData,
-		Capability: auth.capability | mysql.ClientSSL | mysql.ClientPluginAuth,
 		Collation:  auth.collation,
+		AuthData:   authData,
+		Capability: auth.capability | authCap,
+		AuthPlugin: authPlugin,
 	}
-	data := pnet.MakeHandshakeResponse(resp)
 
-	if tls && auth.backendTLSConfig != nil {
-		// write SSL req
-		if err := backendIO.WritePacket(data[:32], true); err != nil {
+	var pkt []byte
+	if backendCapability&pnet.ClientSSL != 0 && backendTLSConfig != nil {
+		pkt = pnet.MakeHandshakeResponse(resp)
+		resp.Capability |= mysql.ClientSSL
+		// write SSL Packet
+		if err := backendIO.WritePacket(pkt[:32], true); err != nil {
 			return err
 		}
 		// Send TLS / SSL request packet. The server must have supported TLS.
-		if err := backendIO.ClientTLSHandshake(auth.backendTLSConfig); err != nil {
+		tcfg := backendTLSConfig.Clone()
+		addr := backendIO.RemoteAddr().String()
+		host, _, err := net.SplitHostPort(addr)
+		if err == nil {
+			tcfg.ServerName = host
+		}
+		if err := backendIO.ClientTLSHandshake(tcfg); err != nil {
 			return err
 		}
+	} else {
+		pkt = pnet.MakeHandshakeResponse(resp)
 	}
 
 	// write handshake resp
-	return backendIO.WritePacket(data, true)
+	return backendIO.WritePacket(pkt, true)
 }
 
 func (auth *Authenticator) handleSecondAuthResult(backendIO *pnet.PacketIO) error {
@@ -353,23 +338,3 @@ func (auth *Authenticator) changeUser(username, db string) {
 func (auth *Authenticator) updateCurrentDB(db string) {
 	auth.dbname = db
 }
-
-func (auth *Authenticator) ClientAddr() string {
-	return auth.clientAddr
-}
-
-func (auth *Authenticator) ServerAddr() string {
-	return auth.serverAddr
-}
-
-func (auth *Authenticator) SetValue(key, val any) {
-	auth.ctxmap.Store(key, val)
-}
-
-func (auth *Authenticator) Value(key any) any {
-	v, ok := auth.ctxmap.Load(key)
-	if !ok {
-		return nil
-	}
-	return v
-}