fix: RestoreRBAC pass wrong role info of user

Signed-off-by: Wei Liu <[email protected]>
milvus-io · Oct 30, 2024 · b5b5c77 · b5b5c77
1 parent c5e54a3
commit b5b5c77
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 8 deletions.
diff --git a/client/rbac.go b/client/rbac.go
@@ -465,9 +465,9 @@ func (c *GrpcClient) RestoreRBAC(ctx context.Context, meta *entity.RBACMeta) err
 	users := make([]*milvuspb.UserInfo, 0, len(meta.Users))
 	for _, user := range meta.Users {
 		roles := make([]*milvuspb.RoleEntity, 0, len(user.Roles))
-		for _, role := range meta.Roles {
+		for _, role := range user.Roles {
 			roles = append(roles, &milvuspb.RoleEntity{
-				Name: role.Name,
+				Name: role,
 			})
 		}
 		users = append(users, &milvuspb.UserInfo{

diff --git a/examples/rbac/main.go b/examples/rbac/main.go
@@ -3,9 +3,12 @@ package main
 import (
 	"context"
 	"log"
+	"strings"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/client"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/utils/crypto"
+	"google.golang.org/grpc/metadata"
 )
 
 const (
@@ -22,13 +25,20 @@ func main() {
 
 	log.Println("start connecting to Milvus")
 	c, err := client.NewClient(ctx, client.Config{
-		Address: milvusAddr,
+		Address:  milvusAddr,
+		Username: "root",
+		Password: "Milvus",
 	})
 	if err != nil {
 		log.Fatalf("failed to connect to milvus, err: %v", err)
 	}
 	defer c.Close()
 
+	// clean rbac
+	c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeCollection, "*", "*")
+	c.DeleteCredential(ctx, "user123")
+	c.DropRole(ctx, "role123")
+
 	// create user
 	err = c.CreateCredential(ctx, "user123", "passwd1")
 	if err != nil {
@@ -40,30 +50,95 @@ func main() {
 	if err != nil {
 		log.Fatalf("failed to create role, err: %v", err)
 	}
-	c.Grant(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+
+	grants, _ := c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
+
+	err = c.Grant(ctx, "role123", entity.PriviledegeObjectTypeCollection, "*", "Search")
+	if err != nil {
+		log.Fatalf("failed to grant role, err: %v", err)
+	}
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
 
 	// grant role to user
 	c.AddUserRole(ctx, "user123", "role123")
+	c.AddUserRole(ctx, "user123", "public")
+	c.AddUserRole(ctx, "user123", "admin")
 
 	// backup rbac
 	meta, err := c.BackupRBAC(ctx)
 	if err != nil {
 		log.Fatalf("failed to backup rbac, err: %v", err)
 	}
+	log.Println("user num: ", len(meta.Users))
+	for _, user := range meta.Users {
+		log.Println("user's role", user.Roles)
+	}
+	log.Println("role num: ", len(meta.Roles))
+	log.Println("grants num: ", len(meta.RoleGrants))
 
 	// clean rbac to make restore works
-	c.DropRole(ctx, "role123")
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
+	err = c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeCollection, "*", "Search")
+	if err != nil {
+		log.Fatalf("failed to revoke, err: %v", err)
+	}
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
 	c.DeleteCredential(ctx, "user123")
-	c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+	err = c.DropRole(ctx, "role123")
+	if err != nil {
+		log.Fatalf("failed to drop role, err: %v", err)
+	}
+
+	log.Println("-----start to restore rbac-----")
 
 	// restore rbac
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
+
 	err = c.RestoreRBAC(ctx, meta)
 	if err != nil {
 		log.Fatalf("failed to restore rbac, err: %v", err)
 	}
 
+	// backup rbac to check
+	log.Println("-----verify restore result-----")
+	meta, err = c.BackupRBAC(ctx)
+	if err != nil {
+		log.Fatalf("failed to backup rbac, err: %v", err)
+	}
+	log.Println("user num: ", len(meta.Users))
+	for _, user := range meta.Users {
+		log.Println("user's role", user.Roles)
+	}
+	log.Println("role num: ", len(meta.Roles))
+	log.Println("grants num: ", len(meta.RoleGrants))
+
 	// clean rbac
-	c.DropRole(ctx, "role123")
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
+	err = c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeCollection, "*", "Search")
+	if err != nil {
+		log.Fatalf("failed to revoke, err: %v", err)
+	}
+	grants, _ = c.ListGrants(ctx, "role123", "default")
+	log.Println("grants: ", len(grants))
 	c.DeleteCredential(ctx, "user123")
-	c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+	err = c.DropRole(ctx, "role123")
+	if err != nil {
+		log.Fatalf("failed to drop role, err: %v", err)
+	}
+}
+
+func GetContext(ctx context.Context, originValue string) context.Context {
+	authKey := strings.ToLower("authorization")
+	authValue := crypto.Base64Encode(originValue)
+	contextMap := map[string]string{
+		authKey: authValue,
+	}
+	md := metadata.New(contextMap)
+	return metadata.NewIncomingContext(ctx, md)
 }