From 1958598a186850c70604e2f33f61cfcb66daa94a Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Fri, 10 Jun 2016 10:59:34 -0700
Subject: [PATCH] auth: add getRole

---
 auth/range_perm_cache.go | 14 ++-------
 auth/store.go            | 67 ++++++++++++++++------------------------
 2 files changed, 30 insertions(+), 51 deletions(-)

diff --git a/auth/range_perm_cache.go b/auth/range_perm_cache.go
index 8ad39aa55..9d29b16b6 100644
--- a/auth/range_perm_cache.go
+++ b/auth/range_perm_cache.go
@@ -85,17 +85,9 @@ func (as *authStore) makeUnifiedPerms(tx backend.BatchTx, userName string) *unif
 	var readPerms, writePerms []*rangePerm
 
 	for _, roleName := range user.Roles {
-		_, vs := tx.UnsafeRange(authRolesBucketName, []byte(roleName), nil, 0)
-		if len(vs) != 1 {
-			plog.Errorf("invalid role name %s", roleName)
-			return nil
-		}
-
-		role := &authpb.Role{}
-		err := role.Unmarshal(vs[0])
-		if err != nil {
-			plog.Errorf("failed to unmarshal a role %s: %s", roleName, err)
-			return nil
+		role := getRole(tx, roleName)
+		if role == nil {
+			continue
 		}
 
 		for _, perm := range role.KeyPermission {
diff --git a/auth/store.go b/auth/store.go
index 2aa1ccbed..acb9ee585 100644
--- a/auth/store.go
+++ b/auth/store.go
@@ -400,17 +400,11 @@ func (as *authStore) RoleGet(r *pb.AuthRoleGetRequest) (*pb.AuthRoleGetResponse,
 	tx.Lock()
 	defer tx.Unlock()
 
-	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(r.Role), nil, 0)
-	if len(vs) != 1 {
+	role := getRole(tx, r.Role)
+	if role == nil {
 		return nil, ErrRoleNotFound
 	}
 
-	role := &authpb.Role{}
-	err := role.Unmarshal(vs[0])
-	if err != nil {
-		return nil, err
-	}
-
 	var resp pb.AuthRoleGetResponse
 	for _, perm := range role.KeyPermission {
 		resp.Perm = append(resp.Perm, perm)
@@ -424,17 +418,11 @@ func (as *authStore) RoleRevokePermission(r *pb.AuthRoleRevokePermissionRequest)
 	tx.Lock()
 	defer tx.Unlock()
 
-	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(r.Role), nil, 0)
-	if len(vs) != 1 {
+	role := getRole(tx, r.Role)
+	if role == nil {
 		return nil, ErrRoleNotFound
 	}
 
-	role := &authpb.Role{}
-	err := role.Unmarshal(vs[0])
-	if err != nil {
-		return nil, err
-	}
-
 	updatedRole := &authpb.Role{}
 	updatedRole.Name = role.Name
 
@@ -483,8 +471,8 @@ func (as *authStore) RoleDelete(r *pb.AuthRoleDeleteRequest) (*pb.AuthRoleDelete
 	tx.Lock()
 	defer tx.Unlock()
 
-	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(r.Role), nil, 0)
-	if len(vs) != 1 {
+	role := getRole(tx, r.Role)
+	if role == nil {
 		return nil, ErrRoleNotFound
 	}
 
@@ -499,8 +487,8 @@ func (as *authStore) RoleAdd(r *pb.AuthRoleAddRequest) (*pb.AuthRoleAddResponse,
 	tx.Lock()
 	defer tx.Unlock()
 
-	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(r.Name), nil, 0)
-	if len(vs) != 0 {
+	role := getRole(tx, r.Name)
+	if role != nil {
 		return nil, ErrRoleAlreadyExist
 	}
 
@@ -546,18 +534,11 @@ func (as *authStore) RoleGrantPermission(r *pb.AuthRoleGrantPermissionRequest) (
 	tx.Lock()
 	defer tx.Unlock()
 
-	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(r.Name), nil, 0)
-	if len(vs) != 1 {
+	role := getRole(tx, r.Name)
+	if role == nil {
 		return nil, ErrRoleNotFound
 	}
 
-	role := &authpb.Role{}
-	err := role.Unmarshal(vs[0])
-	if err != nil {
-		plog.Errorf("failed to unmarshal a role %s: %s", r.Name, err)
-		return nil, err
-	}
-
 	idx := sort.Search(len(role.KeyPermission), func(i int) bool {
 		return bytes.Compare(role.KeyPermission[i].Key, []byte(r.Perm.Key)) >= 0
 	})
@@ -612,17 +593,9 @@ func (as *authStore) isOpPermitted(userName string, key, rangeEnd string, write
 
 	if strings.Compare(rangeEnd, "") == 0 {
 		for _, roleName := range user.Roles {
-			_, vs := tx.UnsafeRange(authRolesBucketName, []byte(roleName), nil, 0)
-			if len(vs) != 1 {
-				plog.Errorf("invalid role name %s for permission checking", roleName)
-				return false
-			}
-
-			role := &authpb.Role{}
-			err := role.Unmarshal(vs[0])
-			if err != nil {
-				plog.Errorf("failed to unmarshal a role %s: %s", roleName, err)
-				return false
+			role := getRole(tx, roleName)
+			if role == nil {
+				continue
 			}
 
 			for _, perm := range role.KeyPermission {
@@ -691,6 +664,20 @@ func getUser(tx backend.BatchTx, username string) *authpb.User {
 	return user
 }
 
+func getRole(tx backend.BatchTx, rolename string) *authpb.Role {
+	_, vs := tx.UnsafeRange(authRolesBucketName, []byte(rolename), nil, 0)
+	if len(vs) == 0 {
+		return nil
+	}
+
+	role := &authpb.Role{}
+	err := role.Unmarshal(vs[0])
+	if err != nil {
+		plog.Panicf("failed to unmarshal role struct (name: %s): %s", rolename, err)
+	}
+	return role
+}
+
 func (as *authStore) isAuthEnabled() bool {
 	as.enabledMu.RLock()
 	defer as.enabledMu.RUnlock()