fix: check store keys length before accessing (#9639)

* check store keys length before accessing

* check store keys length in all modules

* add changelog

* refactoring

* address review comments

* remove commented lines

* small fixes

* address review comments

* fix failing tests

* add godocs

CHANGELOG.md

@@ -75,6 +75,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Bug Fixes
 
+* [\#9639](https://github.com/cosmos/cosmos-sdk/pull/9639) Check store keys length before accessing them by making sure that `key` is of length `m+1` (for `key[n:m]`)
 * (types) [\#9627](https://github.com/cosmos/cosmos-sdk/pull/9627) Fix nil pointer panic on `NewBigIntFromInt`
 * (x/genutil) [\#9574](https://github.com/cosmos/cosmos-sdk/pull/9575) Actually use the `gentx` client tx flags (like `--keyring-dir`)
 * (x/distribution) [\#9599](https://github.com/cosmos/cosmos-sdk/pull/9599) Withdraw rewards event now includes a value attribute even if there are 0 rewards (due to situations like 100% commission).

types/kv/helpers.go

@@ -0,0 +1,17 @@
+package kv
+
+import "fmt"
+
+// AssertKeyAtLeastLength panics when store key length is less than the given length.
+func AssertKeyAtLeastLength(bz []byte, length int) {
+	if len(bz) < length {
+		panic(fmt.Sprintf("expected key of length at least %d, got %d", length, len(bz)))
+	}
+}
+
+// AssertKeyLength panics when store key length is not equal to the given length.
+func AssertKeyLength(bz []byte, length int) {
+	if len(bz) != length {
+		panic(fmt.Sprintf("unexpected key length; got: %d, expected: %d", len(bz), length))
+	}
+}

x/authz/keeper/keys.go

@@ -4,6 +4,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/internal/conv"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	"github.com/cosmos/cosmos-sdk/x/authz"
 )
 
@@ -38,9 +39,12 @@ func grantStoreKey(grantee sdk.AccAddress, granter sdk.AccAddress, msgType strin
 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])
 	granteeAddrLen := int(key[2+granterAddrLen])
+	kv.AssertKeyAtLeastLength(key, 4+int(granterAddrLen+byte(granteeAddrLen)))
 	granteeAddr = sdk.AccAddress(key[3+granterAddrLen : 3+granterAddrLen+byte(granteeAddrLen)])
 
 	return granterAddr, granteeAddr

x/bank/migrations/v040/keys.go

@@ -3,9 +3,8 @@
 package v040
 
 import (
-	"fmt"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/migrations/v040"
 )
 
@@ -40,10 +39,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 {
+	kv.AssertKeyAtLeastLength(key, 1+v040auth.AddrLen)
 	addr := key[:v040auth.AddrLen]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic(fmt.Sprintf("unexpected account address key length; got: %d, expected: %d", len(addr), v040auth.AddrLen))
-	}
-
 	return sdk.AccAddress(addr)
 }

x/bank/types/key.go

@@ -3,6 +3,7 @@ package types
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 )
 
 const (
@@ -43,6 +44,7 @@ func AddressFromBalancesStore(key []byte) (sdk.AccAddress, error) {
 	if len(key) == 0 {
 		return nil, ErrInvalidKey
 	}
+	kv.AssertKeyAtLeastLength(key, 1)
 	addrLen := key[0]
 	bound := int(addrLen)
 	if len(key)-1 < bound {

x/distribution/migrations/v040/keys.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/migrations/v040"
 )
 
@@ -58,78 +59,68 @@ var (
 
 // gets an address from a validator's outstanding rewards key
 func GetValidatorOutstandingRewardsAddress(key []byte) (valAddr sdk.ValAddress) {
+	kv.AssertKeyAtLeastLength(key, 2)
 	addr := key[1:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	return sdk.ValAddress(addr)
 }
 
 // gets an address from a delegator's withdraw info key
 func GetDelegatorWithdrawInfoAddress(key []byte) (delAddr sdk.AccAddress) {
+	kv.AssertKeyAtLeastLength(key, 2)
 	addr := key[1:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	return sdk.AccAddress(addr)
 }
 
 // gets the addresses from a delegator starting info key
 func GetDelegatorStartingInfoAddresses(key []byte) (valAddr sdk.ValAddress, delAddr sdk.AccAddress) {
+	kv.AssertKeyAtLeastLength(key, 2+v040auth.AddrLen)
 	addr := key[1 : 1+v040auth.AddrLen]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	valAddr = sdk.ValAddress(addr)
 	addr = key[1+v040auth.AddrLen:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	delAddr = sdk.AccAddress(addr)
 	return
 }
 
 // gets the address & period from a validator's historical rewards key
 func GetValidatorHistoricalRewardsAddressPeriod(key []byte) (valAddr sdk.ValAddress, period uint64) {
+	kv.AssertKeyAtLeastLength(key, 2+v040auth.AddrLen)
 	addr := key[1 : 1+v040auth.AddrLen]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	valAddr = sdk.ValAddress(addr)
 	b := key[1+v040auth.AddrLen:]
-	if len(b) != 8 {
+	kv.AssertKeyLength(addr, 8)
-		panic("unexpected key length")
-	}
 	period = binary.LittleEndian.Uint64(b)
 	return
 }
 
 // gets the address from a validator's current rewards key
 func GetValidatorCurrentRewardsAddress(key []byte) (valAddr sdk.ValAddress) {
+	kv.AssertKeyAtLeastLength(key, 2)
 	addr := key[1:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	return sdk.ValAddress(addr)
 }
 
 // gets the address from a validator's accumulated commission key
 func GetValidatorAccumulatedCommissionAddress(key []byte) (valAddr sdk.ValAddress) {
+	kv.AssertKeyAtLeastLength(key, 2)
 	addr := key[1:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	return sdk.ValAddress(addr)
 }
 
 // gets the height from a validator's slash event key
 func GetValidatorSlashEventAddressHeight(key []byte) (valAddr sdk.ValAddress, height uint64) {
+	kv.AssertKeyAtLeastLength(key, 2+v040auth.AddrLen)
 	addr := key[1 : 1+v040auth.AddrLen]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	valAddr = sdk.ValAddress(addr)
 	startB := 1 + v040auth.AddrLen
+	kv.AssertKeyAtLeastLength(key, startB+9)
 	b := key[startB : startB+8] // the next 8 bytes represent the height
 	height = binary.BigEndian.Uint64(b)
 	return

x/distribution/types/keys.go

@@ -5,6 +5,7 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 )
 
 const (
@@ -60,10 +61,9 @@ func GetValidatorOutstandingRewardsAddress(key []byte) (valAddr sdk.ValAddress)
 	// 0x02<valAddrLen (1 Byte)><valAddr_Bytes>
 
 	// Remove prefix and address length.
+	kv.AssertKeyAtLeastLength(key, 3)
 	addr := key[2:]
-	if len(addr) != int(key[1]) {
+	kv.AssertKeyLength(addr, int(key[1]))
-		panic("unexpected key length")
-	}
 
 	return sdk.ValAddress(addr)
 }
@@ -74,10 +74,9 @@ func GetDelegatorWithdrawInfoAddress(key []byte) (delAddr sdk.AccAddress) {
 	// 0x03<accAddrLen (1 Byte)><accAddr_Bytes>
 
 	// Remove prefix and address length.
+	kv.AssertKeyAtLeastLength(key, 3)
 	addr := key[2:]
-	if len(addr) != int(key[1]) {
+	kv.AssertKeyLength(addr, int(key[1]))
-		panic("unexpected key length")
-	}
 
 	return sdk.AccAddress(addr)
 }
@@ -86,13 +85,14 @@ func GetDelegatorWithdrawInfoAddress(key []byte) (delAddr sdk.AccAddress) {
 func GetDelegatorStartingInfoAddresses(key []byte) (valAddr sdk.ValAddress, delAddr sdk.AccAddress) {
 	// key is in the format:
 	// 0x04<valAddrLen (1 Byte)><valAddr_Bytes><accAddrLen (1 Byte)><accAddr_Bytes>
+	kv.AssertKeyAtLeastLength(key, 2)
 	valAddrLen := int(key[1])
+	kv.AssertKeyAtLeastLength(key, 3+valAddrLen)
 	valAddr = sdk.ValAddress(key[2 : 2+valAddrLen])
 	delAddrLen := int(key[2+valAddrLen])
+	kv.AssertKeyAtLeastLength(key, 4+valAddrLen)
 	delAddr = sdk.AccAddress(key[3+valAddrLen:])
-	if len(delAddr.Bytes()) != delAddrLen {
+	kv.AssertKeyLength(delAddr.Bytes(), delAddrLen)
-		panic("unexpected key length")
-	}
 
 	return
 }
@@ -101,12 +101,12 @@ func GetDelegatorStartingInfoAddresses(key []byte) (valAddr sdk.ValAddress, delA
 func GetValidatorHistoricalRewardsAddressPeriod(key []byte) (valAddr sdk.ValAddress, period uint64) {
 	// key is in the format:
 	// 0x05<valAddrLen (1 Byte)><valAddr_Bytes><period_Bytes>
+	kv.AssertKeyAtLeastLength(key, 2)
 	valAddrLen := int(key[1])
+	kv.AssertKeyAtLeastLength(key, 3+valAddrLen)
 	valAddr = sdk.ValAddress(key[2 : 2+valAddrLen])
 	b := key[2+valAddrLen:]
-	if len(b) != 8 {
+	kv.AssertKeyLength(b, 8)
-		panic("unexpected key length")
-	}
 	period = binary.LittleEndian.Uint64(b)
 	return
 }
@@ -117,10 +117,9 @@ func GetValidatorCurrentRewardsAddress(key []byte) (valAddr sdk.ValAddress) {
 	// 0x06<valAddrLen (1 Byte)><valAddr_Bytes>: ValidatorCurrentRewards
 
 	// Remove prefix and address length.
+	kv.AssertKeyAtLeastLength(key, 3)
 	addr := key[2:]
-	if len(addr) != int(key[1]) {
+	kv.AssertKeyLength(addr, int(key[1]))
-		panic("unexpected key length")
-	}
 
 	return sdk.ValAddress(addr)
 }
@@ -131,10 +130,9 @@ func GetValidatorAccumulatedCommissionAddress(key []byte) (valAddr sdk.ValAddres
 	// 0x07<valAddrLen (1 Byte)><valAddr_Bytes>: ValidatorCurrentRewards
 
 	// Remove prefix and address length.
+	kv.AssertKeyAtLeastLength(key, 3)
 	addr := key[2:]
-	if len(addr) != int(key[1]) {
+	kv.AssertKeyLength(addr, int(key[1]))
-		panic("unexpected key length")
-	}
 
 	return sdk.ValAddress(addr)
 }
@@ -143,9 +141,12 @@ func GetValidatorAccumulatedCommissionAddress(key []byte) (valAddr sdk.ValAddres
 func GetValidatorSlashEventAddressHeight(key []byte) (valAddr sdk.ValAddress, height uint64) {
 	// key is in the format:
 	// 0x08<valAddrLen (1 Byte)><valAddr_Bytes><height>: ValidatorSlashEvent
+	kv.AssertKeyAtLeastLength(key, 2)
 	valAddrLen := int(key[1])
+	kv.AssertKeyAtLeastLength(key, 3+valAddrLen)
 	valAddr = key[2 : 2+valAddrLen]
 	startB := 2 + valAddrLen
+	kv.AssertKeyAtLeastLength(key, startB+9)
 	b := key[startB : startB+8] // the next 8 bytes represent the height
 	height = binary.BigEndian.Uint64(b)
 	return

x/gov/migrations/v040/keys.go

@@ -4,10 +4,10 @@ package v040
 
 import (
 	"encoding/binary"
-	"fmt"
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/migrations/v040"
 )
 
@@ -113,9 +113,7 @@ func VoteKey(proposalID uint64, voterAddr sdk.AccAddress) []byte {
 
 // SplitProposalKey split the proposal key and returns the proposal id
 func SplitProposalKey(key []byte) (proposalID uint64) {
-	if len(key[1:]) != 8 {
+	kv.AssertKeyLength(key[1:], 8)
-		panic(fmt.Sprintf("unexpected key length (%d ≠ 8)", len(key[1:])))
-	}
 
 	return GetProposalIDFromBytes(key[1:])
 }
@@ -143,9 +141,7 @@ func SplitKeyVote(key []byte) (proposalID uint64, voterAddr sdk.AccAddress) {
 // private functions
 
 func splitKeyWithTime(key []byte) (proposalID uint64, endTime time.Time) {
-	if len(key[1:]) != 8+lenTime {
+	kv.AssertKeyLength(key[1:], 8+lenTime)
-		panic(fmt.Sprintf("unexpected key length (%d ≠ %d)", len(key[1:]), lenTime+8))
-	}
 
 	endTime, err := sdk.ParseTimeBytes(key[1 : 1+lenTime])
 	if err != nil {
@@ -157,10 +153,9 @@ func splitKeyWithTime(key []byte) (proposalID uint64, endTime time.Time) {
 }
 
 func splitKeyWithAddress(key []byte) (proposalID uint64, addr sdk.AccAddress) {
-	if len(key[1:]) != 8+v040auth.AddrLen {
+	kv.AssertKeyLength(key[1:], 8+v040auth.AddrLen)
-		panic(fmt.Sprintf("unexpected key length (%d ≠ %d)", len(key), 8+v040auth.AddrLen))
-	}
 
+	kv.AssertKeyAtLeastLength(key, 10)
 	proposalID = GetProposalIDFromBytes(key[1:9])
 	addr = sdk.AccAddress(key[9:])
 	return

x/gov/types/keys.go

@@ -2,11 +2,11 @@ package types
 
 import (
 	"encoding/binary"
-	"fmt"
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 )
 
 const (
@@ -111,9 +111,7 @@ func VoteKey(proposalID uint64, voterAddr sdk.AccAddress) []byte {
 
 // SplitProposalKey split the proposal key and returns the proposal id
 func SplitProposalKey(key []byte) (proposalID uint64) {
-	if len(key[1:]) != 8 {
+	kv.AssertKeyLength(key[1:], 8)
-		panic(fmt.Sprintf("unexpected key length (%d ≠ 8)", len(key[1:])))
-	}
 
 	return GetProposalIDFromBytes(key[1:])
 }
@@ -141,9 +139,7 @@ func SplitKeyVote(key []byte) (proposalID uint64, voterAddr sdk.AccAddress) {
 // private functions
 
 func splitKeyWithTime(key []byte) (proposalID uint64, endTime time.Time) {
-	if len(key[1:]) != 8+lenTime {
+	kv.AssertKeyLength(key[1:], 8+lenTime)
-		panic(fmt.Sprintf("unexpected key length (%d ≠ %d)", len(key[1:]), lenTime+8))
-	}
 
 	endTime, err := sdk.ParseTimeBytes(key[1 : 1+lenTime])
 	if err != nil {
@@ -157,7 +153,9 @@ func splitKeyWithTime(key []byte) (proposalID uint64, endTime time.Time) {
 func splitKeyWithAddress(key []byte) (proposalID uint64, addr sdk.AccAddress) {
 	// Both Vote and Deposit store keys are of format:
 	// <prefix (1 Byte)><proposalID (8 bytes)><addrLen (1 Byte)><addr_Bytes>
+	kv.AssertKeyAtLeastLength(key, 10)
 	proposalID = GetProposalIDFromBytes(key[1:9])
+	kv.AssertKeyAtLeastLength(key, 11)
 	addr = sdk.AccAddress(key[10:])
 	return
 }

x/slashing/migrations/v040/keys.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/migrations/v040"
 )
 
@@ -44,10 +45,9 @@ func ValidatorSigningInfoKey(v sdk.ConsAddress) []byte {
 
 // ValidatorSigningInfoAddress - extract the address from a validator signing info key
 func ValidatorSigningInfoAddress(key []byte) (v sdk.ConsAddress) {
+	kv.AssertKeyAtLeastLength(key, 2)
 	addr := key[1:]
-	if len(addr) != v040auth.AddrLen {
+	kv.AssertKeyLength(addr, v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 	return sdk.ConsAddress(addr)
 }
 

x/slashing/types/keys.go

@@ -5,6 +5,7 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 )
 
 const (
@@ -43,6 +44,7 @@ func ValidatorSigningInfoKey(v sdk.ConsAddress) []byte {
 // ValidatorSigningInfoAddress - extract the address from a validator signing info key
 func ValidatorSigningInfoAddress(key []byte) (v sdk.ConsAddress) {
 	// Remove prefix and address length.
+	kv.AssertKeyAtLeastLength(key, 3)
 	addr := key[2:]
 
 	return sdk.ConsAddress(addr)

x/staking/migrations/v040/keys.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 	v040auth "github.com/cosmos/cosmos-sdk/x/auth/migrations/v040"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -66,6 +67,7 @@ func GetValidatorByConsAddrKey(addr sdk.ConsAddress) []byte {
 
 // Get the validator operator address from LastValidatorPowerKey
 func AddressFromLastValidatorPowerKey(key []byte) []byte {
+	kv.AssertKeyAtLeastLength(key, 2)
 	return key[1:] // remove prefix bytes
 }
 
@@ -112,9 +114,7 @@ func GetLastValidatorPowerKey(operator sdk.ValAddress) []byte {
 // parse the validators operator address from power rank key
 func ParseValidatorPowerRankKey(key []byte) (operAddr []byte) {
 	powerBytesLen := 8
-	if len(key) != 1+powerBytesLen+v040auth.AddrLen {
+	kv.AssertKeyLength(key, 1+powerBytesLen+v040auth.AddrLen)
-		panic("Invalid validator power rank key length")
-	}
 
 	operAddr = sdk.CopyBytes(key[powerBytesLen+1:])
 
@@ -196,11 +196,11 @@ func GetUBDByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte
 
 // rearranges the ValIndexKey to get the UBDKey
 func GetUBDKeyFromValIndexKey(indexKey []byte) []byte {
+	kv.AssertKeyAtLeastLength(indexKey, 2)
 	addrs := indexKey[1:] // remove prefix bytes
-	if len(addrs) != 2*v040auth.AddrLen {
+	kv.AssertKeyLength(addrs, 2*v040auth.AddrLen)
-		panic("unexpected key length")
-	}
 
+	kv.AssertKeyAtLeastLength(addrs, v040auth.AddrLen+1)
 	valAddr := addrs[:v040auth.AddrLen]
 	delAddr := addrs[v040auth.AddrLen:]
 
@@ -268,9 +268,7 @@ func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.V
 // GetREDKeyFromValSrcIndexKey rearranges the ValSrcIndexKey to get the REDKey
 func GetREDKeyFromValSrcIndexKey(indexKey []byte) []byte {
 	// note that first byte is prefix byte
-	if len(indexKey) != 3*v040auth.AddrLen+1 {
+	kv.AssertKeyLength(indexKey, 3*v040auth.AddrLen+1)
-		panic("unexpected key length")
-	}
 
 	valSrcAddr := indexKey[1 : v040auth.AddrLen+1]
 	delAddr := indexKey[v040auth.AddrLen+1 : 2*v040auth.AddrLen+1]
@@ -282,9 +280,7 @@ func GetREDKeyFromValSrcIndexKey(indexKey []byte) []byte {
 // GetREDKeyFromValDstIndexKey rearranges the ValDstIndexKey to get the REDKey
 func GetREDKeyFromValDstIndexKey(indexKey []byte) []byte {
 	// note that first byte is prefix byte
-	if len(indexKey) != 3*v040auth.AddrLen+1 {
+	kv.AssertKeyLength(indexKey, 3*v040auth.AddrLen+1)
-		panic("unexpected key length")
-	}
 
 	valDstAddr := indexKey[1 : v040auth.AddrLen+1]
 	delAddr := indexKey[v040auth.AddrLen+1 : 2*v040auth.AddrLen+1]

x/staking/types/keys.go

@@ -9,6 +9,7 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/types/kv"
 )
 
 const (
@@ -63,11 +64,13 @@ func GetValidatorByConsAddrKey(addr sdk.ConsAddress) []byte {
 
 // AddressFromValidatorsKey creates the validator operator address from ValidatorsKey
 func AddressFromValidatorsKey(key []byte) []byte {
+	kv.AssertKeyAtLeastLength(key, 3)
 	return key[2:] // remove prefix bytes and address length
 }
 
 // AddressFromLastValidatorPowerKey creates the validator operator address from LastValidatorPowerKey
 func AddressFromLastValidatorPowerKey(key []byte) []byte {
+	kv.AssertKeyAtLeastLength(key, 3)
 	return key[2:] // remove prefix bytes and address length
 }
 
@@ -196,10 +199,13 @@ func GetUBDByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte
 
 // GetUBDKeyFromValIndexKey rearranges the ValIndexKey to get the UBDKey
 func GetUBDKeyFromValIndexKey(indexKey []byte) []byte {
+	kv.AssertKeyAtLeastLength(indexKey, 2)
 	addrs := indexKey[1:] // remove prefix bytes
 
 	valAddrLen := addrs[0]
+	kv.AssertKeyAtLeastLength(addrs, 2+int(valAddrLen))
 	valAddr := addrs[1 : 1+valAddrLen]
+	kv.AssertKeyAtLeastLength(addrs, 3+int(valAddrLen))
 	delAddr := addrs[valAddrLen+2:]
 
 	return GetUBDKey(delAddr, valAddr)
@@ -273,12 +279,16 @@ func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.V
 // GetREDKeyFromValSrcIndexKey rearranges the ValSrcIndexKey to get the REDKey
 func GetREDKeyFromValSrcIndexKey(indexKey []byte) []byte {
 	// note that first byte is prefix byte, which we remove
+	kv.AssertKeyAtLeastLength(indexKey, 2)
 	addrs := indexKey[1:]
 
 	valSrcAddrLen := addrs[0]
+	kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+2)
 	valSrcAddr := addrs[1 : valSrcAddrLen+1]
 	delAddrLen := addrs[valSrcAddrLen+1]
+	kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+int(delAddrLen)+2)
 	delAddr := addrs[valSrcAddrLen+2 : valSrcAddrLen+2+delAddrLen]
+	kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+int(delAddrLen)+4)
 	valDstAddr := addrs[valSrcAddrLen+delAddrLen+3:]
 
 	return GetREDKey(delAddr, valSrcAddr, valDstAddr)
@@ -287,12 +297,16 @@ func GetREDKeyFromValSrcIndexKey(indexKey []byte) []byte {
 // GetREDKeyFromValDstIndexKey rearranges the ValDstIndexKey to get the REDKey
 func GetREDKeyFromValDstIndexKey(indexKey []byte) []byte {
 	// note that first byte is prefix byte, which we remove
+	kv.AssertKeyAtLeastLength(indexKey, 2)
 	addrs := indexKey[1:]
 
 	valDstAddrLen := addrs[0]
+	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+2)
 	valDstAddr := addrs[1 : valDstAddrLen+1]
 	delAddrLen := addrs[valDstAddrLen+1]
+	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+int(delAddrLen)+3)
 	delAddr := addrs[valDstAddrLen+2 : valDstAddrLen+2+delAddrLen]
+	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+int(delAddrLen)+4)
 	valSrcAddr := addrs[valDstAddrLen+delAddrLen+3:]
 
 	return GetREDKey(delAddr, valSrcAddr, valDstAddr)