enhance: Add BackupRBAC/RestoreRBAC API to enable rbac backup (milvus…

…-io#35444)

issue: milvus-io#35443

---------

Signed-off-by: Wei Liu <[email protected]>

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/klauspost/compress v1.17.7
 	github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
-	github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632
+	github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59
 	github.com/minio/minio-go/v7 v7.0.61
 	github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81
 	github.com/prometheus/client_golang v1.14.0
@@ -57,6 +57,7 @@ require (
 
 require (
 	github.com/bits-and-blooms/bitset v1.10.0
+	github.com/bytedance/sonic v1.9.1
 	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/cockroachdb/redact v1.1.3
 	github.com/greatroar/blobloom v0.0.0-00010101000000-000000000000
@@ -91,7 +92,6 @@ require (
 	github.com/benbjohnson/clock v1.1.0 // indirect
 	github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/bytedance/sonic v1.9.1 // indirect
 	github.com/campoy/embedmd v1.0.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect

go.sum

@@ -598,8 +598,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119 h1:9VXijWu
 github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
 github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
 github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632 h1:CXig0DNtUsCLzchCFe3PR2KgOdobbz9gK2nSV7195PM=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59 h1:mKekr0GmCKMpIQh9OJ67TlKVKxDt08600ltARc/JUXY=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
 github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
 github.com/milvus-io/pulsar-client-go v0.6.10/go.mod h1:lQqCkgwDF8YFYjKA+zOheTk1tev2B+bKj5j7+nm8M1w=
 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=

internal/datacoord/mock_test.go

@@ -624,6 +624,14 @@ func (m *mockRootCoordClient) GetMetrics(ctx context.Context, req *milvuspb.GetM
 	}, nil
 }
 
+func (m *mockRootCoordClient) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
+	panic("not implemented") // TODO: Implement
+}
+
+func (m *mockRootCoordClient) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	panic("not implemented") // TODO: Implement
+}
+
 type mockCompactionTrigger struct {
 	methods map[string]interface{}
 }

internal/distributed/proxy/service.go

@@ -1145,6 +1145,14 @@ func (s *Server) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantReque
 	return s.proxy.SelectGrant(ctx, req)
 }
 
+func (s *Server) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
+	return s.proxy.BackupRBAC(ctx, req)
+}
+
+func (s *Server) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
+	return s.proxy.RestoreRBAC(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

@@ -682,3 +682,27 @@ func (c *Client) AlterDatabase(ctx context.Context, request *rootcoordpb.AlterDa
 		return client.AlterDatabase(ctx, request)
 	})
 }
+
+func (c *Client) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, 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.BackupRBACMetaResponse, error) {
+		return client.BackupRBAC(ctx, in)
+	})
+}
+
+func (c *Client) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, 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.RestoreRBAC(ctx, in)
+	})
+}

internal/distributed/rootcoord/service.go

@@ -533,3 +533,11 @@ func (s *Server) AlterCollection(ctx context.Context, request *milvuspb.AlterCol
 func (s *Server) RenameCollection(ctx context.Context, request *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
 	return s.rootCoord.RenameCollection(ctx, request)
 }
+
+func (s *Server) BackupRBAC(ctx context.Context, request *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
+	return s.rootCoord.BackupRBAC(ctx, request)
+}
+
+func (s *Server) RestoreRBAC(ctx context.Context, request *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
+	return s.rootCoord.RestoreRBAC(ctx, request)
+}

internal/metastore/catalog.go

@@ -82,6 +82,10 @@ type RootCoordCatalog interface {
 	// For example []string{"user1/role1"}
 	ListUserRole(ctx context.Context, tenant string) ([]string, error)
 
+	ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error)
+	BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
+	RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
+
 	Close()
 }
 

internal/metastore/kv/rootcoord/kv_catalog.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/cockroachdb/errors"
+	"github.com/samber/lo"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/proto"
 
@@ -740,23 +741,37 @@ func (kc *Catalog) ListAliases(ctx context.Context, dbID int64, ts typeutil.Time
 }
 
 func (kc *Catalog) ListCredentials(ctx context.Context) ([]string, error) {
-	keys, _, err := kc.Txn.LoadWithPrefix(CredentialPrefix)
+	users, err := kc.ListCredentialsWithPasswd(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return lo.Keys(users), nil
+}
+
+func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error) {
+	keys, values, err := kc.Txn.LoadWithPrefix(CredentialPrefix)
 	if err != nil {
 		log.Error("list all credential usernames fail", zap.String("prefix", CredentialPrefix), zap.Error(err))
 		return nil, err
 	}
 
-	var usernames []string
-	for _, path := range keys {
-		username := typeutil.After(path, UserSubPrefix+"/")
+	users := make(map[string]string)
+	for i := range keys {
+		username := typeutil.After(keys[i], UserSubPrefix+"/")
 		if len(username) == 0 {
-			log.Warn("no username extract from path:", zap.String("path", path))
+			log.Warn("no username extract from path:", zap.String("path", keys[i]))
 			continue
 		}
-		usernames = append(usernames, username)
+		credential := &internalpb.CredentialInfo{}
+		err := json.Unmarshal([]byte(values[i]), credential)
+		if err != nil {
+			log.Error("credential unmarshal fail", zap.String("key", keys[i]), zap.Error(err))
+			return nil, err
+		}
+		users[username] = credential.EncryptedPassword
 	}
 
-	return usernames, nil
+	return users, nil
 }
 
 func (kc *Catalog) save(k string) error {
@@ -1210,6 +1225,161 @@ func (kc *Catalog) ListUserRole(ctx context.Context, tenant string) ([]string, e
 	return userRoles, nil
 }
 
+func (kc *Catalog) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
+	users, err := kc.ListUser(ctx, tenant, nil, true)
+	if err != nil {
+		return nil, err
+	}
+
+	credentials, err := kc.ListCredentialsWithPasswd(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	userInfos := lo.FilterMap(users, func(entity *milvuspb.UserResult, _ int) (*milvuspb.UserInfo, bool) {
+		userName := entity.GetUser().GetName()
+		if userName == util.UserRoot {
+			return nil, false
+		}
+		return &milvuspb.UserInfo{
+			User:     userName,
+			Password: credentials[userName],
+			Roles:    entity.GetRoles(),
+		}, true
+	})
+
+	roles, err := kc.ListRole(ctx, tenant, nil, false)
+	if err != nil {
+		return nil, err
+	}
+
+	roleEntity := lo.FilterMap(roles, func(entity *milvuspb.RoleResult, _ int) (*milvuspb.RoleEntity, bool) {
+		roleName := entity.GetRole().GetName()
+		if roleName == util.RoleAdmin || roleName == util.RolePublic {
+			return nil, false
+		}
+
+		return entity.GetRole(), true
+	})
+
+	grantsEntity := make([]*milvuspb.GrantEntity, 0)
+	for _, role := range roleEntity {
+		grants, err := kc.ListGrant(ctx, tenant, &milvuspb.GrantEntity{
+			Role:   role,
+			DbName: util.AnyWord,
+		})
+		if err != nil {
+			return nil, err
+		}
+		grantsEntity = append(grantsEntity, grants...)
+	}
+
+	return &milvuspb.RBACMeta{
+		Users:  userInfos,
+		Roles:  roleEntity,
+		Grants: grantsEntity,
+	}, nil
+}
+
+func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
+	var err error
+	needRollbackUser := make([]*milvuspb.UserInfo, 0)
+	needRollbackRole := make([]*milvuspb.RoleEntity, 0)
+	needRollbackGrants := make([]*milvuspb.GrantEntity, 0)
+	defer func() {
+		if err != nil {
+			log.Warn("failed to restore rbac, try to rollback", zap.Error(err))
+			// roll back role
+			for _, role := range needRollbackRole {
+				err = kc.DropRole(ctx, tenant, role.Name)
+				if err != nil {
+					log.Warn("failed to rollback roles after restore failed", zap.Error(err))
+				}
+			}
+
+			// roll back grant
+			for _, grant := range needRollbackGrants {
+				err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Revoke)
+				if err != nil {
+					log.Warn("failed to rollback grants after restore failed", zap.Error(err))
+				}
+			}
+
+			for _, user := range needRollbackUser {
+				// roll back user
+				err = kc.DropCredential(ctx, user.User)
+				if err != nil {
+					log.Warn("failed to rollback users after restore failed", zap.Error(err))
+				}
+			}
+		}
+	}()
+
+	// restore role
+	existRoles, err := kc.ListRole(ctx, tenant, nil, false)
+	if err != nil {
+		return err
+	}
+	existRoleMap := lo.SliceToMap(existRoles, func(entity *milvuspb.RoleResult) (string, struct{}) { return entity.GetRole().GetName(), struct{}{} })
+	for _, role := range meta.Roles {
+		if _, ok := existRoleMap[role.GetName()]; ok {
+			log.Warn("failed to restore, role already exists", zap.String("role", role.GetName()))
+			err = errors.Newf("role [%s] already exists", role.GetName())
+			return err
+		}
+		err = kc.CreateRole(ctx, tenant, role)
+		if err != nil {
+			return err
+		}
+		needRollbackRole = append(needRollbackRole, role)
+	}
+
+	// restore grant
+	for _, grant := range meta.Grants {
+		err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Grant)
+		if err != nil {
+			return err
+		}
+		needRollbackGrants = append(needRollbackGrants, grant)
+	}
+
+	// need rollback user
+	existUser, err := kc.ListUser(ctx, tenant, nil, false)
+	if err != nil {
+		return err
+	}
+	existUserMap := lo.SliceToMap(existUser, func(entity *milvuspb.UserResult) (string, struct{}) { return entity.GetUser().GetName(), struct{}{} })
+	for _, user := range meta.Users {
+		if _, ok := existUserMap[user.GetUser()]; ok {
+			log.Info("failed to restore, user already exists", zap.String("user", user.GetUser()))
+			err = errors.Newf("user [%s] already exists", user.GetUser())
+			return err
+		}
+		// restore user
+		err = kc.CreateCredential(ctx, &model.Credential{
+			Username:          user.User,
+			EncryptedPassword: user.Password,
+		})
+		if err != nil {
+			return err
+		}
+		needRollbackUser = append(needRollbackUser, user)
+
+		// restore user role mapping
+		entity := &milvuspb.UserEntity{
+			Name: user.User,
+		}
+		for _, role := range user.Roles {
+			err = kc.AlterUserRole(ctx, tenant, entity, role, milvuspb.OperateUserRoleType_AddUserToRole)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return err
+}
+
 func (kc *Catalog) Close() {
 	// do nothing
 }