RBAC new grant/revoke privilege

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>

internal/distributed/proxy/httpserver/constant.go

@@ -43,6 +43,8 @@ const (
 	RevokeRoleAction                = "revoke_role"
 	GrantPrivilegeAction            = "grant_privilege"
 	RevokePrivilegeAction           = "revoke_privilege"
+	GrantPrivilegeActionV2          = "grant_privilege_v2"
+	RevokePrivilegeActionV2         = "revoke_privilege_v2"
 	AlterAction                     = "alter"
 	GetProgressAction               = "get_progress" // deprecated, keep it for compatibility, use `/v2/vectordb/jobs/import/describe` instead
 	AddPrivilegesToGroupAction      = "add_privileges_to_group"

internal/distributed/proxy/httpserver/handler_v2.go

@@ -131,6 +131,8 @@ func (h *HandlersV2) RegisterRoutesToV2(router gin.IRouter) {
 	router.POST(RoleCategory+DropAction, timeoutMiddleware(wrapperPost(func() any { return &RoleReq{} }, wrapperTraceLog(h.dropRole))))
 	router.POST(RoleCategory+GrantPrivilegeAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.addPrivilegeToRole))))
 	router.POST(RoleCategory+RevokePrivilegeAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.removePrivilegeFromRole))))
+	router.POST(RoleCategory+GrantPrivilegeActionV2, timeoutMiddleware(wrapperPost(func() any { return &GrantV2Req{} }, wrapperTraceLog(h.grantV2))))
+	router.POST(RoleCategory+RevokePrivilegeActionV2, timeoutMiddleware(wrapperPost(func() any { return &GrantV2Req{} }, wrapperTraceLog(h.revokeV2))))
 
 	// privilege group
 	router.POST(PrivilegeGroupCategory+CreateAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.createPrivilegeGroup))))
@@ -1702,6 +1704,37 @@ func (h *HandlersV2) operatePrivilegeToRole(ctx context.Context, c *gin.Context,
 	return resp, err
 }
 
+func (h *HandlersV2) operatePrivilegeToRoleV2(ctx context.Context, c *gin.Context, httpReq *GrantV2Req, operateType milvuspb.OperatePrivilegeType) (interface{}, error) {
+	req := &milvuspb.OperatePrivilegeRequest{
+		Entity: &milvuspb.GrantEntity{
+			Role:       &milvuspb.RoleEntity{Name: httpReq.RoleName},
+			Object:     &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
+			ObjectName: httpReq.CollectionName,
+			DbName:     httpReq.DbName,
+			Grantor: &milvuspb.GrantorEntity{
+				Privilege: &milvuspb.PrivilegeEntity{Name: httpReq.Privilege},
+			},
+		},
+		Type:    operateType,
+		Version: "v2",
+	}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/OperatePrivilege", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.OperatePrivilege(reqCtx, req.(*milvuspb.OperatePrivilegeRequest))
+	})
+	if err == nil {
+		HTTPReturn(c, http.StatusOK, wrapperReturnDefault())
+	}
+	return resp, err
+}
+
+func (h *HandlersV2) grantV2(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	return h.operatePrivilegeToRoleV2(ctx, c, anyReq.(*GrantV2Req), milvuspb.OperatePrivilegeType_Grant)
+}
+
+func (h *HandlersV2) revokeV2(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	return h.operatePrivilegeToRoleV2(ctx, c, anyReq.(*GrantV2Req), milvuspb.OperatePrivilegeType_Revoke)
+}
+
 func (h *HandlersV2) addPrivilegeToRole(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
 	return h.operatePrivilegeToRole(ctx, c, anyReq.(*GrantReq), milvuspb.OperatePrivilegeType_Grant, dbName)
 }

internal/distributed/proxy/httpserver/handler_v2_test.go

@@ -1106,7 +1106,7 @@ func TestMethodPost(t *testing.T) {
 	mp.EXPECT().UpdateCredential(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
 	mp.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
 	mp.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
-	mp.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
+	mp.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(4)
 	mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
 	mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonErrorStatus, nil).Once()
 	mp.EXPECT().CreateAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
@@ -1179,6 +1179,12 @@ func TestMethodPost(t *testing.T) {
 	queryTestCases = append(queryTestCases, rawTestCase{
 		path: versionalV2(RoleCategory, RevokePrivilegeAction),
 	})
+	queryTestCases = append(queryTestCases, rawTestCase{
+		path: versionalV2(RoleCategory, GrantPrivilegeActionV2),
+	})
+	queryTestCases = append(queryTestCases, rawTestCase{
+		path: versionalV2(RoleCategory, RevokePrivilegeActionV2),
+	})
 	queryTestCases = append(queryTestCases, rawTestCase{
 		path: versionalV2(IndexCategory, CreateAction),
 	})

internal/distributed/proxy/httpserver/request_v2.go

@@ -273,6 +273,13 @@ type PrivilegeGroupReq struct {
 	Privileges         []string `json:"privileges"`
 }
 
+type GrantV2Req struct {
+	RoleName       string `json:"roleName" binding:"required"`
+	DbName         string `json:"dbName"`
+	CollectionName string `json:"collectionName"`
+	Privilege      string `json:"privilege" binding:"required"`
+}
+
 type GrantReq struct {
 	RoleName   string `json:"roleName" binding:"required"`
 	ObjectType string `json:"objectType" binding:"required"`

internal/proxy/impl.go

@@ -5210,6 +5210,87 @@ func (node *Proxy) validPrivilegeParams(req *milvuspb.OperatePrivilegeRequest) e
 	return nil
 }
 
+func (node *Proxy) validOperatePrivilegeV2Params(req *milvuspb.OperatePrivilegeV2Request) error {
+	if req.Role == nil {
+		return fmt.Errorf("the role in the request is nil")
+	}
+	if err := ValidateRoleName(req.Role.Name); err != nil {
+		return err
+	}
+	if err := ValidatePrivilege(req.Grantor.Privilege.Name); err != nil {
+		return err
+	}
+	if req.Type != milvuspb.OperatePrivilegeType_Grant && req.Type != milvuspb.OperatePrivilegeType_Revoke {
+		return fmt.Errorf("the type in the request not grant or revoke")
+	}
+	if err := ValidateObjectName(req.DbName); err != nil {
+		return err
+	}
+	if err := ValidateObjectName(req.CollectionName); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (node *Proxy) OperatePrivilegeV2(ctx context.Context, req *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error) {
+	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-OperatePrivilegeV2")
+	defer sp.End()
+
+	log := log.Ctx(ctx)
+
+	log.Info("OperatePrivilegeV2",
+		zap.Any("req", req))
+	if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
+		return merr.Status(err), nil
+	}
+	if err := node.validOperatePrivilegeV2Params(req); err != nil {
+		return merr.Status(err), nil
+	}
+	curUser, err := GetCurUserFromContext(ctx)
+	if err != nil {
+		return merr.Status(err), nil
+	}
+	req.Grantor.User = &milvuspb.UserEntity{Name: curUser}
+	request := &milvuspb.OperatePrivilegeRequest{
+		Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_OperatePrivilege},
+		Entity: &milvuspb.GrantEntity{
+			Role:       req.Role,
+			Object:     &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
+			ObjectName: req.CollectionName,
+			DbName:     req.DbName,
+			Grantor:    req.Grantor,
+		},
+		Type:    req.Type,
+		Version: "v2",
+	}
+	req.Grantor.User = &milvuspb.UserEntity{Name: curUser}
+	result, err := node.rootCoord.OperatePrivilege(ctx, request)
+	if err != nil {
+		log.Warn("fail to operate privilege", zap.Error(err))
+		return merr.Status(err), nil
+	}
+	relatedPrivileges := util.RelatedPrivileges[util.PrivilegeNameForMetastore(req.Grantor.Privilege.Name)]
+	if len(relatedPrivileges) != 0 {
+		for _, relatedPrivilege := range relatedPrivileges {
+			relatedReq := proto.Clone(request).(*milvuspb.OperatePrivilegeRequest)
+			relatedReq.Entity.Grantor.Privilege.Name = util.PrivilegeNameForAPI(relatedPrivilege)
+			result, err = node.rootCoord.OperatePrivilege(ctx, relatedReq)
+			if err != nil {
+				log.Warn("fail to operate related privilege", zap.String("related_privilege", relatedPrivilege), zap.Error(err))
+				return merr.Status(err), nil
+			}
+			if !merr.Ok(result) {
+				log.Warn("fail to operate related privilege", zap.String("related_privilege", relatedPrivilege), zap.Any("result", result))
+				return result, nil
+			}
+		}
+	}
+	if merr.Ok(result) {
+		SendReplicateMessagePack(ctx, node.replicateMsgStream, req)
+	}
+	return result, nil
+}
+
 func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-OperatePrivilege")
 	defer sp.End()
@@ -6399,8 +6480,11 @@ func (node *Proxy) CreatePrivilegeGroup(ctx context.Context, req *milvuspb.Creat
 	if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
 		return merr.Status(err), nil
 	}
-	if req.GroupName == "" {
-		return merr.Status(fmt.Errorf("the group name in the drop privilege group request is nil")), nil
+	if err := ValidatePrivilegeGroupName(req.GroupName); err != nil {
+		log.Warn("CreatePrivilegeGroup failed",
+			zap.Error(err),
+		)
+		return getErrResponse(err, "CreatePrivilegeGroup", "", ""), nil
 	}
 	if req.Base == nil {
 		req.Base = &commonpb.MsgBase{}
@@ -6428,8 +6512,11 @@ func (node *Proxy) DropPrivilegeGroup(ctx context.Context, req *milvuspb.DropPri
 	if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
 		return merr.Status(err), nil
 	}
-	if req.GroupName == "" {
-		return merr.Status(fmt.Errorf("the group name in the drop privilege group request is nil")), nil
+	if err := ValidatePrivilegeGroupName(req.GroupName); err != nil {
+		log.Warn("DropPrivilegeGroup failed",
+			zap.Error(err),
+		)
+		return getErrResponse(err, "DropPrivilegeGroup", "", ""), nil
 	}
 	if req.Base == nil {
 		req.Base = &commonpb.MsgBase{}
@@ -6486,8 +6573,11 @@ func (node *Proxy) OperatePrivilegeGroup(ctx context.Context, req *milvuspb.Oper
 	if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
 		return merr.Status(err), nil
 	}
-	if req.GroupName == "" {
-		return merr.Status(fmt.Errorf("the group name in the drop privilege group request is nil")), nil
+	if err := ValidatePrivilegeGroupName(req.GroupName); err != nil {
+		log.Warn("OperatePrivilegeGroup failed",
+			zap.Error(err),
+		)
+		return getErrResponse(err, "OperatePrivilegeGroup", "", ""), nil
 	}
 	for _, priv := range req.GetPrivileges() {
 		if err := ValidatePrivilege(priv.Name); err != nil {

internal/proxy/privilege_interceptor.go

@@ -162,8 +162,8 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
 
 	e := getEnforcer()
 	for _, roleName := range roleNames {
-		permitFunc := func(resName string) (bool, error) {
-			object := funcutil.PolicyForResource(dbName, objectType, resName)
+		permitFunc := func(objectName string) (bool, error) {
+			object := funcutil.PolicyForResource(dbName, objectType, objectName)
 			isPermit, cached, version := GetPrivilegeCache(roleName, object, objectPrivilege)
 			if cached {
 				return isPermit, nil

internal/proxy/privilege_interceptor_test.go

@@ -563,3 +563,45 @@ func TestPrivilegeGroup(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+func TestBuiltinPrivilegeGroup(t *testing.T) {
+	t.Run("ClusterAdmin", func(t *testing.T) {
+		paramtable.Get().Save(Params.CommonCfg.AuthorizationEnabled.Key, "true")
+		initPrivilegeGroups()
+
+		var err error
+		ctx := GetContext(context.Background(), "fooo:123456")
+		client := &MockRootCoordClientInterface{}
+		queryCoord := &mocks.MockQueryCoordClient{}
+		mgr := newShardClientMgr()
+
+		policies := []string{}
+		for _, priv := range util.BuiltinPrivilegeGroups["ClusterReadOnly"] {
+			objectType := util.GetObjectType(priv)
+			policies = append(policies, funcutil.PolicyForPrivilege("role1", objectType, "*", util.PrivilegeNameForMetastore(priv), "default"))
+		}
+		client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
+			return &internalpb.ListPolicyResponse{
+				Status:      merr.Success(),
+				PolicyInfos: policies,
+				UserRoles: []string{
+					funcutil.EncodeUserRoleCache("fooo", "role1"),
+				},
+			}, nil
+		}
+		InitMetaCache(ctx, client, queryCoord, mgr)
+		defer CleanPrivilegeCache()
+
+		_, err = PrivilegeInterceptor(GetContext(context.Background(), "fooo:123456"), &milvuspb.SelectUserRequest{})
+		assert.NoError(t, err)
+
+		_, err = PrivilegeInterceptor(GetContext(context.Background(), "fooo:123456"), &milvuspb.DescribeResourceGroupRequest{})
+		assert.NoError(t, err)
+
+		_, err = PrivilegeInterceptor(GetContext(context.Background(), "fooo:123456"), &milvuspb.ListResourceGroupsRequest{})
+		assert.NoError(t, err)
+
+		_, err = PrivilegeInterceptor(GetContext(context.Background(), "fooo:123456"), &milvuspb.CreateResourceGroupRequest{})
+		assert.Error(t, err)
+	})
+}

internal/proxy/proxy_test.go

@@ -2700,6 +2700,8 @@ func TestProxy(t *testing.T) {
 
 	testProxyRole(ctx, t, proxy)
 	testProxyPrivilege(ctx, t, proxy)
+	testProxyOperatePrivilegeV2(ctx, t, proxy)
+	assert.False(t, false, true)
 	testProxyRefreshPolicyInfoCache(ctx, t, proxy)
 
 	// proxy unhealthy
@@ -4368,6 +4370,106 @@ func testProxyPrivilege(ctx context.Context, t *testing.T, proxy *Proxy) {
 	wg.Wait()
 }
 
+func testProxyOperatePrivilegeV2(ctx context.Context, t *testing.T, proxy *Proxy) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	t.Run("Operate Privilege V2, Select Grant", func(t *testing.T) {
+		defer wg.Done()
+
+		// GrantPrivilege
+		req := &milvuspb.OperatePrivilegeV2Request{}
+		resp, _ := proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.Role = &milvuspb.RoleEntity{}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.Grantor = &milvuspb.GrantorEntity{}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.Grantor.Privilege = &milvuspb.PrivilegeEntity{}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.Grantor.Privilege.Name = util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeAll.String())
+
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.DbName = ""
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		req.CollectionName = ""
+		resp, _ = proxy.OperatePrivilegeV2(ctx, req)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq := &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeLoad.String())}},
+			DbName:         "default",
+			CollectionName: "col1",
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq = &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGroupClusterReadOnly.String())}},
+			DbName:         util.AnyWord,
+			CollectionName: util.AnyWord,
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq = &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGroupClusterReadOnly.String())}},
+			DbName:         "db1",
+			CollectionName: util.AnyWord,
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq = &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGroupCollectionReadOnly.String())}},
+			DbName:         "db1",
+			CollectionName: util.AnyWord,
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq = &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGroupClusterReadOnly.String())}},
+			DbName:         "db1",
+			CollectionName: "col1",
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+
+		roleReq = &milvuspb.OperatePrivilegeV2Request{
+			Role:           &milvuspb.RoleEntity{Name: "public"},
+			Grantor:        &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGroupDatabaseReadOnly.String())}},
+			DbName:         "db1",
+			CollectionName: util.AnyWord,
+			Type:           milvuspb.OperatePrivilegeType_Grant,
+		}
+		resp, _ = proxy.OperatePrivilegeV2(ctx, roleReq)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	wg.Wait()
+}
+
 func testProxyPrivilegeUnhealthy(ctx context.Context, t *testing.T, proxy *Proxy) {
 	testProxyPrivilegeFail(ctx, t, proxy, "unhealthy")
 }