feat: root privileges can be customized (#39191)

- issue: #39184

Signed-off-by: SimFG <bang.fu@zilliz.com>

configs/milvus.yaml

@@ -830,6 +830,7 @@ common:
     # like the old password verification when updating the credential
     superUsers: 
     defaultRootPassword: "Milvus" # default password for root user. The maximum length is 72 characters, and double quotes are required.
+    rootShouldBindRole: false # Whether the root user should bind a role when the authorization is enabled.
     rbac:
       overrideBuiltInPrivilegeGroups:
         enabled: false # Whether to override build-in privilege groups

internal/distributed/proxy/httpserver/handler_v1_test.go

@@ -143,8 +143,11 @@ func initHTTPServer(proxy types.ProxyComponent, needAuth bool) *gin.Engine {
 ----|----------------|----------------|----------------
 */
 func genAuthMiddleWare(needAuth bool) gin.HandlerFunc {
+	InitMockGlobalMetaCache()
+	proxy.AddRootUserToAdminRole()
 	if needAuth {
 		return func(c *gin.Context) {
+			// proxy.RemoveRootUserFromAdminRole()
 			c.Set(ContextUsername, "")
 			username, password, ok := ParseUsernamePassword(c)
 			if !ok {
@@ -161,6 +164,10 @@ func genAuthMiddleWare(needAuth bool) gin.HandlerFunc {
 	}
 }
 
+func InitMockGlobalMetaCache() {
+	proxy.InitEmptyGlobalCache()
+}
+
 func Print(code int32, message string) string {
 	return fmt.Sprintf("{\"%s\":%d,\"%s\":\"%s\"}", HTTPReturnCode, code, HTTPReturnMessage, message)
 }

internal/distributed/proxy/httpserver/handler_v2_test.go

@@ -350,8 +350,8 @@ func TestGrpcWrapper(t *testing.T) {
 		assert.Equal(t, http.StatusForbidden, w.Code)
 		err = json.Unmarshal(w.Body.Bytes(), returnBody)
 		assert.Nil(t, err)
-		assert.Equal(t, int32(2), returnBody.Code)
-		assert.Equal(t, "service unavailable: internal: Milvus Proxy is not ready yet. please wait", returnBody.Message)
+		assert.Equal(t, int32(65535), returnBody.Code)
+		assert.Equal(t, "rpc error: code = PermissionDenied desc = PrivilegeLoad: permission deny to test in the `default` database", returnBody.Message)
 		fmt.Println(w.Body.String())
 	})
 }

internal/proxy/meta_cache_testonly.go

@@ -0,0 +1,63 @@
+//go:build test
+// +build test
+
+/*
+ * Licensed to the LF AI & Data foundation under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package proxy
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/mocks"
+	"github.com/milvus-io/milvus/pkg/common"
+	"github.com/milvus-io/milvus/pkg/util/funcutil"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+func AddRootUserToAdminRole() {
+	err := globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheAddUserToRole, OpKey: funcutil.EncodeUserRoleCache("root", "admin")})
+	if err != nil {
+		panic(err)
+	}
+}
+
+func RemoveRootUserFromAdminRole() {
+	err := globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheRemoveUserFromRole, OpKey: funcutil.EncodeUserRoleCache("root", "admin")})
+	if err != nil {
+		panic(err)
+	}
+}
+
+func InitEmptyGlobalCache() {
+	var err error
+	emptyMock := common.NewEmptyMockT()
+	rootcoord := mocks.NewMockRootCoordClient(emptyMock)
+	rootcoord.EXPECT().DescribeCollection(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("collection not found"))
+	querycoord := mocks.NewMockQueryCoordClient(emptyMock)
+	mgr := newShardClientMgr()
+	globalMetaCache, err = NewMetaCache(rootcoord, querycoord, mgr)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func SetGlobalMetaCache(metaCache *MetaCache) {
+	globalMetaCache = metaCache
+}

internal/proxy/privilege_interceptor.go

@@ -129,7 +129,7 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
 		log.Warn("GetCurUserFromContext fail", zap.Error(err))
 		return ctx, err
 	}
-	if username == util.UserRoot {
+	if !Params.CommonCfg.RootShouldBindRole.GetAsBool() && username == util.UserRoot {
 		return ctx, nil
 	}
 	roleNames, err := GetRole(username)

internal/proxy/privilege_interceptor_test.go

@@ -174,6 +174,40 @@ func TestPrivilegeInterceptor(t *testing.T) {
 	})
 }
 
+func TestRootShouldBindRole(t *testing.T) {
+	paramtable.Init()
+	Params.Save(Params.CommonCfg.AuthorizationEnabled.Key, "true")
+	defer Params.Reset(Params.CommonCfg.AuthorizationEnabled.Key)
+	rootCtx := GetContext(context.Background(), "root:Milvus")
+	t.Run("not bind role", func(t *testing.T) {
+		Params.Save(Params.CommonCfg.RootShouldBindRole.Key, "false")
+		defer Params.Reset(Params.CommonCfg.RootShouldBindRole.Key)
+
+		InitEmptyGlobalCache()
+		_, err := PrivilegeInterceptor(rootCtx, &milvuspb.LoadCollectionRequest{
+			CollectionName: "col1",
+		})
+		assert.NoError(t, err)
+	})
+
+	t.Run("bind role", func(t *testing.T) {
+		Params.Save(Params.CommonCfg.RootShouldBindRole.Key, "true")
+		defer Params.Reset(Params.CommonCfg.RootShouldBindRole.Key)
+
+		InitEmptyGlobalCache()
+		_, err := PrivilegeInterceptor(rootCtx, &milvuspb.LoadCollectionRequest{
+			CollectionName: "col1",
+		})
+		assert.Error(t, err)
+
+		AddRootUserToAdminRole()
+		_, err = PrivilegeInterceptor(rootCtx, &milvuspb.LoadCollectionRequest{
+			CollectionName: "col1",
+		})
+		assert.NoError(t, err)
+	})
+}
+
 func TestResourceGroupPrivilege(t *testing.T) {
 	ctx := context.Background()
 

internal/rootcoord/root_coord.go

@@ -578,7 +578,11 @@ func (c *Core) initCredentials(initCtx context.Context) error {
 		}
 		log.Ctx(initCtx).Info("RootCoord init user root")
 		err = c.meta.AddCredential(initCtx, &internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword})
-		return err
+		if err != nil {
+			log.Ctx(initCtx).Warn("RootCoord init user root failed", zap.Error(err))
+			return err
+		}
+		return nil
 	}
 	return nil
 }

pkg/common/mock_testonly.go

@@ -0,0 +1,52 @@
+//go:build test
+
+/*
+ * Licensed to the LF AI & Data foundation under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package common
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/milvus-io/milvus/pkg/log"
+)
+
+type MockTestingT struct {
+	ctx context.Context
+}
+
+func NewEmptyMockT() *MockTestingT {
+	return &MockTestingT{
+		ctx: context.Background(),
+	}
+}
+
+func (m *MockTestingT) Logf(format string, args ...interface{}) {
+	log.Ctx(m.ctx).Info(fmt.Sprintf(format, args...))
+}
+
+func (m *MockTestingT) Errorf(format string, args ...interface{}) {
+	log.Ctx(m.ctx).Error(fmt.Sprintf(format, args...))
+}
+
+func (m *MockTestingT) FailNow() {
+	log.Ctx(m.ctx).Panic("FailNow called")
+}
+
+func (m *MockTestingT) Cleanup(func()) {}

pkg/util/paramtable/component_param.go

@@ -239,6 +239,7 @@ type commonConfig struct {
 	AuthorizationEnabled ParamItem `refreshable:"false"`
 	SuperUsers           ParamItem `refreshable:"true"`
 	DefaultRootPassword  ParamItem `refreshable:"false"`
+	RootShouldBindRole   ParamItem `refreshable:"true"`
 
 	ClusterName ParamItem `refreshable:"false"`
 
@@ -669,6 +670,15 @@ like the old password verification when updating the credential`,
 	}
 	p.DefaultRootPassword.Init(base.mgr)
 
+	p.RootShouldBindRole = ParamItem{
+		Key:          "common.security.rootShouldBindRole",
+		Version:      "2.5.4",
+		Doc:          "Whether the root user should bind a role when the authorization is enabled.",
+		DefaultValue: "false",
+		Export:       true,
+	}
+	p.RootShouldBindRole.Init(base.mgr)
+
 	p.ClusterName = ParamItem{
 		Key:          "common.cluster.name",
 		Version:      "2.0.0",