!feat: implement secondary index for authz store

Signed-off-by: Artur Troian <troian@users.noreply.github.com>

Author: troian

simapp/go.mod

@@ -1,6 +1,8 @@
 module cosmossdk.io/simapp
 
-go 1.19
+go 1.21
+
+toolchain go1.24.4
 
 require (
 	cosmossdk.io/api v0.3.1

simapp/sim_test.go

@@ -203,6 +203,7 @@ func TestAppImportExport(t *testing.T) {
 				stakingtypes.HistoricalInfoKey, stakingtypes.UnbondingIDKey, stakingtypes.UnbondingIndexKey, stakingtypes.UnbondingTypeKey, stakingtypes.ValidatorUpdatesKey,
 			},
 		}, // ordering may change but it doesn't matter
+
 		{app.GetKey(slashingtypes.StoreKey), newApp.GetKey(slashingtypes.StoreKey), [][]byte{}},
 		{app.GetKey(minttypes.StoreKey), newApp.GetKey(minttypes.StoreKey), [][]byte{}},
 		{app.GetKey(distrtypes.StoreKey), newApp.GetKey(distrtypes.StoreKey), [][]byte{}},
@@ -211,7 +212,7 @@ func TestAppImportExport(t *testing.T) {
 		{app.GetKey(govtypes.StoreKey), newApp.GetKey(govtypes.StoreKey), [][]byte{}},
 		{app.GetKey(evidencetypes.StoreKey), newApp.GetKey(evidencetypes.StoreKey), [][]byte{}},
 		{app.GetKey(capabilitytypes.StoreKey), newApp.GetKey(capabilitytypes.StoreKey), [][]byte{}},
-		{app.GetKey(authzkeeper.StoreKey), newApp.GetKey(authzkeeper.StoreKey), [][]byte{authzkeeper.GrantKey, authzkeeper.GrantQueuePrefix}},
+		{app.GetKey(authzkeeper.StoreKey), newApp.GetKey(authzkeeper.StoreKey), [][]byte{authzkeeper.GrantKey, authzkeeper.GranteeKey, authzkeeper.GrantQueuePrefix}},
 	}
 
 	for _, skp := range storeKeysPrefixes {

tests/go.mod

@@ -1,6 +1,8 @@
 module github.com/cosmos/cosmos-sdk/tests
 
-go 1.19
+go 1.21
+
+toolchain go1.24.4
 
 require (
 	cosmossdk.io/api v0.3.1

x/authz/keeper/grpc_query.go

@@ -146,39 +146,64 @@ func (k Keeper) GranteeGrants(c context.Context, req *authz.QueryGranteeGrantsRe
 	}
 
 	ctx := sdk.UnwrapSDKContext(c)
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), GrantKey)
 
-	authorizations, pageRes, err := query.GenericFilteredPaginate(k.cdc, store, req.Pagination, func(key []byte, auth *authz.Grant) (*authz.GrantAuthorization, error) {
-		auth1, err := auth.GetAuthorization()
+	if req.Pagination == nil {
+		req.Pagination = &query.PageRequest{}
+	}
+
+	if req.Pagination.Limit == 0 {
+		req.Pagination.Limit = query.DefaultLimit
+	}
+
+	iter := sdk.KVStorePrefixIterator(ctx.KVStore(k.storeKey), granteeStoreKey(grantee, nil))
+	defer func() {
+		_ = iter.Close()
+	}()
+
+	var pageRes *query.PageResponse
+	var grants []*authz.GrantAuthorization
+
+	for ; iter.Valid(); iter.Next() {
+		grantee, granter := parseGranteeStoreKey(iter.Key())
+
+		var authorizations []*authz.GrantAuthorization
+
+		store := prefix.NewStore(ctx.KVStore(k.storeKey), grantStoreKey(grantee, granter, ""))
+		authorizations, pageRes, err = query.GenericFilteredPaginate(k.cdc, store, req.Pagination, func(key []byte, auth *authz.Grant) (*authz.GrantAuthorization, error) {
+			auth1, err := auth.GetAuthorization()
+			if err != nil {
+				return nil, err
+			}
+
+			authorizationAny, err := codectypes.NewAnyWithValue(auth1)
+			if err != nil {
+				return nil, status.Errorf(codes.Internal, err.Error())
+			}
+
+			return &authz.GrantAuthorization{
+				Authorization: authorizationAny,
+				Expiration:    auth.Expiration,
+				Granter:       granter.String(),
+				Grantee:       grantee.String(),
+			}, nil
+		}, func() *authz.Grant {
+			return &authz.Grant{}
+		})
+
 		if err != nil {
 			return nil, err
 		}
 
-		granter, g, _ := parseGrantStoreKey(append(GrantKey, key...))
-		if !g.Equals(grantee) {
-			return nil, nil
-		}
+		grants = append(grants, authorizations...)
 
-		authorizationAny, err := codectypes.NewAnyWithValue(auth1)
-		if err != nil {
-			return nil, status.Errorf(codes.Internal, err.Error())
+		req.Pagination.Limit -= uint64(len(authorizations))
+		if req.Pagination.Limit == 0 {
+			break
 		}
-
-		return &authz.GrantAuthorization{
-			Authorization: authorizationAny,
-			Expiration:    auth.Expiration,
-			Granter:       granter.String(),
-			Grantee:       grantee.String(),
-		}, nil
-	}, func() *authz.Grant {
-		return &authz.Grant{}
-	})
-	if err != nil {
-		return nil, err
 	}
 
 	return &authz.QueryGranteeGrantsResponse{
-		Grants:     authorizations,
+		Grants:     grants,
 		Pagination: pageRes,
 	}, nil
 }

x/authz/keeper/keeper.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"fmt"
+	"math/big"
 	"strconv"
 	"time"
 
@@ -202,6 +203,8 @@ func (k Keeper) SaveGrant(ctx sdk.Context, grantee, granter sdk.AccAddress, auth
 	bz := k.cdc.MustMarshal(&grant)
 	store.Set(skey, bz)
 
+	IncGranteeGrants(store, grantee, granter)
+
 	return ctx.EventManager().EmitTypedEvent(&authz.EventGrant{
 		MsgTypeUrl: authorization.MsgTypeURL(),
 		Granter:    granter.String(),
@@ -228,6 +231,8 @@ func (k Keeper) DeleteGrant(ctx sdk.Context, grantee sdk.AccAddress, granter sdk
 
 	store.Delete(skey)
 
+	decGranteeGrants(store, grantee, granter)
+
 	return ctx.EventManager().EmitTypedEvent(&authz.EventRevoke{
 		MsgTypeUrl: msgType,
 		Granter:    granter.String(),
@@ -240,7 +245,9 @@ func (k Keeper) GetAuthorizations(ctx sdk.Context, grantee sdk.AccAddress, grant
 	store := ctx.KVStore(k.storeKey)
 	key := grantStoreKey(grantee, granter, "")
 	iter := sdk.KVStorePrefixIterator(store, key)
-	defer iter.Close()
+	defer func() {
+		_ = iter.Close()
+	}()
 
 	var authorization authz.Grant
 	var authorizations []authz.Authorization
@@ -287,10 +294,12 @@ func (k Keeper) IterateGrants(ctx sdk.Context,
 	handler func(granterAddr sdk.AccAddress, granteeAddr sdk.AccAddress, grant authz.Grant) bool,
 ) {
 	store := ctx.KVStore(k.storeKey)
+
 	iter := sdk.KVStorePrefixIterator(store, GrantKey)
 	defer iter.Close()
 	for ; iter.Valid(); iter.Next() {
 		var grant authz.Grant
+
 		granterAddr, granteeAddr, _ := parseGrantStoreKey(iter.Key())
 		k.cdc.MustUnmarshal(iter.Value(), &grant)
 		if handler(granterAddr, granteeAddr, grant) {
@@ -409,3 +418,34 @@ func (k Keeper) DequeueAndDeleteExpiredGrants(ctx sdk.Context) error {
 
 	return nil
 }
+
+func IncGranteeGrants(store sdk.KVStore, grantee, granter sdk.AccAddress) {
+	key := granteeStoreKey(grantee, granter)
+
+	count := sdk.NewInt(0)
+
+	if store.Has(key) {
+		val := store.Get(key)
+		bi := new(big.Int).SetBytes(val).Int64()
+
+		count = count.AddRaw(bi + 1)
+
+	}
+
+	store.Set(key, count.BigInt().Bytes())
+}
+
+func decGranteeGrants(store sdk.KVStore, grantee, granter sdk.AccAddress) {
+	skey := granteeStoreKey(grantee, granter)
+	val := store.Get(skey)
+
+	bi := new(big.Int).SetBytes(val).Int64()
+	bi--
+
+	if bi == 0 {
+		store.Delete(skey)
+	} else {
+		count := sdk.NewInt(bi)
+		store.Set(skey, count.BigInt().Bytes())
+	}
+}

x/authz/keeper/keys.go

@@ -15,9 +15,11 @@ import (
 //
 // - 0x01<grant_Bytes>: Grant
 // - 0x02<grant_expiration_Bytes>: GrantQueueItem
+// - 0x03<grant_Bytes>: Grant by grantee
 var (
 	GrantKey         = []byte{0x01} // prefix for each key
 	GrantQueuePrefix = []byte{0x02}
+	GranteeKey       = []byte{0x03} // prefix for each key
 )
 
 var lenTime = len(sdk.FormatTimeBytes(time.Now()))
@@ -38,6 +40,19 @@ func grantStoreKey(grantee sdk.AccAddress, granter sdk.AccAddress, msgType strin
 	return key
 }
 
+// granteeStoreKey - return authorization store key
+// Items are stored with the following key: values
+//
+// - 0x03<granteeAddressLen (1 Byte)><granteeAddress_Bytes><granterAddressLen (1 Byte)><granterAddress_Bytes>
+func granteeStoreKey(grantee sdk.AccAddress, granter sdk.AccAddress) []byte {
+	grantee = address.MustLengthPrefix(grantee)
+	granter = address.MustLengthPrefix(granter)
+
+	key := sdk.AppendLengthPrefixedBytes(GranteeKey, grantee, granter)
+
+	return key
+}
+
 // parseGrantStoreKey - split granter, grantee address and msg type from the authorization key
 func parseGrantStoreKey(key []byte) (granterAddr, granteeAddr sdk.AccAddress, msgType string) {
 	// key is of format:
@@ -74,6 +89,27 @@ func parseGrantQueueKey(key []byte) (time.Time, sdk.AccAddress, sdk.AccAddress,
 	return exp, granter, grantee, nil
 }
 
+// parseGranteeStoreKey key is of format:
+// 0x03<granteeAddressLen (1 Byte)><granteeAddress_Bytes><granterAddressLen (1 Byte)><granterAddress_Bytes>
+func parseGranteeStoreKey(key []byte) (sdk.AccAddress, sdk.AccAddress) {
+	// key is of format:
+	// 0x03<granteeAddressLen (1 Byte)><granteeAddress_Bytes><granterAddressLen (1 Byte)><granterAddress_Bytes>
+	kv.AssertKeyAtLeastLength(key, 2)
+	key = key[1:] // remove prefix key
+	granteeAddrLen := key[0]
+	kv.AssertKeyAtLeastLength(key, int(granteeAddrLen))
+	key = key[1:]
+	granteeAddr := key[:granteeAddrLen]
+	kv.AssertKeyAtLeastLength(key, 1)
+	key = key[granteeAddrLen:]
+	granterAddrLen := int(key[0])
+	key = key[1:]
+	kv.AssertKeyLength(key, granterAddrLen)
+	granterAddr := key
+
+	return granteeAddr, granterAddr
+}
+
 // GrantQueueKey - return grant queue store key. If a given grant doesn't have a defined
 // expiration, then it should not be used in the pruning queue.
 // Key format is:
@@ -95,5 +131,5 @@ func GrantQueueTimePrefix(expiration time.Time) []byte {
 // firstAddressFromGrantStoreKey parses the first address only
 func firstAddressFromGrantStoreKey(key []byte) sdk.AccAddress {
 	addrLen := key[0]
-	return sdk.AccAddress(key[1 : 1+addrLen])
+	return key[1 : 1+addrLen]
 }

x/authz/keeper/keys_test.go

@@ -39,3 +39,13 @@ func TestGrantQueueKey(t *testing.T) {
 	require.Equal(t, granter, granter1)
 	require.Equal(t, grantee, grantee1)
 }
+
+func TestGranteeKey(t *testing.T) {
+	key := granteeStoreKey(grantee, granter)
+
+	require.Len(t, key, len(GranteeKey)+len(address.MustLengthPrefix(grantee))+len(address.MustLengthPrefix(granter)))
+
+	grantee1, granter1 := parseGranteeStoreKey(key)
+	require.Equal(t, granter, granter1)
+	require.Equal(t, grantee, grantee1)
+}