server: restore session on change user auth failure (release-7.5-20260724-v7.5.7) (#69692)

Project: pingcap/tidb | Language: Go | Concepts: error-handling, refactoring | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70171

The problem

When a client sends a COM_CHANGE_USER command to change the current session's user or database, if authentication fails (wrong credentials, missing user, plugin error, etc.), the connection could be left in a partially updated and inconsistent state. This means the session might have the new (but unauthenticated) user or database, or a half-initialized context, which can cause subsequent errors or security issues.

Why it's non-trivial

A naive fix might just abort on error, but by that point, some changes (like updating the user, database, or context) may have already been applied to the connection. Simply returning early does not revert these partial changes, so the connection can remain in a corrupted state. Properly restoring the previous session state—including user, database, authentication plugin, and session context—requires careful snapshotting and rollback logic.

The approach

This PR snapshots the old session state before attempting to change the user, and restores it if any part of the change-user/authentication flow fails.

Snapshot old state:

+	oldUser, oldDBName, oldAuthPlugin := cc.user, cc.dbname, cc.authPlugin
+	oldCtx := cc.getCtx()

Before making any changes, the code saves the current user, database name, authentication plugin, and session context. This is necessary to restore them if something goes wrong.

Parse new user/db, but do not apply yet:

-	cc.user = string(hack.String(user))
+	newUser := string(hack.String(user))
...
-	cc.dbname = string(hack.String(dbName))
+	newDBName := string(hack.String(dbName))

The new user and database name are parsed from the packet, but not immediately assigned to the connection.

Apply new state, but with a plan to restore:

+	cc.user = newUser
+	cc.dbname = newDBName
err := cc.openSession()
+	restoreOldSession := func() {
+		if newCtx := cc.getCtx(); newCtx != nil && newCtx != oldCtx {
+			if err := newCtx.Close(); err != nil {
+				logutil.Logger(ctx).Debug("close new context failed", zap.Error(err))
+			}
+		}
+		cc.SetCtx(oldCtx)
+		cc.user = oldUser
+		cc.dbname = oldDBName
+		cc.authPlugin = oldAuthPlugin
+	}
if err != nil {
+	restoreOldSession()
	return err
}

The new user and database are applied, and a new session is opened. If opening the session fails, the restoreOldSession function is called to revert all changes.

Restore on plugin or authentication failure:

if fakeResp.AuthPlugin != "" {
	failpoint.Inject("ChangeUserAuthSwitch", func(val failpoint.Value) {
+		restoreOldSession()
		failpoint.Return(errors.Errorf("%v", val))
	})
	newpass, err := cc.checkAuthPlugin(ctx, fakeResp)
	if err != nil {
+		restoreOldSession()
		return err
	}
	...
}
if err := cc.openSessionAndDoAuth(fakeResp.Auth, fakeResp.AuthPlugin); err != nil {
+	restoreOldSession()
	return err
}

If plugin switching or authentication fails, restoreOldSession is called before returning the error.

Only close the old context after successful authentication:

+	if oldCtx != nil {
+		if err := oldCtx.Close(); err != nil {
+			logutil.Logger(ctx).Debug("close old context failed", zap.Error(err))
+		}
+	}

The old session context is only closed after the new session is fully authenticated and established.

Test: verify restoration on failure

func TestChangeUserAuthFailureRestoresOldSession(t *testing.T) {
    ...
    err = cc.handleChangeUser(context.Background(), data)
    require.Error(t, err)
    require.Same(t, tc, cc.getCtx())
    require.Equal(t, "root", cc.user)
    require.Equal(t, "old_db", cc.dbname)
    require.Equal(t, "root", cc.ctx.GetSessionVars().User.Username)
}

This test triggers a change-user with a non-existent user, expects an error, and checks that the original session context, user, and database are all unchanged.

From the review thread

No notable review discussion on this PR.

Takeaways

• Always snapshot all relevant state before making multi-step changes that could fail partway, so you can safely roll back.
• Error handling for session or authentication changes must account for partially-applied changes; restoring the previous state is safer than leaving the connection in a potentially inconsistent state.
• Unit tests should verify not just error returns, but also that the system's state is unchanged after a failed operation.