Implement AuthenticationHandler for custom auth mechanisms

driver/driver_options_test.go

@@ -266,8 +266,8 @@ func TestDriverOptions_namedValueChecker(t *testing.T) {
 }
 
 func createMockServer(t *testing.T) *testServer {
-	inMemProvider := server.NewInMemoryProvider()
-	require.NoError(t, inMemProvider.AddUser(*testUser, *testPassword))
+	authHandler := server.NewInMemoryAuthenticationHandler()
+	require.NoError(t, authHandler.AddUser(*testUser, *testPassword))
 	defaultServer := server.NewDefaultServer()
 
 	l, err := net.Listen("tcp", "127.0.0.1:3307")
@@ -285,7 +285,7 @@ func createMockServer(t *testing.T) *testServer {
 			}
 
 			go func() {
-				co, err := s.NewCustomizedConn(conn, inMemProvider, handler)
+				co, err := s.NewCustomizedConn(conn, authHandler, handler)
 				if err != nil {
 					return
 				}

mysql/util.go

@@ -56,13 +56,13 @@ func CalcNativePassword(scramble, password []byte) []byte {
 	return Xor(scrambleHash, stage1)
 }
 
-// Xor modifies hash1 in-place with XOR against hash2
+// Xor returns a new slice with hash1 XOR hash2, wrapping hash2 if hash1 is longer.
 func Xor(hash1 []byte, hash2 []byte) []byte {
-	l := min(len(hash1), len(hash2))
-	for i := range l {
-		hash1[i] ^= hash2[i]
+	result := make([]byte, len(hash1))
+	for i := range hash1 {
+		result[i] = hash1[i] ^ hash2[i%len(hash2)]
 	}
-	return hash1
+	return result
 }
 
 // hash_stage1 = xor(reply, sha1(public_seed, hash_stage2))

server/auth.go

@@ -30,15 +30,15 @@ func (c *Conn) compareAuthData(authPluginName string, clientAuthData []byte) err
 	return c.serverConf.authProvider.Authenticate(c, authPluginName, clientAuthData)
 }
 
-func (c *Conn) acquirePassword() error {
-	if c.credential.Password != "" {
+func (c *Conn) acquireCredential() error {
+	if len(c.credential.Passwords) > 0 {
 		return nil
 	}
-	credential, found, err := c.credentialProvider.GetCredential(c.user)
+	credential, found, err := c.authHandler.GetCredential(c.user)
 	if err != nil {
 		return err
 	}
-	if !found {
+	if !found || len(credential.Passwords) == 0 {
 		return mysql.NewDefaultError(mysql.ER_NO_SUCH_USER, c.user, c.RemoteAddr().String())
 	}
 	c.credential = credential
@@ -67,26 +67,32 @@ func scrambleValidation(cached, nonce, scramble []byte) bool {
 
 func (c *Conn) compareNativePasswordAuthData(clientAuthData []byte, credential Credential) error {
 	if len(clientAuthData) == 0 {
-		if credential.Password == "" {
+		if credential.hasEmptyPassword() {
 			return nil
 		}
 		return ErrAccessDeniedNoPassword
 	}
 
-	password, err := mysql.DecodePasswordHex(c.credential.Password)
-	if err != nil {
-		return ErrAccessDenied
-	}
-	if mysql.CompareNativePassword(clientAuthData, password, c.salt) {
-		return nil
+	for _, password := range credential.Passwords {
+		hash, err := credential.hashPassword(password)
+		if err != nil {
+			continue
+		}
+		decoded, err := mysql.DecodePasswordHex(hash)
+		if err != nil {
+			continue
+		}
+		if mysql.CompareNativePassword(clientAuthData, decoded, c.salt) {
+			return nil
+		}
 	}
 	return ErrAccessDenied
 }
 
 func (c *Conn) compareSha256PasswordAuthData(clientAuthData []byte, credential Credential) error {
 	// Empty passwords are not hashed, but sent as empty string
 	if len(clientAuthData) == 0 {
-		if credential.Password == "" {
+		if credential.hasEmptyPassword() {
 			return nil
 		}
 		return ErrAccessDeniedNoPassword
@@ -112,20 +118,26 @@ func (c *Conn) compareSha256PasswordAuthData(clientAuthData []byte, credential C
 			clientAuthData = clientAuthData[:l-1]
 		}
 	}
-	check, err := mysql.Check256HashingPassword([]byte(credential.Password), string(clientAuthData))
-	if err != nil {
-		return err
-	}
-	if check {
-		return nil
+	for _, password := range credential.Passwords {
+		hash, err := credential.hashPassword(password)
+		if err != nil {
+			continue
+		}
+		check, err := mysql.Check256HashingPassword([]byte(hash), string(clientAuthData))
+		if err != nil {
+			continue
+		}
+		if check {
+			return nil
+		}
 	}
 	return ErrAccessDenied
 }
 
 func (c *Conn) compareCacheSha2PasswordAuthData(clientAuthData []byte) error {
 	// Empty passwords are not hashed, but sent as empty string
 	if len(clientAuthData) == 0 {
-		if c.credential.Password == "" {
+		if c.credential.hasEmptyPassword() {
 			return nil
 		}
 		return ErrAccessDeniedNoPassword
@@ -139,10 +151,8 @@ func (c *Conn) compareCacheSha2PasswordAuthData(clientAuthData []byte) error {
 			// 'fast' auth: write "More data" packet (first byte == 0x01) with the second byte = 0x03
 			return c.writeAuthMoreDataFastAuth()
 		}
-
-		return ErrAccessDenied
 	}
-	// cache miss, do full auth
+	// cache miss or validation failed, do full auth
 	if err := c.writeAuthMoreDataFullAuth(); err != nil {
 		return err
 	}

server/auth_switch_response.go

@@ -24,7 +24,7 @@ func (c *Conn) handleAuthSwitchResponse() error {
 }
 
 func (c *Conn) handleCachingSha2PasswordFullAuth(authData []byte) error {
-	if err := c.acquirePassword(); err != nil {
+	if err := c.acquireCredential(); err != nil {
 		return err
 	}
 	if tlsConn, ok := c.Conn.Conn.(*tls.Conn); ok {
@@ -72,15 +72,21 @@ func (c *Conn) handleCachingSha2PasswordFullAuth(authData []byte) error {
 
 func (c *Conn) checkSha2CacheCredentials(clientAuthData []byte, credential Credential) error {
 	if len(clientAuthData) == 0 {
-		if credential.Password == "" {
+		if credential.hasEmptyPassword() {
 			return nil
 		}
 		return ErrAccessDeniedNoPassword
 	}
 
-	match, err := auth.CheckHashingPassword([]byte(credential.Password), string(clientAuthData), mysql.AUTH_CACHING_SHA2_PASSWORD)
-	if match && err == nil {
-		return nil
+	for _, password := range credential.Passwords {
+		hash, err := credential.hashPassword(password)
+		if err != nil {
+			continue
+		}
+		match, err := auth.CheckHashingPassword([]byte(hash), string(clientAuthData), mysql.AUTH_CACHING_SHA2_PASSWORD)
+		if match && err == nil {
+			return nil
+		}
 	}
 	return ErrAccessDenied
 }

server/auth_switch_response_test.go

@@ -30,7 +30,7 @@ func TestCheckSha2CacheCredentials_EmptyPassword(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Conn{
-				credential: Credential{Password: tt.serverPassword},
+				credential: Credential{Passwords: []string{tt.serverPassword}},
 			}
 			err := c.checkSha2CacheCredentials(tt.clientAuthData, c.credential)
 			if tt.wantErr == nil {

server/auth_test.go

@@ -37,7 +37,7 @@ func TestCompareNativePasswordAuthData_EmptyPassword(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Conn{
-				credential: Credential{Password: tt.serverPassword},
+				credential: Credential{Passwords: []string{tt.serverPassword}},
 			}
 			err := c.compareNativePasswordAuthData(tt.clientAuthData, c.credential)
 			if tt.wantErr == nil {
@@ -73,7 +73,7 @@ func TestCompareSha256PasswordAuthData_EmptyPassword(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Conn{
-				credential: Credential{Password: tt.serverPassword},
+				credential: Credential{Passwords: []string{tt.serverPassword}},
 			}
 			err := c.compareSha256PasswordAuthData(tt.clientAuthData, c.credential)
 			if tt.wantErr == nil {
@@ -109,7 +109,7 @@ func TestCompareCacheSha2PasswordAuthData_EmptyPassword(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Conn{
-				credential: Credential{Password: tt.serverPassword},
+				credential: Credential{Passwords: []string{tt.serverPassword}},
 			}
 			err := c.compareCacheSha2PasswordAuthData(tt.clientAuthData)
 			if tt.wantErr == nil {

server/authentication_handler.go

@@ -0,0 +1,123 @@
+package server
+
+import (
+	"slices"
+	"sync"
+
+	"github.com/go-mysql-org/go-mysql/mysql"
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/parser/auth"
+)
+
+// AuthenticationHandler provides user credentials and authentication lifecycle hooks.
+//
+// # Important Note
+//
+// if the password in a third-party auth handler could be updated at runtime, we have to invalidate the caching
+// for 'caching_sha2_password' by calling 'func (s *Server)InvalidateCache(string, string)'.
+type AuthenticationHandler interface {
+	// GetCredential returns the user credential (supports multiple valid passwords per user).
+	// Implementations must be safe for concurrent use.
+	GetCredential(username string) (credential Credential, found bool, err error)
+
+	// OnAuthSuccess is called after successful authentication, before the OK packet.
+	// Return an error to reject the connection (error will be sent to client instead of OK).
+	// Return nil to proceed with sending the OK packet.
+	OnAuthSuccess(conn *Conn) error
+
+	// OnAuthFailure is called after authentication fails, before the error packet.
+	// This is informational only - the connection will be closed regardless.
+	OnAuthFailure(conn *Conn, err error)
+}
+
+func NewInMemoryAuthenticationHandler(defaultAuthMethod ...string) *InMemoryAuthenticationHandler {
+	d := mysql.AUTH_CACHING_SHA2_PASSWORD
+	if len(defaultAuthMethod) > 0 {
+		d = defaultAuthMethod[0]
+	}
+	return &InMemoryAuthenticationHandler{
+		userPool:          sync.Map{},
+		defaultAuthMethod: d,
+	}
+}
+
+type Credential struct {
+	Passwords      []string // raw passwords, hashed on demand during comparison
+	AuthPluginName string
+}
+
+// hashPassword computes the password hash for a given password using the credential's auth plugin.
+func (c Credential) hashPassword(password string) (string, error) {
+	if password == "" {
+		return "", nil
+	}
+
+	switch c.AuthPluginName {
+	case mysql.AUTH_NATIVE_PASSWORD:
+		return mysql.EncodePasswordHex(mysql.NativePasswordHash([]byte(password))), nil
+
+	case mysql.AUTH_CACHING_SHA2_PASSWORD:
+		return auth.NewHashPassword(password, mysql.AUTH_CACHING_SHA2_PASSWORD), nil
+
+	case mysql.AUTH_SHA256_PASSWORD:
+		return mysql.NewSha256PasswordHash(password)
+
+	case mysql.AUTH_CLEAR_PASSWORD:
+		return password, nil
+
+	default:
+		return "", errors.Errorf("unknown authentication plugin name '%s'", c.AuthPluginName)
+	}
+}
+
+// hasEmptyPassword returns true if any password in the credential is empty.
+func (c Credential) hasEmptyPassword() bool {
+	return slices.Contains(c.Passwords, "")
+}
+
+// InMemoryAuthenticationHandler implements AuthenticationHandler with in-memory credential storage.
+type InMemoryAuthenticationHandler struct {
+	userPool          sync.Map // username -> Credential
+	defaultAuthMethod string
+}
+
+func (h *InMemoryAuthenticationHandler) CheckUsername(username string) (found bool, err error) {
+	_, ok := h.userPool.Load(username)
+	return ok, nil
+}
+
+func (h *InMemoryAuthenticationHandler) GetCredential(username string) (credential Credential, found bool, err error) {
+	v, ok := h.userPool.Load(username)
+	if !ok {
+		return Credential{}, false, nil
+	}
+	c, valid := v.(Credential)
+	if !valid {
+		return Credential{}, true, errors.Errorf("invalid credential")
+	}
+	return c, true, nil
+}
+
+func (h *InMemoryAuthenticationHandler) AddUser(username, password string, optionalAuthPluginName ...string) error {
+	authPluginName := h.defaultAuthMethod
+	if len(optionalAuthPluginName) > 0 {
+		authPluginName = optionalAuthPluginName[0]
+	}
+
+	if !isAuthMethodSupported(authPluginName) {
+		return errors.Errorf("unknown authentication plugin name '%s'", authPluginName)
+	}
+
+	h.userPool.Store(username, Credential{
+		Passwords:      []string{password},
+		AuthPluginName: authPluginName,
+	})
+	return nil
+}
+
+func (h *InMemoryAuthenticationHandler) OnAuthSuccess(conn *Conn) error {
+	return nil
+}
+
+func (h *InMemoryAuthenticationHandler) OnAuthFailure(conn *Conn, err error) {
+}