fix legacy query staking

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,21 @@ package staking
 
 import (
 	"context"
+	"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 +25,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,9 +129,93 @@ func (q *LegacyQueryServer) Validator(ctx context.Context, req *stakingtypes.Que
 }
 
 func (q *LegacyQueryServer) ValidatorDelegations(ctx context.Context, req *stakingtypes.QueryValidatorDelegationsRequest) (*stakingtypes.QueryValidatorDelegationsResponse, error) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	if legacyupgrade.IsPreStakingV5(sdkCtx.ChainID(), sdkCtx.BlockHeight()) {
+		return q.validatorDelegationsLegacy(sdkCtx, req)
+	}
 	return q.QueryServer.ValidatorDelegations(q.ensureLegacyParams(ctx), 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) {
 	return q.QueryServer.ValidatorUnbondingDelegations(q.ensureLegacyParams(ctx), req)
 }

custom/staking/query_server_test.go

@@ -0,0 +1,176 @@
+package staking_test
+
+import (
+	"testing"
+
+	"cosmossdk.io/math"
+	storetypes "cosmossdk.io/store/types"
+	apptesting "github.com/classic-terra/core/v4/app/testing"
+	customstaking "github.com/classic-terra/core/v4/custom/staking"
+	"github.com/classic-terra/core/v4/types"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/testutil"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+	"github.com/stretchr/testify/suite"
+)
+
+type ValidatorDelegationsSuite struct {
+	apptesting.KeeperTestHelper
+}
+
+func TestValidatorDelegationsSuite(t *testing.T) {
+	suite.Run(t, new(ValidatorDelegationsSuite))
+}
+
+// seedValidatorWithDelegations creates `numVals` validators (so no single one
+// exceeds the 20% voting-power cap enforced by the custom staking hook),
+// then has `numDels` distinct delegators each delegate 1_000_000 uluna to the
+// FIRST validator. Returns that validator's address.
+func (s *ValidatorDelegationsSuite) seedValidatorWithDelegations(numVals, numDels int) sdk.ValAddress {
+	// Pre-fund the not-bonded-pool with the total self-stake so
+	// TestingUpdateValidator finds the tokens it expects.
+	valOwners := s.RandomAccountAddresses(numVals)
+	for _, o := range valOwners {
+		s.FundAcc(o, sdk.NewCoins(sdk.NewInt64Coin("uluna", 1_000_000)))
+		s.Require().NoError(s.App.BankKeeper.DelegateCoinsFromAccountToModule(
+			s.Ctx, o, stakingtypes.NotBondedPoolName,
+			sdk.NewCoins(sdk.NewInt64Coin("uluna", 1_000_000)),
+		))
+	}
+
+	valAddrs := simtestutil.ConvertAddrsToValAddrs(valOwners)
+	pks := simtestutil.CreateTestPubKeys(numVals)
+	vals := make([]stakingtypes.Validator, numVals)
+	for i := range vals {
+		v := testutil.NewValidator(s.T(), valAddrs[i], pks[i])
+		v, _ = v.AddTokensFromDel(math.NewInt(1_000_000))
+		v = stakingkeeper.TestingUpdateValidator(s.App.StakingKeeper, s.Ctx, v, true)
+		// Distribution rewards state is normally initialized in CreateValidator;
+		// TestingUpdateValidator skips that path, so do it manually.
+		s.Require().NoError(s.App.DistrKeeper.Hooks().AfterValidatorCreated(s.Ctx, valAddrs[i]))
+		vals[i] = v
+	}
+
+	// Delegators all stake to vals[0]. Each delegation is 1M, total stake
+	// across the chain ends up at numVals*1M + numDels*1M; the cap-hook
+	// requires vals[0] tokens / total <= 20%, so caller should pick numVals
+	// large enough to keep the ratio under the threshold.
+	addrDels := s.RandomAccountAddresses(numDels)
+	for _, d := range addrDels {
+		s.FundAcc(d, sdk.NewCoins(sdk.NewInt64Coin("uluna", 1_000_000)))
+		_, err := s.App.StakingKeeper.Delegate(s.Ctx, d, math.NewInt(1_000_000), stakingtypes.Unbonded, vals[0], true)
+		s.Require().NoError(err)
+	}
+
+	_, err := s.App.StakingKeeper.ApplyAndReturnValidatorSetUpdates(s.Ctx)
+	s.Require().NoError(err)
+
+	return valAddrs[0]
+}
+
+// dropReverseIndex deletes every entry under the staking module's
+// DelegationByValIndexKey (0x71) prefix.
+//
+// This simulates the IAVL state at heights *before* the cosmos-sdk staking
+// v4→v5 migration ran (the migration that backfills 0x71 from the primary
+// DelegationKey 0x31). Pre-migration archive state contains delegations under
+// 0x31 but nothing under 0x71 — which is what causes the empty query result
+// reported on the public archive LCDs at heights below 28214400.
+func (s *ValidatorDelegationsSuite) dropReverseIndex() {
+	storeKey := s.App.GetKey(stakingtypes.StoreKey)
+	store := s.Ctx.KVStore(storeKey)
+
+	iter := storetypes.KVStorePrefixIterator(store, stakingtypes.DelegationByValIndexKey)
+	defer iter.Close()
+
+	var keys [][]byte
+	for ; iter.Valid(); iter.Next() {
+		k := make([]byte, len(iter.Key()))
+		copy(k, iter.Key())
+		keys = append(keys, k)
+	}
+	for _, k := range keys {
+		store.Delete(k)
+	}
+}
+
+// TestValidatorDelegations_ReproducesArchiveBug reproduces the symptom seen on
+// public Terra Classic archive LCDs at pre-v5-staking-migration heights:
+// ValidatorDelegations returns an empty list because the SDK query iterates
+// over the 0x71 reverse-index, which has no entries in pre-migration IAVL
+// state.
+//
+// Pre-fix:  expect empty result (bug present).
+// Post-fix: expect populated result (fix routes to a primary-key scan when the
+//
+//	queried height is below MainnetStakingV5Height for Columbus).
+func (s *ValidatorDelegationsSuite) TestValidatorDelegations_ReproducesArchiveBug() {
+	s.Setup(s.T(), types.ColumbusChainID)
+
+	// 30 validators × 1M + 5 × 1M = 35M; vals[0] has 6M = 17.1% < 20% cap.
+	valAddr := s.seedValidatorWithDelegations(30, 5)
+
+	// Build the LegacyQueryServer the same way custom/staking/module.go does.
+	querier := stakingkeeper.Querier{Keeper: s.App.StakingKeeper}
+	ss := s.App.GetSubspace(stakingtypes.ModuleName)
+	qs := customstaking.NewLegacyQueryServer(
+		querier, ss, s.App.StakingKeeper,
+		s.App.AppCodec(), s.App.GetKey(stakingtypes.StoreKey),
+	)
+
+	req := &stakingtypes.QueryValidatorDelegationsRequest{ValidatorAddr: valAddr.String()}
+
+	// Use a query height above the v8 upgrade height so ensureLegacyParams
+	// takes the LegacyHandlingNone path and doesn't try to read non-existent
+	// legacy params from the subspace.
+	queryCtx := s.Ctx.WithBlockHeight(28214399)
+
+	// Sanity: with the reverse-index intact the query returns all 5 delegations.
+	resp, err := qs.ValidatorDelegations(queryCtx, req)
+	s.Require().NoError(err)
+	s.Require().Len(resp.DelegationResponses, 5, "sanity: index intact, should return all delegations")
+
+	// Simulate pre-migration archive state by wiping the 0x71 reverse-index.
+	s.dropReverseIndex()
+
+	resp, err = qs.ValidatorDelegations(queryCtx, req)
+	s.Require().NoError(err)
+
+	// THIS is the assertion that fails before the fix and passes after it.
+	s.Require().Len(
+		resp.DelegationResponses, 5,
+		"pre-migration height must still return delegations (regression of archive-LCD bug)",
+	)
+}
+
+// TestValidatorDelegations_PostMigrationUsesIndex ensures the fix doesn't change
+// behavior at chain-head heights: with the reverse-index intact and queried at
+// a post-v5-staking-migration height, the SDK's normal indexed path runs.
+// (We assert this indirectly by dropping the index at a post-migration height
+// and confirming the wrapper does NOT fall back to legacy iteration — i.e.
+// returns empty, just as the unwrapped SDK query would.)
+func (s *ValidatorDelegationsSuite) TestValidatorDelegations_PostMigrationUsesIndex() {
+	s.Setup(s.T(), types.ColumbusChainID)
+
+	valAddr := s.seedValidatorWithDelegations(30, 5)
+
+	querier := stakingkeeper.Querier{Keeper: s.App.StakingKeeper}
+	ss := s.App.GetSubspace(stakingtypes.ModuleName)
+	qs := customstaking.NewLegacyQueryServer(
+		querier, ss, s.App.StakingKeeper,
+		s.App.AppCodec(), s.App.GetKey(stakingtypes.StoreKey),
+	)
+
+	req := &stakingtypes.QueryValidatorDelegationsRequest{ValidatorAddr: valAddr.String()}
+	postCtx := s.Ctx.WithBlockHeight(28214400)
+
+	s.dropReverseIndex()
+	resp, err := qs.ValidatorDelegations(postCtx, req)
+	s.Require().NoError(err)
+	s.Require().Len(
+		resp.DelegationResponses, 0,
+		"at post-migration heights the legacy fallback must NOT trigger",
+	)
+}