extension: add support for extension to custom table and sysvar privileges

executor/set.go

@@ -24,14 +24,17 @@ import (
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/charset"
 	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/plugin"
+	"github.com/pingcap/tidb/privilege"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/table/temptable"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/gcutil"
 	"github.com/pingcap/tidb/util/logutil"
+	"github.com/pingcap/tidb/util/sem"
 	"go.uber.org/zap"
 )
 
@@ -109,6 +112,22 @@ func (e *SetExecutor) setSysVariable(ctx context.Context, name string, v *expres
 		}
 		return variable.ErrUnknownSystemVar.GenWithStackByArgs(name)
 	}
+
+	if sysVar.RequireDynamicPrivileges != nil {
+		semEnabled := sem.IsEnabled()
+		pm := privilege.GetPrivilegeManager(e.ctx)
+		privs := sysVar.RequireDynamicPrivileges(v.IsGlobal, semEnabled)
+		for _, priv := range privs {
+			if !pm.RequestDynamicVerification(sessionVars.ActiveRoles, priv, false) {
+				msg := priv
+				if !semEnabled {
+					msg = "SUPER or " + msg
+				}
+				return core.ErrSpecificAccessDenied.GenWithStackByArgs(msg)
+			}
+		}
+	}
+
 	if sysVar.IsNoop && !variable.EnableNoopVariables.Load() {
 		// The variable is a noop. For compatibility we allow it to still
 		// be changed, but we append a warning since users might be expecting

expression/extension.go

@@ -114,11 +114,13 @@ func (c *extensionFuncClass) getFunction(ctx sessionctx.Context, args []Expressi
 }
 
 func (c *extensionFuncClass) checkPrivileges(ctx sessionctx.Context) error {
-	privs := c.funcDef.RequireDynamicPrivileges
-	if semPrivs := c.funcDef.SemRequireDynamicPrivileges; len(semPrivs) > 0 && sem.IsEnabled() {
-		privs = semPrivs
+	fn := c.funcDef.RequireDynamicPrivileges
+	if fn == nil {
+		return nil
 	}
 
+	semEnabled := sem.IsEnabled()
+	privs := fn(semEnabled)
 	if len(privs) == 0 {
 		return nil
 	}
@@ -129,7 +131,7 @@ func (c *extensionFuncClass) checkPrivileges(ctx sessionctx.Context) error {
 	for _, priv := range privs {
 		if !manager.RequestDynamicVerification(activeRoles, priv, false) {
 			msg := priv
-			if !sem.IsEnabled() {
+			if !semEnabled {
 				msg = "SUPER or " + msg
 			}
 			return errSpecificAccessDenied.GenWithStackByArgs(msg)

extension/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//parser",
         "//parser/ast",
         "//parser/auth",
+        "//parser/mysql",
         "//sessionctx/stmtctx",
         "//sessionctx/variable",
         "//types",

extension/extensions.go

@@ -45,6 +45,21 @@ func (es *Extensions) Bootstrap(ctx BootstrapContext) error {
 	return nil
 }
 
+// GetAccessCheckFuncs returns spec functions of the custom access check
+func (es *Extensions) GetAccessCheckFuncs() (funcs []AccessCheckFunc) {
+	if es == nil {
+		return nil
+	}
+
+	for _, m := range es.manifests {
+		if m.accessCheckFunc != nil {
+			funcs = append(funcs, m.accessCheckFunc)
+		}
+	}
+
+	return funcs
+}
+
 // NewSessionExtensions creates a new ConnExtensions object
 func (es *Extensions) NewSessionExtensions() *SessionExtensions {
 	if es == nil {

extension/function.go

@@ -48,12 +48,8 @@ type FunctionDef struct {
 	EvalStringFunc func(ctx FunctionContext, row chunk.Row) (string, bool, error)
 	// EvalIntFunc is the eval function when `EvalTp` is `types.ETInt`
 	EvalIntFunc func(ctx FunctionContext, row chunk.Row) (int64, bool, error)
-	// RequireDynamicPrivileges is the dynamic privileges needed to invoke the function
-	// If `RequireDynamicPrivileges` is empty, it means every one can invoke this function
-	RequireDynamicPrivileges []string
-	// SemRequireDynamicPrivileges is the dynamic privileges needed to invoke the function in sem mode
-	// If `SemRequireDynamicPrivileges` is empty, `DynamicPrivileges` will be used in sem mode
-	SemRequireDynamicPrivileges []string
+	// RequireDynamicPrivileges is a function to return a list of dynamic privileges to check.
+	RequireDynamicPrivileges func(sem bool) []string
 }
 
 // Validate validates the function definition

extension/function_test.go

@@ -283,26 +283,37 @@ func TestExtensionFuncPrivilege(t *testing.T) {
 				},
 			},
 			{
-				Name:                     "custom_only_dyn_priv_func",
-				EvalTp:                   types.ETString,
-				RequireDynamicPrivileges: []string{"CUSTOM_DYN_PRIV_1"},
+				Name:   "custom_only_dyn_priv_func",
+				EvalTp: types.ETString,
+				RequireDynamicPrivileges: func(sem bool) []string {
+					return []string{"CUSTOM_DYN_PRIV_1"}
+				},
 				EvalStringFunc: func(ctx extension.FunctionContext, row chunk.Row) (string, bool, error) {
 					return "abc", false, nil
 				},
 			},
 			{
-				Name:                        "custom_only_sem_dyn_priv_func",
-				EvalTp:                      types.ETString,
-				SemRequireDynamicPrivileges: []string{"RESTRICTED_CUSTOM_DYN_PRIV_2"},
+				Name:   "custom_only_sem_dyn_priv_func",
+				EvalTp: types.ETString,
+				RequireDynamicPrivileges: func(sem bool) []string {
+					if sem {
+						return []string{"RESTRICTED_CUSTOM_DYN_PRIV_2"}
+					}
+					return nil
+				},
 				EvalStringFunc: func(ctx extension.FunctionContext, row chunk.Row) (string, bool, error) {
 					return "def", false, nil
 				},
 			},
 			{
-				Name:                        "custom_both_dyn_priv_func",
-				EvalTp:                      types.ETString,
-				RequireDynamicPrivileges:    []string{"CUSTOM_DYN_PRIV_1"},
-				SemRequireDynamicPrivileges: []string{"RESTRICTED_CUSTOM_DYN_PRIV_2"},
+				Name:   "custom_both_dyn_priv_func",
+				EvalTp: types.ETString,
+				RequireDynamicPrivileges: func(sem bool) []string {
+					if sem {
+						return []string{"RESTRICTED_CUSTOM_DYN_PRIV_2"}
+					}
+					return []string{"CUSTOM_DYN_PRIV_1"}
+				},
 				EvalStringFunc: func(ctx extension.FunctionContext, row chunk.Row) (string, bool, error) {
 					return "ghi", false, nil
 				},

extension/manifest.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/ngaut/pools"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/chunk"
 	clientv3 "go.etcd.io/etcd/client/v3"
@@ -54,6 +55,16 @@ func WithCustomFunctions(funcs []*FunctionDef) Option {
 	}
 }
 
+// AccessCheckFunc is a function that returns a dynamic privilege list for db/tbl/column access
+type AccessCheckFunc func(db, tbl, column string, priv mysql.PrivilegeType, sem bool) []string
+
+// WithCustomAccessCheck specifies the custom db/tbl/column dynamic privilege check
+func WithCustomAccessCheck(fn AccessCheckFunc) Option {
+	return func(m *Manifest) {
+		m.accessCheckFunc = fn
+	}
+}
+
 // WithSessionHandlerFactory specifies a factory function to handle session
 func WithSessionHandlerFactory(factory func() *SessionHandler) Option {
 	return func(m *Manifest) {
@@ -106,6 +117,7 @@ type Manifest struct {
 	dynPrivs              []string
 	bootstrap             func(BootstrapContext) error
 	funcs                 []*FunctionDef
+	accessCheckFunc       AccessCheckFunc
 	sessionHandlerFactory func() *SessionHandler
 	close                 func()
 }