Support RBAC custom privilege group

Signed-off-by: shaoting-huang <[email protected]>

go.mod

@@ -137,6 +137,7 @@ require (
 	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
 	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/mock v1.6.0
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/flatbuffers v2.0.8+incompatible // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
@@ -274,3 +275,5 @@ replace (
 )
 
 exclude github.com/apache/pulsar-client-go/oauth2 v0.0.0-20211108044248-fe3b7c4e445b
+
+replace github.com/milvus-io/milvus-proto/go-api/v2 => github.com/shaoting-huang/milvus-proto/go-api/v2 v2.0.0-20241024035527-5500ce3757be

go.sum

@@ -363,6 +363,8 @@ github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt
 github.com/golang/mock v1.4.4/go.mod h1:l3mdAwkq5BuhzHwde/uurv3sEJeZMXNpwsxVWU71h+4=
 github.com/golang/mock v1.5.0 h1:jlYHihg//f7RRwuPfptm04yp4s7O6Kw8EZiVYIGcH0g=
 github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
+github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
+github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -803,6 +805,8 @@ github.com/schollz/closestmatch v2.1.0+incompatible/go.mod h1:RtP1ddjLong6gTkbtm
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
 github.com/sergi/go-diff v1.0.0/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=
 github.com/sergi/go-diff v1.1.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
+github.com/shaoting-huang/milvus-proto/go-api/v2 v2.0.0-20241024035527-5500ce3757be h1:Sp42nFLt9t0CsFEDzbMyIhyLzpUUcLyJEc4akCyD6B4=
+github.com/shaoting-huang/milvus-proto/go-api/v2 v2.0.0-20241024035527-5500ce3757be/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
 github.com/shirou/gopsutil/v3 v3.22.9 h1:yibtJhIVEMcdw+tCTbOPiF1VcsuDeTE4utJ8Dm4c5eA=
 github.com/shirou/gopsutil/v3 v3.22.9/go.mod h1:bBYl1kjgEJpWpxeHmLI+dVHWtyAwfcmSBLDsp2TNT8A=
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=
@@ -1323,6 +1327,7 @@ golang.org/x/tools v0.0.0-20201208233053-a543418bbed2/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.2/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.3/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=

internal/distributed/proxy/httpserver/constant.go

@@ -9,14 +9,15 @@ import (
 // v2
 const (
 	// --- category ---
-	CollectionCategory = "/collections/"
-	EntityCategory     = "/entities/"
-	PartitionCategory  = "/partitions/"
-	UserCategory       = "/users/"
-	RoleCategory       = "/roles/"
-	IndexCategory      = "/indexes/"
-	AliasCategory      = "/aliases/"
-	ImportJobCategory  = "/jobs/import/"
+	CollectionCategory     = "/collections/"
+	EntityCategory         = "/entities/"
+	PartitionCategory      = "/partitions/"
+	UserCategory           = "/users/"
+	RoleCategory           = "/roles/"
+	IndexCategory          = "/indexes/"
+	AliasCategory          = "/aliases/"
+	ImportJobCategory      = "/jobs/import/"
+	PrivilegeGroupCategory = "/privilege_groups/"
 
 	ListAction           = "list"
 	HasAction            = "has"
@@ -37,13 +38,15 @@ const (
 	AdvancedSearchAction = "advanced_search"
 	HybridSearchAction   = "hybrid_search"
 
-	UpdatePasswordAction  = "update_password"
-	GrantRoleAction       = "grant_role"
-	RevokeRoleAction      = "revoke_role"
-	GrantPrivilegeAction  = "grant_privilege"
-	RevokePrivilegeAction = "revoke_privilege"
-	AlterAction           = "alter"
-	GetProgressAction     = "get_progress" // deprecated, keep it for compatibility, use `/v2/vectordb/jobs/import/describe` instead
+	UpdatePasswordAction          = "update_password"
+	GrantRoleAction               = "grant_role"
+	RevokeRoleAction              = "revoke_role"
+	GrantPrivilegeAction          = "grant_privilege"
+	RevokePrivilegeAction         = "revoke_privilege"
+	AlterAction                   = "alter"
+	GetProgressAction             = "get_progress" // deprecated, keep it for compatibility, use `/v2/vectordb/jobs/import/describe` instead
+	AddPrivilegesToGroupAction    = "add_privileges_to_group"
+	DropPrivilegesFromGroupAction = "drop_privileges_from_group"
 )
 
 const (

internal/distributed/proxy/httpserver/handler_v2.go

@@ -123,6 +123,13 @@ func (h *HandlersV2) RegisterRoutesToV2(router gin.IRouter) {
 	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))))
 
+	// privilege group
+	router.POST(PrivilegeGroupCategory+CreateAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.createPrivilegeGroup))))
+	router.POST(PrivilegeGroupCategory+DropAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.dropPrivilegeGroup))))
+	router.POST(PrivilegeGroupCategory+ListAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.listPrivilegeGroups))))
+	router.POST(PrivilegeGroupCategory+AddPrivilegesToGroupAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.addPrivilegesToGroup))))
+	router.POST(PrivilegeGroupCategory+DropPrivilegesFromGroupAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.dropPrivilegesFromGroup))))
+
 	router.POST(IndexCategory+ListAction, timeoutMiddleware(wrapperPost(func() any { return &CollectionNameReq{} }, wrapperTraceLog(h.wrapperCheckDatabase(h.listIndexes)))))
 	router.POST(IndexCategory+DescribeAction, timeoutMiddleware(wrapperPost(func() any { return &IndexReq{} }, wrapperTraceLog(h.wrapperCheckDatabase(h.describeIndex)))))
 
@@ -1760,6 +1767,80 @@ func (h *HandlersV2) removePrivilegeFromRole(ctx context.Context, c *gin.Context
 	return h.operatePrivilegeToRole(ctx, c, anyReq.(*GrantReq), milvuspb.OperatePrivilegeType_Revoke, dbName)
 }
 
+func (h *HandlersV2) createPrivilegeGroup(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	httpReq := anyReq.(*PrivilegeGroupReq)
+	req := &milvuspb.CreatePrivilegeGroupRequest{
+		GroupName:  httpReq.PrivilegeGroupName,
+		Privileges: httpReq.Privileges,
+	}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/CreatePrivilegeGroup", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.CreatePrivilegeGroup(reqCtx, req.(*milvuspb.CreatePrivilegeGroupRequest))
+	})
+	if err == nil {
+		HTTPReturn(c, http.StatusOK, wrapperReturnDefault())
+	}
+	return resp, err
+}
+
+func (h *HandlersV2) dropPrivilegeGroup(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	httpReq := anyReq.(*PrivilegeGroupReq)
+	req := &milvuspb.DropPrivilegeGroupRequest{
+		GroupName: httpReq.PrivilegeGroupName,
+	}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/DropPrivilegeGroup", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.DropPrivilegeGroup(reqCtx, req.(*milvuspb.DropPrivilegeGroupRequest))
+	})
+	if err == nil {
+		HTTPReturn(c, http.StatusOK, wrapperReturnDefault())
+	}
+	return resp, err
+}
+
+func (h *HandlersV2) listPrivilegeGroups(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	req := &milvuspb.ListPrivilegeGroupsRequest{}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/ListPrivilegeGroups", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.ListPrivilegeGroups(reqCtx, req.(*milvuspb.ListPrivilegeGroupsRequest))
+	})
+	if err == nil {
+		groupNames := []string{}
+		for _, group := range resp.(*milvuspb.ListPrivilegeGroupsResponse).Groups {
+			groupNames = append(groupNames, group)
+		}
+		HTTPReturn(c, http.StatusOK, wrapperReturnList(groupNames))
+	}
+	return resp, err
+}
+
+func (h *HandlersV2) addPrivilegesToGroup(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	httpReq := anyReq.(*PrivilegeGroupReq)
+	req := &milvuspb.AddPrivilegesToGroupRequest{
+		GroupName:  httpReq.PrivilegeGroupName,
+		Privileges: httpReq.Privileges,
+	}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/AddPrivilegesToGroup", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.AddPrivilegesToGroup(reqCtx, req.(*milvuspb.AddPrivilegesToGroupRequest))
+	})
+	if err == nil {
+		HTTPReturn(c, http.StatusOK, wrapperReturnDefault())
+	}
+	return resp, err
+}
+
+func (h *HandlersV2) dropPrivilegesFromGroup(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
+	httpReq := anyReq.(*PrivilegeGroupReq)
+	req := &milvuspb.DropPrivilegesFromGroupRequest{
+		GroupName:  httpReq.PrivilegeGroupName,
+		Privileges: httpReq.Privileges,
+	}
+	resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/DropPrivilegesFromGroups", func(reqCtx context.Context, req any) (interface{}, error) {
+		return h.proxy.DropPrivilegesFromGroup(reqCtx, req.(*milvuspb.DropPrivilegesFromGroupRequest))
+	})
+	if err == nil {
+		HTTPReturn(c, http.StatusOK, wrapperReturnDefault())
+	}
+	return resp, err
+}
+
 func (h *HandlersV2) listIndexes(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) {
 	collectionGetter, _ := anyReq.(requestutil.CollectionNameGetter)
 	indexNames := []string{}

internal/distributed/proxy/httpserver/request_v2.go

@@ -6,6 +6,7 @@ import (
 	"github.com/gin-gonic/gin"
 
 	"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/merr"
 )
 
@@ -273,6 +274,11 @@ func (req *RoleReq) GetRoleName() string {
 	return req.RoleName
 }
 
+type PrivilegeGroupReq struct {
+	PrivilegeGroupName string                      `json:"privilegeGroupName" binding:"required"`
+	Privileges         []*milvuspb.PrivilegeEntity `json:"privileges"`
+}
+
 type GrantReq struct {
 	RoleName   string `json:"roleName" binding:"required"`
 	ObjectType string `json:"objectType" binding:"required"`

internal/distributed/proxy/service.go

@@ -993,6 +993,26 @@ func (s *Server) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaR
 	return s.proxy.RestoreRBAC(ctx, req)
 }
 
+func (s *Server) CreatePrivilegeGroup(ctx context.Context, req *milvuspb.CreatePrivilegeGroupRequest) (*commonpb.Status, error) {
+	return s.proxy.CreatePrivilegeGroup(ctx, req)
+}
+
+func (s *Server) DropPrivilegeGroup(ctx context.Context, req *milvuspb.DropPrivilegeGroupRequest) (*commonpb.Status, error) {
+	return s.proxy.DropPrivilegeGroup(ctx, req)
+}
+
+func (s *Server) ListPrivilegeGroups(ctx context.Context, req *milvuspb.ListPrivilegeGroupsRequest) (*milvuspb.ListPrivilegeGroupsResponse, error) {
+	return s.proxy.ListPrivilegeGroups(ctx, req)
+}
+
+func (s *Server) AddPrivilegesToGroup(ctx context.Context, req *milvuspb.AddPrivilegesToGroupRequest) (*commonpb.Status, error) {
+	return s.proxy.AddPrivilegesToGroup(ctx, req)
+}
+
+func (s *Server) DropPrivilegesFromGroup(ctx context.Context, req *milvuspb.DropPrivilegesFromGroupRequest) (*commonpb.Status, error) {
+	return s.proxy.DropPrivilegesFromGroup(ctx, req)
+}
+
 func (s *Server) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
 	return s.proxy.RefreshPolicyInfoCache(ctx, req)
 }

internal/distributed/rootcoord/client/client.go

@@ -706,3 +706,63 @@ func (c *Client) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRe
 		return client.RestoreRBAC(ctx, in)
 	})
 }
+
+func (c *Client) CreatePrivilegeGroup(ctx context.Context, in *milvuspb.CreatePrivilegeGroupRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	in = typeutil.Clone(in)
+	commonpbutil.UpdateMsgBase(
+		in.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
+	)
+
+	return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*commonpb.Status, error) {
+		return client.CreatePrivilegeGroup(ctx, in)
+	})
+}
+
+func (c *Client) DropPrivilegeGroup(ctx context.Context, in *milvuspb.DropPrivilegeGroupRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	in = typeutil.Clone(in)
+	commonpbutil.UpdateMsgBase(
+		in.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
+	)
+
+	return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*commonpb.Status, error) {
+		return client.DropPrivilegeGroup(ctx, in)
+	})
+}
+
+func (c *Client) ListPrivilegeGroups(ctx context.Context, in *milvuspb.ListPrivilegeGroupsRequest, opts ...grpc.CallOption) (*milvuspb.ListPrivilegeGroupsResponse, error) {
+	in = typeutil.Clone(in)
+	commonpbutil.UpdateMsgBase(
+		in.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
+	)
+
+	return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*milvuspb.ListPrivilegeGroupsResponse, error) {
+		return client.ListPrivilegeGroups(ctx, in)
+	})
+}
+
+func (c *Client) AddPrivilegesToGroup(ctx context.Context, in *milvuspb.AddPrivilegesToGroupRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	in = typeutil.Clone(in)
+	commonpbutil.UpdateMsgBase(
+		in.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
+	)
+
+	return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*commonpb.Status, error) {
+		return client.AddPrivilegesToGroup(ctx, in)
+	})
+}
+
+func (c *Client) DropPrivilegesFromGroup(ctx context.Context, in *milvuspb.DropPrivilegesFromGroupRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	in = typeutil.Clone(in)
+	commonpbutil.UpdateMsgBase(
+		in.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
+	)
+
+	return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*commonpb.Status, error) {
+		return client.DropPrivilegesFromGroup(ctx, in)
+	})
+}

internal/distributed/rootcoord/service.go

@@ -550,3 +550,23 @@ func (s *Server) BackupRBAC(ctx context.Context, request *milvuspb.BackupRBACMet
 func (s *Server) RestoreRBAC(ctx context.Context, request *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
 	return s.rootCoord.RestoreRBAC(ctx, request)
 }
+
+func (s *Server) CreatePrivilegeGroup(ctx context.Context, request *milvuspb.CreatePrivilegeGroupRequest) (*commonpb.Status, error) {
+	return s.rootCoord.CreatePrivilegeGroup(ctx, request)
+}
+
+func (s *Server) DropPrivilegeGroup(ctx context.Context, request *milvuspb.DropPrivilegeGroupRequest) (*commonpb.Status, error) {
+	return s.rootCoord.DropPrivilegeGroup(ctx, request)
+}
+
+func (s *Server) ListPrivilegeGroups(ctx context.Context, request *milvuspb.ListPrivilegeGroupsRequest) (*milvuspb.ListPrivilegeGroupsResponse, error) {
+	return s.rootCoord.ListPrivilegeGroups(ctx, request)
+}
+
+func (s *Server) AddPrivilegesToGroup(ctx context.Context, request *milvuspb.AddPrivilegesToGroupRequest) (*commonpb.Status, error) {
+	return s.rootCoord.AddPrivilegesToGroup(ctx, request)
+}
+
+func (s *Server) DropPrivilegesFromGroup(ctx context.Context, request *milvuspb.DropPrivilegesFromGroupRequest) (*commonpb.Status, error) {
+	return s.rootCoord.DropPrivilegesFromGroup(ctx, request)
+}

internal/metastore/catalog.go

@@ -86,6 +86,11 @@ type RootCoordCatalog interface {
 	BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
 	RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
 
+	GetPrivilegeGroup(ctx context.Context, groupName string) ([]*milvuspb.PrivilegeEntity, error)
+	DropPrivilegeGroup(ctx context.Context, groupName string) error
+	AlterPrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
+	ListPrivilegeGroups(ctx context.Context) ([]string, error)
+
 	Close()
 }