!feat: implement secondary index for authz store

Signed-off-by: Artur Troian <artur@troian.io>

x/authz/keeper/grpc_query.go

@@ -15,7 +15,7 @@ import (
 
 var _ authz.QueryServer = Keeper{}
 
-// Authorizations implements the Query/Grants gRPC method.
+// Grants implements the Query/Grants gRPC method.
 func (k Keeper) Grants(c context.Context, req *authz.QueryGrantsRequest) (*authz.QueryGrantsResponse, error) {
 	if req == nil {
 		return nil, status.Errorf(codes.InvalidArgument, "empty request")
@@ -138,39 +138,59 @@ 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 := auth.GetAuthorization()
-		if err != nil {
-			return nil, err
-		}
+	if req.Pagination == nil {
+		req.Pagination = &query.PageRequest{}
+	}
 
-		granter, g := addressesFromGrantStoreKey(append(GrantKey, key...))
-		if !g.Equals(grantee) {
-			return nil, nil
-		}
+	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 := AddressesFromGranteeStoreKey(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 := 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{}
+		})
 
-		authorizationAny, err := codectypes.NewAnyWithValue(auth1)
 		if err != nil {
-			return nil, status.Errorf(codes.Internal, err.Error())
+			return nil, err
 		}
 
-		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
+		grants = append(grants, authorizations...)
 	}
 
 	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"
 
@@ -152,6 +153,9 @@ func (k Keeper) SaveGrant(ctx sdk.Context, grantee, granter sdk.AccAddress, auth
 	bz := k.cdc.MustMarshal(&grant)
 	skey := grantStoreKey(grantee, granter, authorization.MsgTypeURL())
 	store.Set(skey, bz)
+
+	IncGranteeGrants(store, grantee, granter)
+
 	return ctx.EventManager().EmitTypedEvent(&authz.EventGrant{
 		MsgTypeUrl: authorization.MsgTypeURL(),
 		Granter:    granter.String(),
@@ -168,7 +172,10 @@ func (k Keeper) DeleteGrant(ctx sdk.Context, grantee sdk.AccAddress, granter sdk
 	if !found {
 		return sdkerrors.ErrNotFound.Wrap("authorization not found")
 	}
+
 	store.Delete(skey)
+	decGranteeGrants(store, grantee, granter)
+
 	return ctx.EventManager().EmitTypedEvent(&authz.EventRevoke{
 		MsgTypeUrl: msgType,
 		Granter:    granter.String(),
@@ -181,12 +188,16 @@ 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
 	for ; iter.Valid(); iter.Next() {
 		k.cdc.MustUnmarshal(iter.Value(), &authorization)
 		authorizations = append(authorizations, authorization.GetAuthorization())
 	}
+
 	return authorizations
 }
 
@@ -199,7 +210,7 @@ func (k Keeper) GetCleanAuthorization(ctx sdk.Context, grantee sdk.AccAddress, g
 		return nil, time.Time{}
 	}
 	if grant.Expiration.Before(ctx.BlockHeader().Time) {
-		k.DeleteGrant(ctx, grantee, granter, msgType)
+		_ = k.DeleteGrant(ctx, grantee, granter, msgType)
 		return nil, time.Time{}
 	}
 
@@ -213,11 +224,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 := addressesFromGrantStoreKey(iter.Key())
+		granterAddr, granteeAddr := AddressesFromGrantStoreKey(iter.Key())
 		k.cdc.MustUnmarshal(iter.Value(), &grant)
 		if handler(granterAddr, granteeAddr, grant) {
 			break
@@ -258,3 +270,34 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data *authz.GenesisState) {
 		}
 	}
 }
+
+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

@@ -8,9 +8,11 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/authz"
 )
 
-// Keys for store prefixes
+// - 0x01<grant_Bytes>: Grant
+// - 0x03<grant_Bytes>: Grant
 var (
-	GrantKey = []byte{0x01} // prefix for each key
+	GrantKey   = []byte{0x01} // prefix for each key
+	GranteeKey = []byte{0x03} // prefix for each key
 )
 
 // StoreKey is the store key string for authz
@@ -31,27 +33,64 @@ func grantStoreKey(grantee sdk.AccAddress, granter sdk.AccAddress, msgType strin
 	copy(key[1:], granter)
 	copy(key[1+len(granter):], grantee)
 	copy(key[l-len(m):], m)
-	//	fmt.Println(">>>> len", l, key)
+
+	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)
+
+	l := 1 + len(grantee) + len(granter)
+	key := make([]byte, l)
+
+	copy(key, GranteeKey)
+	copy(key[1:], grantee)
+	copy(key[1+len(grantee):], granter)
+
 	return key
 }
 
-// addressesFromGrantStoreKey - split granter & grantee address from the authorization key
-func addressesFromGrantStoreKey(key []byte) (granterAddr, granteeAddr sdk.AccAddress) {
+// AddressesFromGrantStoreKey - split granter & grantee address from the authorization key
+func AddressesFromGrantStoreKey(key []byte) (granterAddr, granteeAddr sdk.AccAddress) {
 	// key is of format:
 	// 0x01<granterAddressLen (1 Byte)><granterAddress_Bytes><granteeAddressLen (1 Byte)><granteeAddress_Bytes><msgType_Bytes>
 	kv.AssertKeyAtLeastLength(key, 2)
 	granterAddrLen := key[1] // remove prefix key
 	kv.AssertKeyAtLeastLength(key, int(3+granterAddrLen))
-	granterAddr = sdk.AccAddress(key[2 : 2+granterAddrLen])
+	granterAddr = key[2 : 2+granterAddrLen]
 	granteeAddrLen := int(key[2+granterAddrLen])
 	kv.AssertKeyAtLeastLength(key, 4+int(granterAddrLen+byte(granteeAddrLen)))
-	granteeAddr = sdk.AccAddress(key[3+granterAddrLen : 3+granterAddrLen+byte(granteeAddrLen)])
+	granteeAddr = key[3+granterAddrLen : 3+granterAddrLen+byte(granteeAddrLen)]
+
+	return granterAddr, granteeAddr
+}
+
+func AddressesFromGranteeStoreKey(key []byte) (granteeAddr, granterAddr sdk.AccAddress) {
+	// key is of format:
+	// 0x03<granteeAddressLen (1 Byte)><granteeAddress_Bytes><granterAddressLen (1 Byte)><granterAddress_Bytes>
+	kv.AssertKeyAtLeastLength(key, 2)
+	key = key[1:]
+	granteeAddrLen := key[0] // remove prefix key
+	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)
+	granteeAddr = key
 
 	return granterAddr, granteeAddr
 }
 
 // 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/module/module.go

@@ -151,7 +151,7 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 }
 
 // ConsensusVersion implements AppModule/ConsensusVersion.
-func (AppModule) ConsensusVersion() uint64 { return 1 }
+func (AppModule) ConsensusVersion() uint64 { return 2 }
 
 func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {}