fix the issue that some privileges cannot be limited

Signed-off-by: Yang Keao <yangkeao@chunibyo.icu>

Author: YangKeao

pkg/planner/core/planbuilder.go

@@ -3963,7 +3963,7 @@ func collectVisitInfoFromGrantStmt(sctx base.PlanContext, vi []visitInfo, stmt *
 	for _, item := range stmt.Privs {
 		if semv2.IsEnabled() {
 			if (len(item.Name) > 0 && semv2.IsRestrictedPrivilege(item.Name)) ||
-				(len(item.Name) == 0 && semv2.IsRestrictedPrivilege(item.Priv.String())) {
+				(len(item.Name) == 0 && semv2.IsRestrictedPrivilege(strings.ToUpper(item.Priv.String()))) {
 				// In `semv2`, we'll support to limit non-dynamic privileges unless the user has the `RESTRICTED_PRIV_ADMIN` privilege.
 				// For example, `File` privilege might be restricted.
 				// It's also controlled by the `GRANT OPTION`, so the user will also need the `GRANT OPTION` for this privilege.

pkg/privilege/privileges/BUILD.bazel

@@ -83,6 +83,7 @@ go_test(
         "//pkg/util/hack",
         "//pkg/util/sem",
         "//pkg/util/sem/compat",
+        "//pkg/util/sem/v2:sem",
         "//pkg/util/sqlescape",
         "@com_github_lestrrat_go_jwx_v2//jwa",
         "@com_github_lestrrat_go_jwx_v2//jwk",

pkg/privilege/privileges/privileges_test.go

@@ -50,6 +50,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
 	semv1 "github.com/pingcap/tidb/pkg/util/sem"
 	sem "github.com/pingcap/tidb/pkg/util/sem/compat"
+	semv2 "github.com/pingcap/tidb/pkg/util/sem/v2"
 	"github.com/pingcap/tidb/pkg/util/sqlescape"
 	"github.com/stretchr/testify/require"
 )
@@ -2230,16 +2231,19 @@ func TestGrantOptionWithSEMv2(t *testing.T) {
 	rootTk.MustExec("CREATE USER varuser2")
 	rootTk.MustExec("CREATE USER varuser3")
 	rootTk.MustExec("CREATE USER varuser4")
+	rootTk.MustExec("CREATE USER varuser5")
 	rootTk.MustExec("CREATE USER grantee")
 
 	rootTk.MustExec("GRANT SYSTEM_VARIABLES_ADMIN, FILE ON *.* TO varuser1")
 	rootTk.MustExec("GRANT SYSTEM_VARIABLES_ADMIN, FILE ON *.* TO varuser2 WITH GRANT OPTION")
 	rootTk.MustExec("GRANT RESTRICTED_PRIV_ADMIN ON *.* TO varuser3")
 	rootTk.MustExec("GRANT RESTRICTED_PRIV_ADMIN ON *.* TO varuser4")
 	rootTk.MustExec("GRANT SYSTEM_VARIABLES_ADMIN, FILE ON *.* TO varuser4 WITH GRANT OPTION")
+	rootTk.MustExec("GRANT SYSTEM_VARIABLES_ADMIN, DROP ON *.* TO varuser5 WITH GRANT OPTION")
 
-	// SYSTEM_VARIABLES_ADMIN is not restricted, FILE is restricted.
+	// SYSTEM_VARIABLES_ADMIN is not restricted, FILE and Drop are restricted.
 	defer sem.SwitchToSEMForTest(t, sem.V2)()
+	semv2.AddRestrictedPrivilegesForTest("Drop")
 	// try to grant SYSTEM_VARIABLES_ADMIN and FILE privilege to grantee with different user
 	tk1 := testkit.NewTestKit(t, store)
 	require.NoError(t, tk1.Session().Auth(&auth.UserIdentity{Username: "varuser1", Hostname: "%"}, nil, nil, nil))
@@ -2268,4 +2272,10 @@ func TestGrantOptionWithSEMv2(t *testing.T) {
 	require.NoError(t, err)
 	err = tk4.ExecToErr("GRANT FILE ON *.* TO grantee")
 	require.NoError(t, err)
+
+	// Test grant drop
+	tk5 := testkit.NewTestKit(t, store)
+	require.NoError(t, tk5.Session().Auth(&auth.UserIdentity{Username: "varuser5", Hostname: "%"}, nil, nil, nil))
+	err = tk5.ExecToErr("GRANT drop ON *.* TO grantee")
+	require.EqualError(t, err, "[planner:1227]Access denied; you need (at least one of) the RESTRICTED_PRIV_ADMIN privilege(s) for this operation")
 }

pkg/util/sem/v2/testhelper.go

@@ -14,7 +14,11 @@
 
 package sem
 
-import "github.com/pingcap/tidb/pkg/sessionctx/variable"
+import (
+	"strings"
+
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+)
 
 // EnableFromPathForTest enables SEM v2 in test using a configuration file.
 func EnableFromPathForTest(configPath string) (func(), error) {
@@ -48,3 +52,12 @@ func EnableFromPathForTest(configPath string) (func(), error) {
 		}
 	}, nil
 }
+
+// AddRestrictedPrivilegesForTest adds restricted privileges for test.
+func AddRestrictedPrivilegesForTest(privilege string) {
+	if sem == nil {
+		return
+	}
+
+	sem.restrictedPrivileges[strings.ToUpper(privilege)] = struct{}{}
+}