Skip to content

Commit ff00a12

Browse files
shaoting-huangshaoting-huang
authored
enhance: RBAC custom privilege group ut coverage (#37558)
issue: #37031 Signed-off-by: shaoting-huang <[email protected]>
1 parent ffdde39 commit ff00a12

File tree

7 files changed

+126
-13
lines changed

7 files changed

+126
-13
lines changed

‎internal/distributed/proxy/httpserver/handler_v2.go

+10-8
Original file line numberDiff line numberDiff line change
@@ -133,11 +133,11 @@ func (h *HandlersV2) RegisterRoutesToV2(router gin.IRouter) {
133133
router.POST(RoleCategory+RevokePrivilegeAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.removePrivilegeFromRole))))
134134

135135
// privilege group
136-
router.POST(PrivilegeGroupCategory+CreateAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.createPrivilegeGroup))))
137-
router.POST(PrivilegeGroupCategory+DropAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.dropPrivilegeGroup))))
138-
router.POST(PrivilegeGroupCategory+ListAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.listPrivilegeGroups))))
139-
router.POST(PrivilegeGroupCategory+AddPrivilegesToGroupAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.addPrivilegesToGroup))))
140-
router.POST(PrivilegeGroupCategory+RemovePrivilegesFromGroupAction, timeoutMiddleware(wrapperPost(func() any { return &GrantReq{} }, wrapperTraceLog(h.removePrivilegesFromGroup))))
136+
router.POST(PrivilegeGroupCategory+CreateAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.createPrivilegeGroup))))
137+
router.POST(PrivilegeGroupCategory+DropAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.dropPrivilegeGroup))))
138+
router.POST(PrivilegeGroupCategory+ListAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.listPrivilegeGroups))))
139+
router.POST(PrivilegeGroupCategory+AddPrivilegesToGroupAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.addPrivilegesToGroup))))
140+
router.POST(PrivilegeGroupCategory+RemovePrivilegesFromGroupAction, timeoutMiddleware(wrapperPost(func() any { return &PrivilegeGroupReq{} }, wrapperTraceLog(h.removePrivilegesFromGroup))))
141141

142142
router.POST(IndexCategory+ListAction, timeoutMiddleware(wrapperPost(func() any { return &CollectionNameReq{} }, wrapperTraceLog(h.wrapperCheckDatabase(h.listIndexes)))))
143143
router.POST(IndexCategory+DescribeAction, timeoutMiddleware(wrapperPost(func() any { return &IndexReq{} }, wrapperTraceLog(h.wrapperCheckDatabase(h.describeIndex)))))
@@ -1857,9 +1857,11 @@ func (h *HandlersV2) removePrivilegesFromGroup(ctx context.Context, c *gin.Conte
18571857
func (h *HandlersV2) operatePrivilegeGroup(ctx context.Context, c *gin.Context, anyReq any, dbName string, operateType milvuspb.OperatePrivilegeGroupType) (interface{}, error) {
18581858
httpReq := anyReq.(*PrivilegeGroupReq)
18591859
req := &milvuspb.OperatePrivilegeGroupRequest{
1860-
GroupName: httpReq.PrivilegeGroupName,
1861-
Privileges: httpReq.Privileges,
1862-
Type: operateType,
1860+
GroupName: httpReq.PrivilegeGroupName,
1861+
Privileges: lo.Map(httpReq.Privileges, func(p string, _ int) *milvuspb.PrivilegeEntity {
1862+
return &milvuspb.PrivilegeEntity{Name: p}
1863+
}),
1864+
Type: operateType,
18631865
}
18641866
resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/OperatePrivilegeGroup", func(reqCtx context.Context, req any) (interface{}, error) {
18651867
return h.proxy.OperatePrivilegeGroup(reqCtx, req.(*milvuspb.OperatePrivilegeGroupRequest))

‎internal/distributed/proxy/httpserver/handler_v2_test.go

+26-2
Original file line numberDiff line numberDiff line change
@@ -1230,6 +1230,10 @@ func TestMethodGet(t *testing.T) {
12301230
Status: &StatusSuccess,
12311231
Alias: DefaultAliasName,
12321232
}, nil).Once()
1233+
mp.EXPECT().ListPrivilegeGroups(mock.Anything, mock.Anything).Return(&milvuspb.ListPrivilegeGroupsResponse{
1234+
Status: &StatusSuccess,
1235+
PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{{GroupName: "group1", Privileges: []*milvuspb.PrivilegeEntity{{Name: "*"}}}},
1236+
}, nil).Once()
12331237

12341238
testEngine := initHTTPServerV2(mp, false)
12351239
queryTestCases := []rawTestCase{}
@@ -1320,6 +1324,9 @@ func TestMethodGet(t *testing.T) {
13201324
queryTestCases = append(queryTestCases, rawTestCase{
13211325
path: versionalV2(AliasCategory, DescribeAction),
13221326
})
1327+
queryTestCases = append(queryTestCases, rawTestCase{
1328+
path: versionalV2(PrivilegeGroupCategory, ListAction),
1329+
})
13231330

13241331
for _, testcase := range queryTestCases {
13251332
t.Run(testcase.path, func(t *testing.T) {
@@ -1329,7 +1336,8 @@ func TestMethodGet(t *testing.T) {
13291336
`"indexName": "` + DefaultIndexName + `",` +
13301337
`"userName": "` + util.UserRoot + `",` +
13311338
`"roleName": "` + util.RoleAdmin + `",` +
1332-
`"aliasName": "` + DefaultAliasName + `"` +
1339+
`"aliasName": "` + DefaultAliasName + `",` +
1340+
`"privilegeGroupName": "pg"` +
13331341
`}`))
13341342
req := httptest.NewRequest(http.MethodPost, testcase.path, bodyReader)
13351343
w := httptest.NewRecorder()
@@ -1369,6 +1377,7 @@ func TestMethodDelete(t *testing.T) {
13691377
mp.EXPECT().DropRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
13701378
mp.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
13711379
mp.EXPECT().DropAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
1380+
mp.EXPECT().DropPrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
13721381
testEngine := initHTTPServerV2(mp, false)
13731382
queryTestCases := []rawTestCase{}
13741383
queryTestCases = append(queryTestCases, rawTestCase{
@@ -1389,10 +1398,13 @@ func TestMethodDelete(t *testing.T) {
13891398
queryTestCases = append(queryTestCases, rawTestCase{
13901399
path: versionalV2(AliasCategory, DropAction),
13911400
})
1401+
queryTestCases = append(queryTestCases, rawTestCase{
1402+
path: versionalV2(PrivilegeGroupCategory, DropAction),
1403+
})
13921404
for _, testcase := range queryTestCases {
13931405
t.Run(testcase.path, func(t *testing.T) {
13941406
bodyReader := bytes.NewReader([]byte(`{"collectionName": "` + DefaultCollectionName + `", "partitionName": "` + DefaultPartitionName +
1395-
`", "userName": "` + util.UserRoot + `", "roleName": "` + util.RoleAdmin + `", "indexName": "` + DefaultIndexName + `", "aliasName": "` + DefaultAliasName + `"}`))
1407+
`", "userName": "` + util.UserRoot + `", "roleName": "` + util.RoleAdmin + `", "indexName": "` + DefaultIndexName + `", "aliasName": "` + DefaultAliasName + `", "privilegeGroupName": "pg"}`))
13961408
req := httptest.NewRequest(http.MethodPost, testcase.path, bodyReader)
13971409
w := httptest.NewRecorder()
13981410
testEngine.ServeHTTP(w, req)
@@ -1431,6 +1443,8 @@ func TestMethodPost(t *testing.T) {
14311443
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonErrorStatus, nil).Once()
14321444
mp.EXPECT().CreateAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
14331445
mp.EXPECT().AlterAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
1446+
mp.EXPECT().CreatePrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
1447+
mp.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
14341448
mp.EXPECT().ImportV2(mock.Anything, mock.Anything).Return(&internalpb.ImportResponse{
14351449
Status: commonSuccessStatus, JobID: "1234567890",
14361450
}, nil).Once()
@@ -1523,6 +1537,15 @@ func TestMethodPost(t *testing.T) {
15231537
queryTestCases = append(queryTestCases, rawTestCase{
15241538
path: versionalV2(ImportJobCategory, DescribeAction),
15251539
})
1540+
queryTestCases = append(queryTestCases, rawTestCase{
1541+
path: versionalV2(PrivilegeGroupCategory, CreateAction),
1542+
})
1543+
queryTestCases = append(queryTestCases, rawTestCase{
1544+
path: versionalV2(PrivilegeGroupCategory, AddPrivilegesToGroupAction),
1545+
})
1546+
queryTestCases = append(queryTestCases, rawTestCase{
1547+
path: versionalV2(PrivilegeGroupCategory, RemovePrivilegesFromGroupAction),
1548+
})
15261549

15271550
for _, testcase := range queryTestCases {
15281551
t.Run(testcase.path, func(t *testing.T) {
@@ -1533,6 +1556,7 @@ func TestMethodPost(t *testing.T) {
15331556
`"indexParams": [{"indexName": "` + DefaultIndexName + `", "fieldName": "book_intro", "metricType": "L2", "params": {"nlist": 30, "index_type": "IVF_FLAT"}}],` +
15341557
`"userName": "` + util.UserRoot + `", "password": "Milvus", "newPassword": "milvus", "roleName": "` + util.RoleAdmin + `",` +
15351558
`"roleName": "` + util.RoleAdmin + `", "objectType": "Global", "objectName": "*", "privilege": "*",` +
1559+
`"privilegeGroupName": "pg", "privileges": ["create", "drop"],` +
15361560
`"aliasName": "` + DefaultAliasName + `",` +
15371561
`"jobId": "1234567890",` +
15381562
`"files": [["book.json"]]` +

‎internal/distributed/proxy/httpserver/request_v2.go

+2-3
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66
"github.com/gin-gonic/gin"
77

88
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
9-
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
109
"github.com/milvus-io/milvus/pkg/util/merr"
1110
)
1211

@@ -275,8 +274,8 @@ func (req *RoleReq) GetRoleName() string {
275274
}
276275

277276
type PrivilegeGroupReq struct {
278-
PrivilegeGroupName string `json:"privilegeGroupName" binding:"required"`
279-
Privileges []*milvuspb.PrivilegeEntity `json:"privileges"`
277+
PrivilegeGroupName string `json:"privilegeGroupName" binding:"required"`
278+
Privileges []string `json:"privileges"`
280279
}
281280

282281
type GrantReq struct {

‎internal/rootcoord/meta_table_test.go

+12
Original file line numberDiff line numberDiff line change
@@ -2093,10 +2093,22 @@ func TestMetaTable_PrivilegeGroup(t *testing.T) {
20932093
}
20942094
err := mt.CreatePrivilegeGroup("pg1")
20952095
assert.Error(t, err)
2096+
err = mt.CreatePrivilegeGroup("")
2097+
assert.Error(t, err)
2098+
err = mt.CreatePrivilegeGroup("Insert")
2099+
assert.Error(t, err)
20962100
err = mt.CreatePrivilegeGroup("pg2")
20972101
assert.NoError(t, err)
2102+
err = mt.DropPrivilegeGroup("")
2103+
assert.Error(t, err)
20982104
err = mt.DropPrivilegeGroup("pg1")
20992105
assert.NoError(t, err)
2106+
err = mt.OperatePrivilegeGroup("", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
2107+
assert.Error(t, err)
2108+
err = mt.OperatePrivilegeGroup("pg3", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
2109+
assert.Error(t, err)
2110+
_, err = mt.GetPrivilegeGroupRoles("")
2111+
assert.Error(t, err)
21002112
_, err = mt.ListPrivilegeGroups()
21012113
assert.NoError(t, err)
21022114
}

‎internal/rootcoord/mock_test.go

+8
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ type mockMetaTable struct {
101101
DropPrivilegeGroupFunc func(groupName string) error
102102
ListPrivilegeGroupsFunc func() ([]*milvuspb.PrivilegeGroupInfo, error)
103103
OperatePrivilegeGroupFunc func(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
104+
GetPrivilegeGroupRolesFunc func(groupName string) ([]*milvuspb.RoleEntity, error)
104105
}
105106

106107
func (m mockMetaTable) GetDatabaseByName(ctx context.Context, dbName string, ts Timestamp) (*model.Database, error) {
@@ -271,6 +272,10 @@ func (m mockMetaTable) OperatePrivilegeGroup(groupName string, privileges []*mil
271272
return m.OperatePrivilegeGroupFunc(groupName, privileges, operateType)
272273
}
273274

275+
func (m mockMetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error) {
276+
return m.GetPrivilegeGroupRolesFunc(groupName)
277+
}
278+
274279
func newMockMetaTable() *mockMetaTable {
275280
return &mockMetaTable{}
276281
}
@@ -559,6 +564,9 @@ func withInvalidMeta() Opt {
559564
meta.OperatePrivilegeGroupFunc = func(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
560565
return errors.New("error mock OperatePrivilegeGroup")
561566
}
567+
meta.GetPrivilegeGroupRolesFunc = func(groupName string) ([]*milvuspb.RoleEntity, error) {
568+
return nil, errors.New("error mock GetPrivilegeGroupRoles")
569+
}
562570
return withMeta(meta)
563571
}
564572

‎internal/rootcoord/root_coord_test.go

+33
Original file line numberDiff line numberDiff line change
@@ -1787,6 +1787,39 @@ func TestRootCoord_RBACError(t *testing.T) {
17871787
}
17881788
})
17891789

1790+
t.Run("operate privilege group failed", func(t *testing.T) {
1791+
mockMeta := c.meta.(*mockMetaTable)
1792+
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
1793+
return nil, errors.New("mock error")
1794+
}
1795+
mockMeta.CreatePrivilegeGroupFunc = func(groupName string) error {
1796+
return errors.New("mock error")
1797+
}
1798+
mockMeta.GetPrivilegeGroupRolesFunc = func(groupName string) ([]*milvuspb.RoleEntity, error) {
1799+
return nil, errors.New("mock error")
1800+
}
1801+
{
1802+
resp, err := c.OperatePrivilegeGroup(ctx, &milvuspb.OperatePrivilegeGroupRequest{})
1803+
assert.NoError(t, err)
1804+
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
1805+
}
1806+
{
1807+
resp, err := c.ListPrivilegeGroups(ctx, &milvuspb.ListPrivilegeGroupsRequest{})
1808+
assert.NoError(t, err)
1809+
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
1810+
}
1811+
{
1812+
resp, err := c.OperatePrivilegeGroup(ctx, &milvuspb.OperatePrivilegeGroupRequest{})
1813+
assert.NoError(t, err)
1814+
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
1815+
}
1816+
{
1817+
resp, err := c.CreatePrivilegeGroup(ctx, &milvuspb.CreatePrivilegeGroupRequest{})
1818+
assert.NoError(t, err)
1819+
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
1820+
}
1821+
})
1822+
17901823
t.Run("select grant failed", func(t *testing.T) {
17911824
{
17921825
resp, err := c.SelectGrant(ctx, &milvuspb.SelectGrantRequest{})

‎tests/integration/rbac/privilege_group_test.go

+35
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,41 @@ func (s *PrivilegeGroupTestSuite) TestCustomPrivilegeGroup() {
199199
s.True(merr.Ok(dropRoleResp))
200200
}
201201

202+
func (s *PrivilegeGroupTestSuite) TestInvalidPrivilegeGroup() {
203+
ctx := GetContext(context.Background(), "root:123456")
204+
205+
createResp, err := s.Cluster.Proxy.CreatePrivilegeGroup(ctx, &milvuspb.CreatePrivilegeGroupRequest{
206+
GroupName: "",
207+
})
208+
s.NoError(err)
209+
s.False(merr.Ok(createResp))
210+
211+
dropResp, err := s.Cluster.Proxy.DropPrivilegeGroup(ctx, &milvuspb.DropPrivilegeGroupRequest{
212+
GroupName: "group1",
213+
})
214+
s.NoError(err)
215+
s.True(merr.Ok(dropResp))
216+
217+
dropResp, err = s.Cluster.Proxy.DropPrivilegeGroup(ctx, &milvuspb.DropPrivilegeGroupRequest{
218+
GroupName: "",
219+
})
220+
s.NoError(err)
221+
s.False(merr.Ok(dropResp))
222+
223+
operateResp, err := s.Cluster.Proxy.OperatePrivilegeGroup(ctx, &milvuspb.OperatePrivilegeGroupRequest{
224+
GroupName: "",
225+
})
226+
s.NoError(err)
227+
s.False(merr.Ok(operateResp))
228+
229+
operateResp, err = s.Cluster.Proxy.OperatePrivilegeGroup(ctx, &milvuspb.OperatePrivilegeGroupRequest{
230+
GroupName: "group1",
231+
Privileges: []*milvuspb.PrivilegeEntity{{Name: "123"}},
232+
})
233+
s.NoError(err)
234+
s.False(merr.Ok(operateResp))
235+
}
236+
202237
func (s *PrivilegeGroupTestSuite) operatePrivilege(ctx context.Context, role, privilege, objectType string, operateType milvuspb.OperatePrivilegeType) {
203238
resp, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{
204239
Type: operateType,

0 commit comments

Comments
 (0)