session, variable: add privilege check to set session_states (#59601)

close #59600

Author: djshow832

pkg/session/session.go

@@ -4395,6 +4395,7 @@ func (s *session) EncodeSessionStates(ctx context.Context,
 	if err := s.sessionVars.EncodeSessionStates(ctx, sessionStates); err != nil {
 		return err
 	}
+	sessionStates.ResourceGroupName = s.sessionVars.ResourceGroupName
 
 	hasRestrictVarPriv := false
 	checker := privilege.GetPrivilegeManager(s)
@@ -4474,6 +4475,25 @@ func (s *session) DecodeSessionStates(ctx context.Context,
 		}
 	}
 
+	// Put resource group privilege check from sessionVars to session to avoid circular dependency.
+	if sessionStates.ResourceGroupName != s.sessionVars.ResourceGroupName {
+		hasPriv := true
+		if vardef.EnableResourceControlStrictMode.Load() {
+			checker := privilege.GetPrivilegeManager(s)
+			if checker != nil {
+				hasRgAdminPriv := checker.RequestDynamicVerification(s.sessionVars.ActiveRoles, "RESOURCE_GROUP_ADMIN", false)
+				hasRgUserPriv := checker.RequestDynamicVerification(s.sessionVars.ActiveRoles, "RESOURCE_GROUP_USER", false)
+				hasPriv = hasRgAdminPriv || hasRgUserPriv
+			}
+		}
+		if hasPriv {
+			s.sessionVars.SetResourceGroupName(sessionStates.ResourceGroupName)
+		} else {
+			logutil.Logger(ctx).Warn("set session states error, no privilege to set resource group, skip changing resource group",
+				zap.String("source_resource_group", s.sessionVars.ResourceGroupName), zap.String("target_resource_group", sessionStates.ResourceGroupName))
+		}
+	}
+
 	// Decoding session vars / prepared statements may override stmt ctx, such as warnings,
 	// so we decode stmt ctx at last.
 	return s.sessionVars.DecodeSessionStates(ctx, sessionStates)

pkg/sessionctx/variable/session.go

@@ -2826,7 +2826,6 @@ func (s *SessionVars) EncodeSessionStates(_ context.Context, sessionStates *sess
 	sessionStates.SequenceLatestValues = s.SequenceState.GetAllStates()
 	sessionStates.FoundInPlanCache = s.PrevFoundInPlanCache
 	sessionStates.FoundInBinding = s.PrevFoundInBinding
-	sessionStates.ResourceGroupName = s.ResourceGroupName
 	sessionStates.HypoIndexes = s.HypoIndexes
 	sessionStates.HypoTiFlashReplicas = s.HypoTiFlashReplicas
 
@@ -2862,7 +2861,6 @@ func (s *SessionVars) DecodeSessionStates(_ context.Context, sessionStates *sess
 	s.SequenceState.SetAllStates(sessionStates.SequenceLatestValues)
 	s.FoundInPlanCache = sessionStates.FoundInPlanCache
 	s.FoundInBinding = sessionStates.FoundInBinding
-	s.SetResourceGroupName(sessionStates.ResourceGroupName)
 	s.HypoIndexes = sessionStates.HypoIndexes
 	s.HypoTiFlashReplicas = sessionStates.HypoTiFlashReplicas
 

tests/integrationtest/r/privilege/privileges.result

@@ -101,6 +101,30 @@ id
 REVOKE SELECT on test_rc.* FROM resource_group_admin;
 REVOKE SELECT on test_rc.* FROM resource_group_user;
 DROP DATABASE test_rc;
+DROP USER resource_group_admin;
+DROP USER resource_group_user;
+DROP RESOURCE GROUP test;
+CREATE USER resource_group_user;
+CREATE USER no_resource_group;
+CREATE RESOURCE GROUP test RU_PER_SEC = 666;
+GRANT RESOURCE_GROUP_USER ON *.* TO resource_group_user;
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+CURRENT_RESOURCE_GROUP()
+default
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+CURRENT_RESOURCE_GROUP()
+test
+set @@global.tidb_resource_control_strict_mode = 0;
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+CURRENT_RESOURCE_GROUP()
+test
+set @@global.tidb_resource_control_strict_mode = default;
+DROP RESOURCE GROUP test;
+DROP USER resource_group_user;
+DROP USER no_resource_group;
 CREATE SCHEMA IF NOT EXISTS privilege__privileges;
 USE privilege__privileges;
 CREATE TABLE reftest (a int);

tests/integrationtest/t/privilege/privileges.test

@@ -139,6 +139,40 @@ connection default;
 REVOKE SELECT on test_rc.* FROM resource_group_admin;
 REVOKE SELECT on test_rc.* FROM resource_group_user;
 DROP DATABASE test_rc;
+DROP USER resource_group_admin;
+DROP USER resource_group_user;
+DROP RESOURCE GROUP test;
+
+# TestSetSessionStatesPriv
+CREATE USER resource_group_user;
+CREATE USER no_resource_group;
+
+connection default;
+CREATE RESOURCE GROUP test RU_PER_SEC = 666;
+GRANT RESOURCE_GROUP_USER ON *.* TO resource_group_user;
+
+connect (no_resource_group,localhost,no_resource_group,,);
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+
+connect (resource_group_user,localhost,resource_group_user,,);
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+
+connection default;
+set @@global.tidb_resource_control_strict_mode = 0;
+
+connection no_resource_group;
+SET SESSION_STATES '{"rs-group":"test"}';
+SELECT CURRENT_RESOURCE_GROUP();
+
+disconnect resource_group_user;
+disconnect no_resource_group;
+connection default;
+set @@global.tidb_resource_control_strict_mode = default;
+DROP RESOURCE GROUP test;
+DROP USER resource_group_user;
+DROP USER no_resource_group;
 
 # TestGrantReferences
 CREATE SCHEMA IF NOT EXISTS privilege__privileges;