privileges: improve dynamic privs registration and tests

pingcap · May 19, 2021 · a343afd · a343afd
1 parent b546036
commit a343afd
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/privilege/privileges/privileges.go b/privilege/privileges/privileges.go
@@ -531,7 +531,8 @@ func (p *UserPrivileges) GetAllRoles(user, host string) []*auth.RoleIdentity {
 }
 
 // IsDynamicPrivilege returns true if the DYNAMIC privilege is built-in or has been registered by a plugin
-func (p *UserPrivileges) IsDynamicPrivilege(privNameInUpper string) bool {
+func (p *UserPrivileges) IsDynamicPrivilege(privName string) bool {
+	privNameInUpper := strings.ToUpper(privName)
 	for _, priv := range dynamicPrivs {
 		if privNameInUpper == priv {
 			return true
@@ -541,7 +542,11 @@ func (p *UserPrivileges) IsDynamicPrivilege(privNameInUpper string) bool {
 }
 
 // RegisterDynamicPrivilege is used by plugins to add new privileges to TiDB
-func RegisterDynamicPrivilege(privNameInUpper string) error {
+func RegisterDynamicPrivilege(privName string) error {
+	privNameInUpper := strings.ToUpper(privName)
+	if len(privNameInUpper) > 32 {
+		return errors.New("privilege name is longer than 32 characters")
+	}
 	dynamicPrivLock.Lock()
 	defer dynamicPrivLock.Unlock()
 	for _, priv := range dynamicPrivs {

diff --git a/privilege/privileges/privileges_test.go b/privilege/privileges/privileges_test.go
@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/store/mockstore"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/sem"
+	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/pingcap/tidb/util/testkit"
 	"github.com/pingcap/tidb/util/testleak"
 	"github.com/pingcap/tidb/util/testutil"
@@ -1579,7 +1580,33 @@ func (s *testPrivilegeSuite) TestDynamicPrivsRegistration(c *C) {
 	count := len(privileges.GetDynamicPrivileges())
 
 	c.Assert(pm.IsDynamicPrivilege("ACDC_ADMIN"), IsFalse)
-	privileges.RegisterDynamicPrivilege("ACDC_ADMIN")
+	c.Assert(privileges.RegisterDynamicPrivilege("ACDC_ADMIN"), IsNil)
 	c.Assert(pm.IsDynamicPrivilege("ACDC_ADMIN"), IsTrue)
 	c.Assert(len(privileges.GetDynamicPrivileges()), Equals, count+1)
+
+	c.Assert(pm.IsDynamicPrivilege("iAmdynamIC"), IsFalse)
+	c.Assert(privileges.RegisterDynamicPrivilege("IAMdynamic"), IsNil)
+	c.Assert(pm.IsDynamicPrivilege("IAMdyNAMIC"), IsTrue)
+	c.Assert(len(privileges.GetDynamicPrivileges()), Equals, count+2)
+
+	c.Assert(privileges.RegisterDynamicPrivilege("THIS_PRIVILEGE_NAME_IS_TOO_LONG_THE_MAX_IS_32_CHARS").Error(), Equals, "privilege name is longer than 32 characters")
+	c.Assert(pm.IsDynamicPrivilege("THIS_PRIVILEGE_NAME_IS_TOO_LONG_THE_MAX_IS_32_CHARS"), IsFalse)
+
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("CREATE USER privassigntest")
+	tk.MustExec("SET tidb_enable_dynamic_privileges=1")
+
+	// Check that all privileges registered are assignable to users,
+	// including the recently registered ACDC_ADMIN
+	for _, priv := range privileges.GetDynamicPrivileges() {
+		sqlGrant, err := sqlexec.EscapeSQL("GRANT %n ON *.* TO privassigntest", priv)
+		c.Assert(err, IsNil)
+		tk.MustExec(sqlGrant)
+	}
+	// Check that all privileges registered are revokable
+	for _, priv := range privileges.GetDynamicPrivileges() {
+		sqlGrant, err := sqlexec.EscapeSQL("REVOKE %n ON *.* FROM privassigntest", priv)
+		c.Assert(err, IsNil)
+		tk.MustExec(sqlGrant)
+	}
 }