|
| 1 | +package server |
| 2 | + |
| 3 | +import ( |
| 4 | + "slices" |
| 5 | + "sync" |
| 6 | + |
| 7 | + "github.com/go-mysql-org/go-mysql/mysql" |
| 8 | + "github.com/pingcap/errors" |
| 9 | + "github.com/pingcap/tidb/pkg/parser/auth" |
| 10 | +) |
| 11 | + |
| 12 | +// AuthenticationHandler provides user credentials and authentication lifecycle hooks. |
| 13 | +// |
| 14 | +// # Important Note |
| 15 | +// |
| 16 | +// if the password in a third-party auth handler could be updated at runtime, we have to invalidate the caching |
| 17 | +// for 'caching_sha2_password' by calling 'func (s *Server)InvalidateCache(string, string)'. |
| 18 | +type AuthenticationHandler interface { |
| 19 | + // GetCredential returns the user credential (supports multiple valid passwords per user). |
| 20 | + // Implementations must be safe for concurrent use. |
| 21 | + GetCredential(username string) (credential Credential, found bool, err error) |
| 22 | + |
| 23 | + // OnAuthSuccess is called after successful authentication, before the OK packet. |
| 24 | + // Return an error to reject the connection (error will be sent to client instead of OK). |
| 25 | + // Return nil to proceed with sending the OK packet. |
| 26 | + OnAuthSuccess(conn *Conn) error |
| 27 | + |
| 28 | + // OnAuthFailure is called after authentication fails, before the error packet. |
| 29 | + // This is informational only - the connection will be closed regardless. |
| 30 | + OnAuthFailure(conn *Conn, err error) |
| 31 | +} |
| 32 | + |
| 33 | +func NewInMemoryAuthenticationHandler(defaultAuthMethod ...string) *InMemoryAuthenticationHandler { |
| 34 | + d := mysql.AUTH_CACHING_SHA2_PASSWORD |
| 35 | + if len(defaultAuthMethod) > 0 { |
| 36 | + d = defaultAuthMethod[0] |
| 37 | + } |
| 38 | + return &InMemoryAuthenticationHandler{ |
| 39 | + userPool: sync.Map{}, |
| 40 | + defaultAuthMethod: d, |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +// Credential holds authentication settings for a user. |
| 45 | +// Passwords contains all valid raw passwords for the user. They are hashed on demand during comparison. |
| 46 | +// If empty password authentication is allowed, Passwords must contain an empty string (e.g., []string{""}) |
| 47 | +// rather than being a zero-length slice. A zero-length slice means no valid passwords are configured. |
| 48 | +type Credential struct { |
| 49 | + Passwords []string |
| 50 | + AuthPluginName string |
| 51 | +} |
| 52 | + |
| 53 | +// hashPassword computes the password hash for a given password using the credential's auth plugin. |
| 54 | +func (c Credential) hashPassword(password string) (string, error) { |
| 55 | + if password == "" { |
| 56 | + return "", nil |
| 57 | + } |
| 58 | + |
| 59 | + switch c.AuthPluginName { |
| 60 | + case mysql.AUTH_NATIVE_PASSWORD: |
| 61 | + return mysql.EncodePasswordHex(mysql.NativePasswordHash([]byte(password))), nil |
| 62 | + |
| 63 | + case mysql.AUTH_CACHING_SHA2_PASSWORD: |
| 64 | + return auth.NewHashPassword(password, mysql.AUTH_CACHING_SHA2_PASSWORD), nil |
| 65 | + |
| 66 | + case mysql.AUTH_SHA256_PASSWORD: |
| 67 | + return mysql.NewSha256PasswordHash(password) |
| 68 | + |
| 69 | + case mysql.AUTH_CLEAR_PASSWORD: |
| 70 | + return password, nil |
| 71 | + |
| 72 | + default: |
| 73 | + return "", errors.Errorf("unknown authentication plugin name '%s'", c.AuthPluginName) |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +// hasEmptyPassword returns true if any password in the credential is empty. |
| 78 | +func (c Credential) hasEmptyPassword() bool { |
| 79 | + return slices.Contains(c.Passwords, "") |
| 80 | +} |
| 81 | + |
| 82 | +// InMemoryAuthenticationHandler implements AuthenticationHandler with in-memory credential storage. |
| 83 | +type InMemoryAuthenticationHandler struct { |
| 84 | + userPool sync.Map // username -> Credential |
| 85 | + defaultAuthMethod string |
| 86 | +} |
| 87 | + |
| 88 | +func (h *InMemoryAuthenticationHandler) CheckUsername(username string) (found bool, err error) { |
| 89 | + _, ok := h.userPool.Load(username) |
| 90 | + return ok, nil |
| 91 | +} |
| 92 | + |
| 93 | +func (h *InMemoryAuthenticationHandler) GetCredential(username string) (credential Credential, found bool, err error) { |
| 94 | + v, ok := h.userPool.Load(username) |
| 95 | + if !ok { |
| 96 | + return Credential{}, false, nil |
| 97 | + } |
| 98 | + c, valid := v.(Credential) |
| 99 | + if !valid { |
| 100 | + return Credential{}, true, errors.Errorf("invalid credential") |
| 101 | + } |
| 102 | + return c, true, nil |
| 103 | +} |
| 104 | + |
| 105 | +func (h *InMemoryAuthenticationHandler) AddUser(username, password string, optionalAuthPluginName ...string) error { |
| 106 | + authPluginName := h.defaultAuthMethod |
| 107 | + if len(optionalAuthPluginName) > 0 { |
| 108 | + authPluginName = optionalAuthPluginName[0] |
| 109 | + } |
| 110 | + |
| 111 | + if !isAuthMethodSupported(authPluginName) { |
| 112 | + return errors.Errorf("unknown authentication plugin name '%s'", authPluginName) |
| 113 | + } |
| 114 | + |
| 115 | + h.userPool.Store(username, Credential{ |
| 116 | + Passwords: []string{password}, |
| 117 | + AuthPluginName: authPluginName, |
| 118 | + }) |
| 119 | + return nil |
| 120 | +} |
| 121 | + |
| 122 | +func (h *InMemoryAuthenticationHandler) OnAuthSuccess(conn *Conn) error { |
| 123 | + return nil |
| 124 | +} |
| 125 | + |
| 126 | +func (h *InMemoryAuthenticationHandler) OnAuthFailure(conn *Conn, err error) { |
| 127 | +} |
0 commit comments