From 3a5550a938f97b1e49f436cc55e523403726c86f Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@orijtech.com>
Date: Wed, 7 Apr 2021 23:20:38 +0700
Subject: [PATCH] x/bank/types: fix AddressFromBalancesStore panics with
 invalid keys (#9061)

Currently, AddressFromBalancesStore uses the input key without any
validation, so an empty key or an invalid key length cause it panics.

This commit fixes the problem, by returning an error in case of invalid
key was passed.

Found by fuzzing added in #9060.

Fixed #9062
---
 CHANGELOG.md             |  1 +
 x/bank/keeper/view.go    |  8 +++++++-
 x/bank/types/errors.go   |  1 +
 x/bank/types/key.go      | 12 ++++++++++--
 x/bank/types/key_test.go | 25 ++++++++++++++++++++++++-
 5 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 661d3f2ddd15..ef35980ae30a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -80,6 +80,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (client) [\#8926](https://github.com/cosmos/cosmos-sdk/pull/8926) `client/tx.PrepareFactory` has been converted to a private function, as it's only used internally.
 * (auth/tx) [\#8926](https://github.com/cosmos/cosmos-sdk/pull/8926) The `ProtoTxProvider` interface used as a workaround for transaction simulation has been removed.
 * (x/bank) [\#8798](https://github.com/cosmos/cosmos-sdk/pull/8798) `GetTotalSupply` is removed in favour of `GetPaginatedTotalSupply`
+* (x/bank/types) [\#9061](https://github.com/cosmos/cosmos-sdk/pull/9061) `AddressFromBalancesStore` now returns an error for invalid key instead of panic.
 
 ### State Machine Breaking
 
diff --git a/x/bank/keeper/view.go b/x/bank/keeper/view.go
index 7271bba05945..d2a560cb7834 100644
--- a/x/bank/keeper/view.go
+++ b/x/bank/keeper/view.go
@@ -140,7 +140,13 @@ func (k BaseViewKeeper) IterateAllBalances(ctx sdk.Context, cb func(sdk.AccAddre
 	defer iterator.Close()
 
 	for ; iterator.Valid(); iterator.Next() {
-		address := types.AddressFromBalancesStore(iterator.Key())
+		address, err := types.AddressFromBalancesStore(iterator.Key())
+		if err != nil {
+			k.Logger(ctx).With("key", iterator.Key(), "err", err).Error("failed to get address from balances store")
+			// TODO: revisit, for now, panic here to keep same behavior as in 0.42
+			// ref: https://github.com/cosmos/cosmos-sdk/issues/7409
+			panic(err)
+		}
 
 		var balance sdk.Coin
 		k.cdc.MustUnmarshalBinaryBare(iterator.Value(), &balance)
diff --git a/x/bank/types/errors.go b/x/bank/types/errors.go
index ab88b40e9b1d..8446d957b678 100644
--- a/x/bank/types/errors.go
+++ b/x/bank/types/errors.go
@@ -11,4 +11,5 @@ var (
 	ErrInputOutputMismatch   = sdkerrors.Register(ModuleName, 4, "sum inputs != sum outputs")
 	ErrSendDisabled          = sdkerrors.Register(ModuleName, 5, "send transactions are disabled")
 	ErrDenomMetadataNotFound = sdkerrors.Register(ModuleName, 6, "client denom metadata not found")
+	ErrInvalidKey            = sdkerrors.Register(ModuleName, 7, "invalid key")
 )
diff --git a/x/bank/types/key.go b/x/bank/types/key.go
index 858fd2480ac9..0d8ec96a0de3 100644
--- a/x/bank/types/key.go
+++ b/x/bank/types/key.go
@@ -37,11 +37,19 @@ func DenomMetadataKey(denom string) []byte {
 // AddressFromBalancesStore returns an account address from a balances prefix
 // store. The key must not contain the perfix BalancesPrefix as the prefix store
 // iterator discards the actual prefix.
-func AddressFromBalancesStore(key []byte) sdk.AccAddress {
+//
+// If invalid key is passed, AddressFromBalancesStore returns ErrInvalidKey.
+func AddressFromBalancesStore(key []byte) (sdk.AccAddress, error) {
+	if len(key) == 0 {
+		return nil, ErrInvalidKey
+	}
 	addrLen := key[0]
+	if len(key[1:]) < int(addrLen) {
+		return nil, ErrInvalidKey
+	}
 	addr := key[1 : addrLen+1]
 
-	return sdk.AccAddress(addr)
+	return sdk.AccAddress(addr), nil
 }
 
 // CreateAccountBalancesPrefix creates the prefix for an account's balances.
diff --git a/x/bank/types/key_test.go b/x/bank/types/key_test.go
index 206ef8335c91..c54037a2279e 100644
--- a/x/bank/types/key_test.go
+++ b/x/bank/types/key_test.go
@@ -1,8 +1,10 @@
 package types_test
 
 import (
+	"errors"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -24,6 +26,27 @@ func TestAddressFromBalancesStore(t *testing.T) {
 	require.Equal(t, 20, addrLen)
 
 	key := cloneAppend(address.MustLengthPrefix(addr), []byte("stake"))
-	res := types.AddressFromBalancesStore(key)
+	res, err := types.AddressFromBalancesStore(key)
+	require.NoError(t, err)
 	require.Equal(t, res, addr)
 }
+
+func TestInvalidAddressFromBalancesStore(t *testing.T) {
+	tests := []struct {
+		name string
+		key  []byte
+	}{
+		{"empty", []byte("")},
+		{"invalid", []byte("3AA")},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			_, err := types.AddressFromBalancesStore(tc.key)
+			assert.Error(t, err)
+			assert.True(t, errors.Is(types.ErrInvalidKey, err))
+		})
+	}
+}