fix legacy query staking (#79)

* fix legacy query staking

* bet

* fix legacy params for validator delegation fallback

* Bet

* bet

---------

Co-authored-by: Khanh Hoa <hoa@notional.ventures>
Co-authored-by: Tuan Tran <tropicaldog17@gmail.com>

Author: 3 people

app/modules.go

@@ -156,7 +156,7 @@ func appModules(
 		mint.NewAppModule(appCodec, app.MintKeeper, app.AccountKeeper, nil, app.GetSubspace(minttypes.ModuleName)),
 		slashing.NewAppModule(appCodec, app.SlashingKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper, app.GetSubspace(slashingtypes.ModuleName), app.InterfaceRegistry()),
 		distr.NewAppModule(appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper, app.GetSubspace(distrtypes.ModuleName)),
-		customstaking.NewAppModule(appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper, app.ParamsKeeper, app.GetSubspace(stakingtypes.ModuleName)),
+		customstaking.NewAppModule(appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper, app.ParamsKeeper, app.GetSubspace(stakingtypes.ModuleName), app.GetKey(stakingtypes.StoreKey)),
 		upgrade.NewAppModule(app.UpgradeKeeper, app.AccountKeeper.AddressCodec()),
 		evidence.NewAppModule(app.EvidenceKeeper),
 		params.NewAppModule(app.ParamsKeeper),

custom/staking/module.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	storetypes "cosmossdk.io/store/types"
 	customtypes "github.com/classic-terra/core/v4/custom/staking/types"
 	core "github.com/classic-terra/core/v4/types"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -44,9 +45,11 @@ func (am AppModuleBasic) DefaultGenesis(cdc codec.JSONCodec) json.RawMessage {
 type AppModule struct {
 	staking.AppModule
 
+	cdc          codec.Codec
 	keeper       *keeper.Keeper
 	paramsKeeper paramskeeper.Keeper
 	ss           paramtypes.Subspace
+	storeKey     storetypes.StoreKey
 }
 
 // NewAppModule creates a new AppModule object
@@ -56,12 +59,15 @@ func NewAppModule(cdc codec.Codec,
 	bk stakingtypes.BankKeeper,
 	pk paramskeeper.Keeper,
 	ss paramtypes.Subspace,
+	storeKey storetypes.StoreKey,
 ) AppModule {
 	return AppModule{
 		AppModule:    staking.NewAppModule(cdc, keeper, ak, bk, ss),
+		cdc:          cdc,
 		keeper:       keeper,
 		paramsKeeper: pk,
 		ss:           ss,
+		storeKey:     storeKey,
 	}
 }
 
@@ -72,7 +78,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 	querier := keeper.Querier{Keeper: am.keeper}
 	stakingtypes.RegisterQueryServer(
 		cfg.QueryServer(),
-		NewLegacyQueryServer(querier, am.ss, am.keeper),
+		NewLegacyQueryServer(querier, am.ss, am.keeper, am.cdc, am.storeKey),
 	)
 
 	m := keeper.NewMigrator(am.keeper, am.ss)

custom/staking/query_server.go

@@ -2,14 +2,22 @@ package staking
 
 import (
 	"context"
+	"strconv"
+	"strings"
 
 	"cosmossdk.io/math"
+	"cosmossdk.io/store/prefix"
+	storetypes "cosmossdk.io/store/types"
 	legacytypes "github.com/classic-terra/core/v4/custom/staking/types"
 	legacyupgrade "github.com/classic-terra/core/v4/custom/upgrade/legacy"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/query"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	"github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // LegacyQueryServer wraps the staking QueryServer and sets legacy parameters for pre-upgrade height queries
@@ -18,18 +26,29 @@ type LegacyQueryServer struct {
 	stakingtypes.QueryServer
 	keeper         *keeper.Keeper
 	legacySubspace paramtypes.Subspace
+	cdc            codec.BinaryCodec
+	storeKey       storetypes.StoreKey
 }
 
-// NewLegacyQueryServer creates a new LegacyQueryServer instance
+// NewLegacyQueryServer creates a new LegacyQueryServer instance.
+//
+// `cdc` and `storeKey` are required for the pre-v5-staking-migration
+// ValidatorDelegations fallback path, which scans the primary DelegationKey
+// (0x31) prefix directly when the SDK's reverse-index (0x71) hasn't been
+// backfilled at the queried height.
 func NewLegacyQueryServer(
 	originalServer stakingtypes.QueryServer,
 	legacySubspace paramtypes.Subspace,
 	keeper *keeper.Keeper,
+	cdc codec.BinaryCodec,
+	storeKey storetypes.StoreKey,
 ) stakingtypes.QueryServer {
 	return &LegacyQueryServer{
 		QueryServer:    originalServer,
 		keeper:         keeper,
 		legacySubspace: legacySubspace,
+		cdc:            cdc,
+		storeKey:       storeKey,
 	}
 }
 
@@ -111,7 +130,92 @@ func (q *LegacyQueryServer) Validator(ctx context.Context, req *stakingtypes.Que
 }
 
 func (q *LegacyQueryServer) ValidatorDelegations(ctx context.Context, req *stakingtypes.QueryValidatorDelegationsRequest) (*stakingtypes.QueryValidatorDelegationsResponse, error) {
-	return q.QueryServer.ValidatorDelegations(q.ensureLegacyParams(ctx), req)
+	ensuredCtx := q.ensureLegacyParams(ctx)
+	sdkCtx := sdk.UnwrapSDKContext(ensuredCtx)
+	if legacyupgrade.IsPreStakingV5(sdkCtx.ChainID(), sdkCtx.BlockHeight()) {
+		return q.validatorDelegationsLegacy(sdkCtx, req)
+	}
+	return q.QueryServer.ValidatorDelegations(ensuredCtx, req)
+}
+
+// validatorDelegationsLegacy reproduces cosmos-sdk's unexported
+// `getValidatorDelegationsLegacy` (x/staking/keeper/grpc_query.go): it scans
+// the primary DelegationKey (0x31) prefix and filters by validator. Used for
+// archive queries at heights before the v4→v5 staking migration backfilled the
+// DelegationByValIndexKey (0x71) reverse-index that the SDK's default
+// ValidatorDelegations now relies on.
+func (q *LegacyQueryServer) validatorDelegationsLegacy(
+	ctx sdk.Context, req *stakingtypes.QueryValidatorDelegationsRequest,
+) (*stakingtypes.QueryValidatorDelegationsResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+	if req.ValidatorAddr == "" {
+		return nil, status.Error(codes.InvalidArgument, "validator address cannot be empty")
+	}
+	if _, err := sdk.ValAddressFromBech32(req.ValidatorAddr); err != nil {
+		return nil, status.Error(codes.InvalidArgument, err.Error())
+	}
+
+	store := ctx.KVStore(q.storeKey)
+	delStore := prefix.NewStore(store, stakingtypes.DelegationKey)
+
+	dels, pageRes, err := query.GenericFilteredPaginate(
+		q.cdc, delStore, req.Pagination,
+		func(_ []byte, d *stakingtypes.Delegation) (*stakingtypes.Delegation, error) {
+			if !strings.EqualFold(d.GetValidatorAddr(), req.ValidatorAddr) {
+				return nil, nil
+			}
+			return d, nil
+		},
+		func() *stakingtypes.Delegation { return &stakingtypes.Delegation{} },
+	)
+	if err != nil {
+		return nil, status.Error(codes.Internal, err.Error())
+	}
+
+	delegations := make(stakingtypes.Delegations, 0, len(dels))
+	for _, d := range dels {
+		delegations = append(delegations, *d)
+	}
+
+	delResps, err := q.delegationsToDelegationResponses(ctx, delegations)
+	if err != nil {
+		return nil, status.Error(codes.Internal, err.Error())
+	}
+
+	return &stakingtypes.QueryValidatorDelegationsResponse{
+		DelegationResponses: delResps,
+		Pagination:          pageRes,
+	}, nil
+}
+
+// delegationsToDelegationResponses mirrors the unexported helper of the same
+// name in cosmos-sdk's staking keeper: it looks up the validator for each
+// delegation and converts shares to bonded balance.
+func (q *LegacyQueryServer) delegationsToDelegationResponses(
+	ctx sdk.Context, delegations stakingtypes.Delegations,
+) (stakingtypes.DelegationResponses, error) {
+	bondDenom, err := q.keeper.BondDenom(ctx)
+	if err != nil {
+		return nil, err
+	}
+	resps := make(stakingtypes.DelegationResponses, 0, len(delegations))
+	for _, d := range delegations {
+		valAddr, err := sdk.ValAddressFromBech32(d.GetValidatorAddr())
+		if err != nil {
+			return nil, err
+		}
+		val, err := q.keeper.GetValidator(ctx, valAddr)
+		if err != nil {
+			return nil, err
+		}
+		balance := val.TokensFromShares(d.Shares).TruncateInt()
+		resps = append(resps, stakingtypes.NewDelegationResp(
+			d.GetDelegatorAddr(), d.GetValidatorAddr(), d.Shares, sdk.NewCoin(bondDenom, balance),
+		))
+	}
+	return resps, nil
 }
 
 func (q *LegacyQueryServer) ValidatorUnbondingDelegations(ctx context.Context, req *stakingtypes.QueryValidatorUnbondingDelegationsRequest) (*stakingtypes.QueryValidatorUnbondingDelegationsResponse, error) {
@@ -147,7 +251,44 @@ func (q *LegacyQueryServer) DelegatorValidator(ctx context.Context, req *staking
 }
 
 func (q *LegacyQueryServer) HistoricalInfo(ctx context.Context, req *stakingtypes.QueryHistoricalInfoRequest) (*stakingtypes.QueryHistoricalInfoResponse, error) {
-	return q.QueryServer.HistoricalInfo(q.ensureLegacyParams(ctx), req)
+	ensuredCtx := q.ensureLegacyParams(ctx)
+	sdkCtx := sdk.UnwrapSDKContext(ensuredCtx)
+	if legacyupgrade.IsPreStakingV5(sdkCtx.ChainID(), sdkCtx.BlockHeight()) {
+		return q.historicalInfoLegacy(sdkCtx, req)
+	}
+	return q.QueryServer.HistoricalInfo(ensuredCtx, req)
+}
+
+// historicalInfoLegacy reads HistoricalInfo using the pre-v5-staking-migration
+// key encoding: prefix 0x50 followed by the ASCII-decimal height string. The
+// v5 migration (cosmos-sdk@v0.53.6/x/staking/migrations/v5/store.go:39) re-keys
+// every entry to a big-endian uint64; before that migration ran (block
+// 28214400 on Columbus, 28917279 on Rebel-2) IAVL state contains only the old
+// string-format keys, so the SDK's GetHistoricalInfo — which constructs the
+// new binary key — misses and returns NotFound.
+func (q *LegacyQueryServer) historicalInfoLegacy(
+	ctx sdk.Context, req *stakingtypes.QueryHistoricalInfoRequest,
+) (*stakingtypes.QueryHistoricalInfoResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+	if req.Height < 0 {
+		return nil, status.Error(codes.InvalidArgument, "height cannot be negative")
+	}
+
+	store := ctx.KVStore(q.storeKey)
+	legacyKey := append([]byte{}, stakingtypes.HistoricalInfoKey...)
+	legacyKey = append(legacyKey, []byte(strconv.FormatInt(req.Height, 10))...)
+	bz := store.Get(legacyKey)
+	if bz == nil {
+		return nil, status.Errorf(codes.NotFound, "historical info for height %d not found", req.Height)
+	}
+
+	var hi stakingtypes.HistoricalInfo
+	if err := q.cdc.Unmarshal(bz, &hi); err != nil {
+		return nil, status.Error(codes.Internal, err.Error())
+	}
+	return &stakingtypes.QueryHistoricalInfoResponse{Hist: &hi}, nil
 }
 
 func (q *LegacyQueryServer) Pool(ctx context.Context, req *stakingtypes.QueryPoolRequest) (*stakingtypes.QueryPoolResponse, error) {