diff --git a/pkg/session/session.go b/pkg/session/session.go index 5e1f01158680c..05d7dc7787ed1 100644 --- a/pkg/session/session.go +++ b/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) diff --git a/pkg/sessionctx/variable/session.go b/pkg/sessionctx/variable/session.go index d6500986bc4cb..8d685273ff9b5 100644 --- a/pkg/sessionctx/variable/session.go +++ b/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 diff --git a/tests/integrationtest/r/privilege/privileges.result b/tests/integrationtest/r/privilege/privileges.result index c6025a020cc39..4ca52456dd12a 100644 --- a/tests/integrationtest/r/privilege/privileges.result +++ b/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); diff --git a/tests/integrationtest/t/privilege/privileges.test b/tests/integrationtest/t/privilege/privileges.test index b0d6f67f606ea..d7bd1b15e6348 100644 --- a/tests/integrationtest/t/privilege/privileges.test +++ b/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;