Refactor store keys for variable-length addresses

types/address.go

@@ -256,6 +256,21 @@ func (aa AccAddress) Format(s fmt.State, verb rune) {
 	}
 }
 
+// LengthPrefixAddress prefixes the address bytes with its length, this is used
+// for variable-length components in store keys. Note: All addresses should be
+// max 255 bytes, or else this function panics.
+func LengthPrefixAddress(bz []byte) []byte {
+	if len(bz) == 0 {
+		return bz
+	}
+
+	if len(bz) > 255 {
+		panic("address length should be max 255 bytes")
+	}
+
+	return append([]byte{byte(len(bz))}, bz...)
+}
+
 // ----------------------------------------------------------------------------
 // validator operator
 // ----------------------------------------------------------------------------

types/query/filtered_pagination_test.go

@@ -111,7 +111,7 @@ func ExampleFilteredPaginate() {
 	pageReq := &query.PageRequest{Key: nil, Limit: 1, CountTotal: true}
 	store := ctx.KVStore(app.GetKey(authtypes.StoreKey))
 	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr1.Bytes())
+	accountStore := prefix.NewStore(balancesStore, sdk.LengthPrefixAddress(addr1))
 
 	var balResult sdk.Coins
 	pageRes, err := query.FilteredPaginate(accountStore, pageReq, func(key []byte, value []byte, accumulate bool) (bool, error) {
@@ -143,7 +143,7 @@ func ExampleFilteredPaginate() {
 
 func execFilterPaginate(store sdk.KVStore, pageReq *query.PageRequest, appCodec codec.Marshaler) (balances sdk.Coins, res *query.PageResponse, err error) {
 	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr1.Bytes())
+	accountStore := prefix.NewStore(balancesStore, sdk.LengthPrefixAddress(addr1))
 
 	var balResult sdk.Coins
 	res, err = query.FilteredPaginate(accountStore, pageReq, func(key []byte, value []byte, accumulate bool) (bool, error) {

types/query/pagination_test.go

@@ -193,7 +193,7 @@ func ExamplePaginate() {
 	balResult := sdk.NewCoins()
 	authStore := ctx.KVStore(app.GetKey(authtypes.StoreKey))
 	balancesStore := prefix.NewStore(authStore, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr1.Bytes())
+	accountStore := prefix.NewStore(balancesStore, sdk.LengthPrefixAddress(addr1))
 	pageRes, err := query.Paginate(accountStore, request.Pagination, func(key []byte, value []byte) error {
 		var tempRes sdk.Coin
 		err := app.AppCodec().UnmarshalBinaryBare(value, &tempRes)

x/bank/keeper/grpc_query.go

@@ -57,9 +57,7 @@ func (k BaseKeeper) AllBalances(ctx context.Context, req *types.QueryAllBalances
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
 	balances := sdk.NewCoins()
-	store := sdkCtx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr.Bytes())
+	accountStore := k.getAccountStore(sdkCtx, addr)
 
 	pageRes, err := query.Paginate(accountStore, req.Pagination, func(_, value []byte) error {
 		var result sdk.Coin

x/bank/keeper/send.go

@@ -2,7 +2,6 @@ package keeper
 
 import (
 	"github.com/cosmos/cosmos-sdk/codec"
-	"github.com/cosmos/cosmos-sdk/store/prefix"
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -233,9 +232,7 @@ func (k BaseSendKeeper) ClearBalances(ctx sdk.Context, addr sdk.AccAddress) {
 		return false
 	})
 
-	store := ctx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr.Bytes())
+	accountStore := k.getAccountStore(ctx, addr)
 
 	for _, key := range keys {
 		accountStore.Delete(key)
@@ -264,9 +261,7 @@ func (k BaseSendKeeper) SetBalance(ctx sdk.Context, addr sdk.AccAddress, balance
 		return sdkerrors.Wrap(sdkerrors.ErrInvalidCoins, balance.String())
 	}
 
-	store := ctx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr.Bytes())
+	accountStore := k.getAccountStore(ctx, addr)
 
 	bz := k.cdc.MustMarshalBinaryBare(&balance)
 	accountStore.Set([]byte(balance.Denom), bz)

x/bank/keeper/view.go

@@ -97,9 +97,7 @@ func (k BaseViewKeeper) GetAccountsBalances(ctx sdk.Context) []types.Balance {
 // GetBalance returns the balance of a specific denomination for a given account
 // by address.
 func (k BaseViewKeeper) GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin {
-	store := ctx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr.Bytes())
+	accountStore := k.getAccountStore(ctx, addr)
 
 	bz := accountStore.Get([]byte(denom))
 	if bz == nil {
@@ -116,9 +114,7 @@ func (k BaseViewKeeper) GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom s
 // provides the token balance to a callback. If true is returned from the
 // callback, iteration is halted.
 func (k BaseViewKeeper) IterateAccountBalances(ctx sdk.Context, addr sdk.AccAddress, cb func(sdk.Coin) bool) {
-	store := ctx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
-	accountStore := prefix.NewStore(balancesStore, addr.Bytes())
+	accountStore := k.getAccountStore(ctx, addr)
 
 	iterator := accountStore.Iterator(nil, nil)
 	defer iterator.Close()
@@ -214,3 +210,11 @@ func (k BaseViewKeeper) ValidateBalance(ctx sdk.Context, addr sdk.AccAddress) er
 
 	return nil
 }
+
+// getAccountStore gets the account store of the given address.
+func (k BaseViewKeeper) getAccountStore(ctx sdk.Context, addr sdk.AccAddress) prefix.Store {
+	store := ctx.KVStore(k.storeKey)
+	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
+
+	return prefix.NewStore(balancesStore, sdk.LengthPrefixAddress(addr))
+}

x/bank/types/key.go

@@ -1,8 +1,6 @@
 package types
 
 import (
-	"fmt"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -22,7 +20,7 @@ const (
 
 // KVStore keys
 var (
-	BalancesPrefix      = []byte("balances")
+	BalancesPrefix      = []byte{0x02}
 	SupplyKey           = []byte{0x00}
 	DenomMetadataPrefix = []byte{0x1}
 )
@@ -37,10 +35,8 @@ func DenomMetadataKey(denom string) []byte {
 // store. The key must not contain the perfix BalancesPrefix as the prefix store
 // iterator discards the actual prefix.
 func AddressFromBalancesStore(key []byte) sdk.AccAddress {
-	addr := key[:sdk.AddrLen]
-	if len(addr) != sdk.AddrLen {
-		panic(fmt.Sprintf("unexpected account address key length; got: %d, expected: %d", len(addr), sdk.AddrLen))
-	}
+	addrLen := key[0]
+	addr := key[1 : addrLen+1]
 
 	return sdk.AccAddress(addr)
 }

x/bank/types/key_test.go

@@ -19,8 +19,10 @@ func cloneAppend(bz []byte, tail []byte) (res []byte) {
 func TestAddressFromBalancesStore(t *testing.T) {
 	addr, err := sdk.AccAddressFromBech32("cosmos1n88uc38xhjgxzw9nwre4ep2c8ga4fjxcar6mn7")
 	require.NoError(t, err)
+	addrLen := len(addr)
+	require.Equal(t, 20, addrLen)
 
-	key := cloneAppend(addr.Bytes(), []byte("stake"))
+	key := cloneAppend(sdk.LengthPrefixAddress(addr), []byte("stake"))
 	res := types.AddressFromBalancesStore(key)
 	require.Equal(t, res, addr)
 }

x/staking/keeper/val_state_change.go

@@ -12,7 +12,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
-// Calculate the ValidatorUpdates for the current block
+// BlockValidatorUpdates calculates the ValidatorUpdates for the current block
 // Called in each EndBlock
 func (k Keeper) BlockValidatorUpdates(ctx sdk.Context) []abci.ValidatorUpdate {
 	// Calculate validator set changes.
@@ -97,7 +97,7 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context) []abci.ValidatorUpdate {
 	return validatorUpdates
 }
 
-// Apply and return accumulated updates to the bonded validator set. Also,
+// ApplyAndReturnValidatorSetUpdates applies and return saccumulated updates to the bonded validator set. Also,
 // * Updates the active valset as keyed by LastValidatorPowerKey.
 // * Updates the total power as keyed by LastTotalPowerKey.
 // * Updates validator status' according to updated powers.
@@ -117,7 +117,10 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx sdk.Context) (updates []ab
 	// Retrieve the last validator set.
 	// The persistent set is updated later in this function.
 	// (see LastValidatorPowerKey).
-	last := k.getLastValidatorsByAddr(ctx)
+	last, err := k.getLastValidatorsByAddr(ctx)
+	if err != nil {
+		return nil, err
+	}
 
 	// Iterate over validators, highest power to lowest.
 	iterator := k.ValidatorsPowerStoreIterator(ctx)
@@ -160,10 +163,11 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx sdk.Context) (updates []ab
 		}
 
 		// fetch the old power bytes
-		var valAddrBytes [sdk.AddrLen]byte
-
-		copy(valAddrBytes[:], valAddr[:])
-		oldPowerBytes, found := last[valAddrBytes]
+		valAddrStr, err := sdk.Bech32ifyAddressBytes(sdk.Bech32PrefixValAddr, valAddr)
+		if err != nil {
+			return nil, err
+		}
+		oldPowerBytes, found := last[valAddrStr]
 		newPower := validator.ConsensusPower()
 		newPowerBytes := k.cdc.MustMarshalBinaryBare(&gogotypes.Int64Value{Value: newPower})
 
@@ -174,13 +178,16 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx sdk.Context) (updates []ab
 			k.SetLastValidatorPower(ctx, valAddr, newPower)
 		}
 
-		delete(last, valAddrBytes)
+		delete(last, valAddrStr)
 		count++
 
 		totalPower = totalPower.Add(sdk.NewInt(newPower))
 	}
 
-	noLongerBonded := sortNoLongerBonded(last)
+	noLongerBonded, err := sortNoLongerBonded(last)
+	if err != nil {
+		return nil, err
+	}
 	for _, valAddrBytes := range noLongerBonded {
 		validator := k.mustGetValidator(ctx, sdk.ValAddress(valAddrBytes))
 		validator, err = k.bondedToUnbonding(ctx, validator)
@@ -339,38 +346,47 @@ func (k Keeper) completeUnbondingValidator(ctx sdk.Context, validator types.Vali
 	return validator
 }
 
-// map of operator addresses to serialized power
-type validatorsByAddr map[[sdk.AddrLen]byte][]byte
+// map of operator bech32-addresses to serialized power
+// We use bech32 strings here, because we can't have slices as keys: map[[]byte][]byte
+type validatorsByAddr map[string][]byte
 
 // get the last validator set
-func (k Keeper) getLastValidatorsByAddr(ctx sdk.Context) validatorsByAddr {
+func (k Keeper) getLastValidatorsByAddr(ctx sdk.Context) (validatorsByAddr, error) {
 	last := make(validatorsByAddr)
 
 	iterator := k.LastValidatorsIterator(ctx)
 	defer iterator.Close()
 
 	for ; iterator.Valid(); iterator.Next() {
-		var valAddr [sdk.AddrLen]byte
-		// extract the validator address from the key (prefix is 1-byte)
-		copy(valAddr[:], iterator.Key()[1:])
+		// extract the validator address from the key (prefix is 1-byte, addrLen is 1-byte)
+		valAddr := iterator.Key()[2:]
+		valAddrStr, err := sdk.Bech32ifyAddressBytes(sdk.Bech32PrefixValAddr, valAddr)
+		if err != nil {
+			return nil, err
+		}
+
 		powerBytes := iterator.Value()
-		last[valAddr] = make([]byte, len(powerBytes))
-		copy(last[valAddr], powerBytes)
+		last[valAddrStr] = make([]byte, len(powerBytes))
+		copy(last[valAddrStr], powerBytes)
 	}
 
-	return last
+	return last, nil
 }
 
 // given a map of remaining validators to previous bonded power
 // returns the list of validators to be unbonded, sorted by operator address
-func sortNoLongerBonded(last validatorsByAddr) [][]byte {
+func sortNoLongerBonded(last validatorsByAddr) ([][]byte, error) {
 	// sort the map keys for determinism
 	noLongerBonded := make([][]byte, len(last))
 	index := 0
 
-	for valAddrBytes := range last {
-		valAddr := make([]byte, sdk.AddrLen)
-		copy(valAddr, valAddrBytes[:])
+	for valAddrStr := range last {
+		valAddrBytes, err := sdk.ValAddressFromBech32(valAddrStr)
+		if err != nil {
+			return nil, err
+		}
+		valAddr := make([]byte, len(valAddrBytes))
+		copy(valAddr, valAddrBytes)
 		noLongerBonded[index] = valAddr
 		index++
 	}
@@ -380,5 +396,5 @@ func sortNoLongerBonded(last validatorsByAddr) [][]byte {
 		return bytes.Compare(noLongerBonded[i], noLongerBonded[j]) == -1
 	})
 
-	return noLongerBonded
+	return noLongerBonded, nil
 }