enhance: support user-defined builtin roles

add new configuration: builtinRoles
user can define roles in config file: milvus.yaml

there is an example:
1. db_ro, only have read privileges, include load
2. db_rw, read and write privileges, include create/drop/rename collection
3. db_admin, not only read and write privileges, but also user administration

Signed-off-by: PowderLi <[email protected]>

internal/proxy/impl.go

@@ -4225,7 +4225,7 @@ func (node *Proxy) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest)
 	if err := ValidateRoleName(req.RoleName); err != nil {
 		return merr.Status(err), nil
 	}
-	if IsDefaultRole(req.RoleName) {
+	if IsBuiltinRole(req.RoleName) {
 		err := merr.WrapErrPrivilegeNotPermitted("the role[%s] is a default role, which can't be droped", req.GetRoleName())
 		return merr.Status(err), nil
 	}

internal/proxy/util.go

@@ -835,6 +835,18 @@ func IsDefaultRole(roleName string) bool {
 	return false
 }
 
+func IsBuiltinRole(roleName string) bool {
+	if IsDefaultRole(roleName) {
+		return true
+	}
+	for _, builtinRole := range util.BuiltinRoles {
+		if builtinRole == roleName {
+			return true
+		}
+	}
+	return false
+}
+
 func ValidateObjectName(entity string) error {
 	if util.IsAnyWord(entity) {
 		return nil

internal/proxy/util_test.go

@@ -816,8 +816,13 @@ func TestValidateName(t *testing.T) {
 
 func TestIsDefaultRole(t *testing.T) {
 	assert.Equal(t, true, IsDefaultRole(util.RoleAdmin))
+	assert.Equal(t, true, IsBuiltinRole(util.RoleAdmin))
 	assert.Equal(t, true, IsDefaultRole(util.RolePublic))
+	assert.Equal(t, true, IsBuiltinRole(util.RolePublic))
 	assert.Equal(t, false, IsDefaultRole("manager"))
+	assert.Equal(t, false, IsBuiltinRole("manager"))
+	util.BuiltinRoles = append(util.BuiltinRoles, "manager")
+	assert.Equal(t, true, IsBuiltinRole("manager"))
 }
 
 func GetContext(ctx context.Context, originValue string) context.Context {

internal/rootcoord/root_coord.go

@@ -613,6 +613,43 @@ func (c *Core) initRbac() error {
 			return errors.Wrap(err, "failed to grant collection privilege")
 		}
 	}
+	if Params.RoleCfg.Enabled.GetAsBool() {
+		return c.initBuiltinRoles()
+	}
+	return nil
+}
+
+func (c *Core) initBuiltinRoles() error {
+	rolePrivilegesMap := Params.RoleCfg.Roles.GetAsRoleDetails()
+	for role, privilegesJSON := range rolePrivilegesMap {
+		err := c.meta.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
+		if err != nil && !common.IsIgnorableError(err) {
+			log.Error("create a builtin role fail", zap.Any("error", err), zap.String("roleName", role))
+			return errors.Wrapf(err, "failed to create a builtin role: %s", role)
+		}
+		for _, privilege := range privilegesJSON[util.RoleConfigPrivileges] {
+			privilegeName := privilege[util.RoleConfigPrivilege]
+			if !util.IsAnyWord(privilege[util.RoleConfigPrivilege]) {
+				privilegeName = util.PrivilegeNameForMetastore(privilege[util.RoleConfigPrivilege])
+			}
+			err := c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
+				Role:       &milvuspb.RoleEntity{Name: role},
+				Object:     &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]},
+				ObjectName: privilege[util.RoleConfigObjectName],
+				DbName:     privilege[util.RoleConfigDBName],
+				Grantor: &milvuspb.GrantorEntity{
+					User:      &milvuspb.UserEntity{Name: util.UserRoot},
+					Privilege: &milvuspb.PrivilegeEntity{Name: privilegeName},
+				},
+			}, milvuspb.OperatePrivilegeType_Grant)
+			if err != nil && !common.IsIgnorableError(err) {
+				log.Error("grant privilege to builtin role fail", zap.Any("error", err), zap.String("roleName", role), zap.Any("privilege", privilege))
+				return errors.Wrapf(err, "failed to grant privilege: <%s, %s, %s> of db: %s to role: %s", privilege[util.RoleConfigObjectType], privilege[util.RoleConfigObjectName], privilege[util.RoleConfigPrivilege], privilege[util.RoleConfigDBName], role)
+			}
+		}
+		util.BuiltinRoles = append(util.BuiltinRoles, role)
+		log.Info("init a builtin role successfully", zap.String("roleName", role))
+	}
 	return nil
 }
 

internal/rootcoord/root_coord_test.go

@@ -2068,6 +2068,45 @@ func TestRootCoord_RBACError(t *testing.T) {
 	})
 }
 
+func TestRootCoord_BuiltinRoles(t *testing.T) {
+	paramtable.Init()
+	paramtable.Get().Save(paramtable.Get().RoleCfg.Enabled.Key, "true")
+	paramtable.Get().Save(paramtable.Get().RoleCfg.Roles.Key, `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}}`)
+	t.Run("init builtin roles success", func(t *testing.T) {
+		c := newTestCore(withHealthyCode(), withInvalidMeta())
+		mockMeta := c.meta.(*mockMetaTable)
+		mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
+			return nil
+		}
+		mockMeta.OperatePrivilegeFunc = func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
+			return nil
+		}
+		err := c.initBuiltinRoles()
+		assert.Equal(t, nil, err)
+	})
+	t.Run("init builtin roles fail to create role", func(t *testing.T) {
+		c := newTestCore(withHealthyCode(), withInvalidMeta())
+		mockMeta := c.meta.(*mockMetaTable)
+		mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
+			return merr.ErrPrivilegeNotPermitted
+		}
+		err := c.initBuiltinRoles()
+		assert.Error(t, err)
+	})
+	t.Run("init builtin roles fail to operate privileg", func(t *testing.T) {
+		c := newTestCore(withHealthyCode(), withInvalidMeta())
+		mockMeta := c.meta.(*mockMetaTable)
+		mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
+			return nil
+		}
+		mockMeta.OperatePrivilegeFunc = func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
+			return merr.ErrPrivilegeNotPermitted
+		}
+		err := c.initBuiltinRoles()
+		assert.Error(t, err)
+	})
+}
+
 func TestCore_Stop(t *testing.T) {
 	t.Run("abnormal stop before component is ready", func(t *testing.T) {
 		c := &Core{}

pkg/util/constant.go

@@ -61,6 +61,12 @@ const (
 
 	IdentifierKey = "identifier"
 	HeaderDBName  = "dbName"
+
+	RoleConfigPrivileges = "privileges"
+	RoleConfigObjectType = "object_type"
+	RoleConfigObjectName = "object_name"
+	RoleConfigDBName     = "db_name"
+	RoleConfigPrivilege  = "privilege"
 )
 
 const (
@@ -70,6 +76,7 @@ const (
 
 var (
 	DefaultRoles = []string{RoleAdmin, RolePublic}
+	BuiltinRoles = []string{}
 
 	ObjectPrivileges = map[string][]string{
 		commonpb.ObjectType_Collection.String(): {

pkg/util/funcutil/func.go

@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -93,6 +94,34 @@ func MapToJSON(m map[string]string) []byte {
 	return bs
 }
 
+func JSONToRoleDetails(mStr string) (map[string](map[string]([](map[string]string))), error) {
+	buffer := make(map[string](map[string]([](map[string]string))), 0)
+	err := json.Unmarshal([]byte(mStr), &buffer)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal `builtinRoles.Roles` failed, %w", err)
+	}
+	ret := make(map[string](map[string]([](map[string]string))), 0)
+	for role, privilegesJSON := range buffer {
+		ret[role] = make(map[string]([](map[string]string)), 0)
+		privilegesArray := make([]map[string]string, 0)
+		for _, privileges := range privilegesJSON[util.RoleConfigPrivileges] {
+			privilegesArray = append(privilegesArray, map[string]string{
+				util.RoleConfigObjectType: privileges[util.RoleConfigObjectType],
+				util.RoleConfigObjectName: privileges[util.RoleConfigObjectName],
+				util.RoleConfigPrivilege:  privileges[util.RoleConfigPrivilege],
+				util.RoleConfigDBName:     privileges[util.RoleConfigDBName],
+			})
+		}
+		ret[role]["privileges"] = privilegesArray
+	}
+	return ret, nil
+}
+
+func RoleDetailsToJSON(m map[string](map[string]([](map[string]string)))) []byte {
+	bs, _ := json.Marshal(m)
+	return bs
+}
+
 const (
 	// PulsarMaxMessageSizeKey is the key of config item
 	PulsarMaxMessageSizeKey = "maxMessageSize"

pkg/util/funcutil/func_test.go

@@ -33,6 +33,7 @@ import (
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus/pkg/util"
 )
 
 func Test_CheckGrpcReady(t *testing.T) {
@@ -89,6 +90,25 @@ func Test_ParseIndexParamsMap(t *testing.T) {
 	assert.NotEqual(t, err, nil)
 }
 
+func Test_ParseBuiltinRolesMap(t *testing.T) {
+	t.Run("correct format", func(t *testing.T) {
+		builtinRoles := `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}}`
+		rolePrivilegesMap, err := JSONToRoleDetails(builtinRoles)
+		assert.Nil(t, err)
+		for role, privilegesJSON := range rolePrivilegesMap {
+			assert.Contains(t, []string{"db_admin", "db_rw", "db_ro"}, role)
+			for _, privileges := range privilegesJSON[util.RoleConfigPrivileges] {
+				assert.Equal(t, privileges[util.RoleConfigObjectType], "Global")
+			}
+		}
+	})
+	t.Run("wrong format", func(t *testing.T) {
+		builtinRoles := `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}`
+		_, err := JSONToRoleDetails(builtinRoles)
+		assert.NotNil(t, err)
+	})
+}
+
 func TestGetAttrByKeyFromRepeatedKV(t *testing.T) {
 	kvs := []*commonpb.KeyValuePair{
 		{Key: "Key1", Value: "Value1"},

pkg/util/paramtable/component_param.go

@@ -69,6 +69,7 @@ type ComponentParam struct {
 	IndexNodeCfg  indexNodeConfig
 	HTTPCfg       httpConfig
 	LogCfg        logConfig
+	RoleCfg       roleConfig
 
 	RootCoordGrpcServerCfg  GrpcServerConfig
 	ProxyGrpcServerCfg      GrpcServerConfig
@@ -116,6 +117,7 @@ func (p *ComponentParam) init(bt *BaseTable) {
 	p.IndexNodeCfg.init(bt)
 	p.HTTPCfg.init(bt)
 	p.LogCfg.init(bt)
+	p.RoleCfg.init(bt)
 
 	p.RootCoordGrpcServerCfg.Init("rootCoord", bt)
 	p.ProxyGrpcServerCfg.Init("proxy", bt)

pkg/util/paramtable/param_item.go

@@ -141,6 +141,10 @@ func (pi *ParamItem) GetAsJSONMap() map[string]string {
 	return getAndConvert(pi.GetValue(), funcutil.JSONToMap, nil)
 }
 
+func (pi *ParamItem) GetAsRoleDetails() map[string](map[string]([](map[string]string))) {
+	return getAndConvert(pi.GetValue(), funcutil.JSONToRoleDetails, nil)
+}
+
 type CompositeParamItem struct {
 	Items  []*ParamItem
 	Format func(map[string]string) string

pkg/util/paramtable/role_param.go

@@ -0,0 +1,45 @@
+package paramtable
+
+import (
+	"github.com/milvus-io/milvus/pkg/config"
+	"github.com/milvus-io/milvus/pkg/util/funcutil"
+)
+
+type roleConfig struct {
+	Enabled ParamItem `refreshable:"false"`
+	Roles   ParamItem `refreshable:"false"`
+}
+
+func (p *roleConfig) init(base *BaseTable) {
+	p.Enabled = ParamItem{
+		Key:          "builtinRoles.enable",
+		DefaultValue: "false",
+		Version:      "2.3.4",
+		Doc:          "Whether to init builtin roles",
+		Export:       true,
+	}
+	p.Enabled.Init(base.mgr)
+
+	p.Roles = ParamItem{
+		Key:          "builtinRoles.roles",
+		DefaultValue: `{}`,
+		Version:      "2.3.4",
+		Doc:          "what builtin roles should be init",
+		Export:       true,
+	}
+	p.Roles.Init(base.mgr)
+
+	p.panicIfNotValid(base.mgr)
+}
+
+func (p *roleConfig) panicIfNotValid(mgr *config.Manager) {
+	if p.Enabled.GetAsBool() {
+		m := p.Roles.GetAsRoleDetails()
+		if m == nil {
+			panic("builtinRoles.roles not invalid, should be json format")
+		}
+
+		j := funcutil.RoleDetailsToJSON(m)
+		mgr.SetConfig("builtinRoles.roles", string(j))
+	}
+}

pkg/util/paramtable/role_param_test.go

@@ -0,0 +1,57 @@
+package paramtable
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus/pkg/config"
+)
+
+func TestRoleConfig_Init(t *testing.T) {
+	params := ComponentParam{}
+	params.Init(NewBaseTable(SkipRemote(true)))
+	cfg := &params.RoleCfg
+	assert.Equal(t, cfg.Enabled.GetAsBool(), false)
+	assert.Equal(t, cfg.Roles.GetValue(), "{}")
+	assert.Equal(t, len(cfg.Roles.GetAsJSONMap()), 0)
+}
+
+func TestRoleConfig_Invalid(t *testing.T) {
+	t.Run("valid roles", func(t *testing.T) {
+		mgr := config.NewManager()
+		mgr.SetConfig("builtinRoles.enable", "true")
+		mgr.SetConfig("builtinRoles.roles", `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}}`)
+		p := &roleConfig{
+			Enabled: ParamItem{
+				Key: "builtinRoles.enable",
+			},
+			Roles: ParamItem{
+				Key: "builtinRoles.roles",
+			},
+		}
+		p.Enabled.Init(mgr)
+		p.Roles.Init(mgr)
+		assert.NotPanics(t, func() {
+			p.panicIfNotValid(mgr)
+		})
+	})
+	t.Run("invalid roles", func(t *testing.T) {
+		mgr := config.NewManager()
+		mgr.SetConfig("builtinRoles.enable", "true")
+		mgr.SetConfig("builtinRoles.roles", `{"db_admin": {"privileges": {"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}}}`)
+		p := &roleConfig{
+			Enabled: ParamItem{
+				Key: "builtinRoles.enable",
+			},
+			Roles: ParamItem{
+				Key: "builtinRoles.roles",
+			},
+		}
+		p.Enabled.Init(mgr)
+		p.Roles.Init(mgr)
+		assert.Panics(t, func() {
+			p.panicIfNotValid(mgr)
+		})
+	})
+}