Chore: Make a state entry to find total amount of native assets IBC'd out

modules/apps/transfer/client/cli/query.go

@@ -167,3 +167,36 @@ func GetCmdQueryDenomHash() *cobra.Command {
 	flags.AddQueryFlagsToCmd(cmd)
 	return cmd
 }
+
+// GetCmdQueryTotalEscrowForDenom defines the command to query the total amount of tokens in escrow for a native denom
+func GetCmdQueryTotalEscrowForDenom() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "token-escrow [denom]",
+		Short:   "Query the total amount of tokens of a native denom in escrow",
+		Long:    "Query the total amount of tokens of a native denom in escrow",
+		Example: fmt.Sprintf("%s query ibc-transfer token-escrow uosmo", version.AppName),
+		Args:    cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			clientCtx, err := client.GetClientQueryContext(cmd)
+			if err != nil {
+				return err
+			}
+
+			queryClient := types.NewQueryClient(clientCtx)
+
+			req := &types.QueryTotalEscrowFormDenomRequest{
+				Denom: args[0],
+			}
+
+			res, err := queryClient.TotalEscrowForDenom(cmd.Context(), req)
+			if err != nil {
+				return err
+			}
+
+			return clientCtx.PrintProto(res)
+		},
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+	return cmd
+}

modules/apps/transfer/keeper/grpc_query.go

@@ -121,3 +121,22 @@ func (q Keeper) EscrowAddress(c context.Context, req *types.QueryEscrowAddressRe
 		EscrowAddress: addr.String(),
 	}, nil
 }
+
+// TotalEscrowForDenom implements the TotalEscrowForDenom gRPC method.
+func (q Keeper) TotalEscrowForDenom(c context.Context, req *types.QueryTotalEscrowFormDenomRequest) (*types.QueryTotalEscrowForDenomResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+
+	ctx := sdk.UnwrapSDKContext(c)
+
+	if q.IsIBCDenom(ctx, req.Denom) {
+		return nil, status.Error(codes.InvalidArgument, "denom is not a native token denomination")
+	}
+
+	denomAmount := q.GetTotalEscrowForDenom(ctx, req.Denom)
+
+	return &types.QueryTotalEscrowForDenomResponse{
+		Amount: denomAmount.Int64(),
+	}, nil
+}

modules/apps/transfer/keeper/grpc_query_test.go

@@ -261,3 +261,80 @@ func (suite *KeeperTestSuite) TestEscrowAddress() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestTotalEscrowForDenom() {
+	var req *types.QueryTotalEscrowFormDenomRequest
+
+	testCases := []struct {
+		msg      string
+		malleate func()
+		expPass  bool
+	}{
+		{
+			"success",
+			func() {
+				req = &types.QueryTotalEscrowFormDenomRequest{
+					Denom: sdk.DefaultBondDenom,
+				}
+			},
+			true,
+		},
+		{
+			"not found denom trace",
+			func() {
+				denomTrace := types.DenomTrace{
+					Path:      "transfer/channel-0",
+					BaseDenom: sdk.DefaultBondDenom,
+				}
+
+				req = &types.QueryTotalEscrowFormDenomRequest{
+					Denom: denomTrace.IBCDenom(),
+				}
+			},
+			true, // consider the denom is of a native token
+		},
+		{
+			"invalid ibc denom",
+			func() {
+				req = &types.QueryTotalEscrowFormDenomRequest{
+					Denom: "ibc/𓃠🐾",
+				}
+			},
+			true, // consider the denom is of a native token
+		},
+		{
+			"non-native denom",
+			func() {
+				denomTrace := types.DenomTrace{
+					Path:      "transfer/channel-0",
+					BaseDenom: sdk.DefaultBondDenom,
+				}
+
+				suite.chainA.GetSimApp().TransferKeeper.SetDenomTrace(suite.chainA.GetContext(), denomTrace)
+
+				req = &types.QueryTotalEscrowFormDenomRequest{
+					Denom: denomTrace.IBCDenom(),
+				}
+			},
+			false,
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.msg), func() {
+			suite.SetupTest() // reset
+
+			tc.malleate()
+			ctx := sdk.WrapSDKContext(suite.chainA.GetContext())
+
+			res, err := suite.queryClient.TotalEscrowForDenom(ctx, req)
+
+			if tc.expPass {
+				suite.Require().NoError(err)
+				suite.Require().Equal(int64(0), res.Amount)
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}

modules/apps/transfer/keeper/keeper.go

@@ -1,6 +1,10 @@
 package keeper
 
 import (
+	"fmt"
+	"strings"
+
+	"cosmossdk.io/math"
 	tmbytes "github.com/cometbft/cometbft/libs/bytes"
 	"github.com/cometbft/cometbft/libs/log"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -151,3 +155,44 @@ func (k Keeper) AuthenticateCapability(ctx sdk.Context, cap *capabilitytypes.Cap
 func (k Keeper) ClaimCapability(ctx sdk.Context, cap *capabilitytypes.Capability, name string) error {
 	return k.scopedKeeper.ClaimCapability(ctx, cap, name)
 }
+
+// GetTotalEscrowForDenom gets the total amount of source chain tokens that are in escrow.
+func (k Keeper) GetTotalEscrowForDenom(ctx sdk.Context, denom string) math.Int {
+	store := ctx.KVStore(k.storeKey)
+	bz := store.Get(types.TotalEscrowForDenomKey(denom))
+	if bz == nil {
+		return math.ZeroInt()
+	}
+
+	var amount math.Int
+	if err := amount.Unmarshal(bz); err != nil {
+		panic(err)
+	}
+	return amount
+}
+
+// SetTotalEscrowForDenom stores the total amount of source chain tokens that are in escrow.
+func (k Keeper) SetTotalEscrowForDenom(ctx sdk.Context, denom string, amount math.Int) {
+	if amount.LT(math.ZeroInt()) {
+		panic(fmt.Sprintf("amount cannot be negative: %s", amount))
+	}
+
+	store := ctx.KVStore(k.storeKey)
+	bz, err := amount.Marshal()
+	if err != nil {
+		panic(err)
+	}
+
+	store.Set(types.TotalEscrowForDenomKey(denom), bz)
+}
+
+// IsIBCDenom returns true if the denomination is an known on-chain IBC denomination.
+func (k Keeper) IsIBCDenom(ctx sdk.Context, denom string) bool {
+	if strings.HasPrefix(denom, fmt.Sprintf("%s/", types.DenomPrefix)) {
+		_, err := k.DenomPathFromHash(ctx, denom)
+		if err == nil {
+			return true
+		}
+	}
+	return false
+}

modules/apps/transfer/keeper/migrations.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	"fmt"
 
+	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/cosmos/ibc-go/v7/modules/apps/transfer/types"
@@ -53,6 +54,40 @@ func (m Migrator) MigrateTraces(ctx sdk.Context) error {
 	return nil
 }
 
+// MigrateTotalEscrowForDenom migrates the total amount of source chain tokens in escrow.
+func (m Migrator) MigrateTotalEscrowForDenom(ctx sdk.Context) error {
+	nativeTokens := make(map[string]math.Int)
+
+	transferChannels := m.keeper.channelKeeper.GetAllChannelsWithPortPrefix(ctx, types.PortID)
+	for _, channel := range transferChannels {
+		escrowAddress := types.GetEscrowAddress(types.PortID, channel.ChannelId)
+		escrowBalances := m.keeper.bankKeeper.GetAllBalances(ctx, escrowAddress)
+
+		for _, escrowBalance := range escrowBalances {
+			// Denom possibilities:
+			// - "atom" = native denom
+			// - "ibc/<hash>" = non native denom
+			// - "atom/ibc/osmo" = native denom
+
+			if !m.keeper.IsIBCDenom(ctx, escrowBalance.Denom) {
+				if val, ok := nativeTokens[escrowBalance.Denom]; ok {
+					nativeTokens[escrowBalance.Denom] = val.Add(escrowBalance.Amount)
+				} else {
+					nativeTokens[escrowBalance.Denom] = escrowBalance.Amount
+				}
+			}
+		}
+	}
+
+	if len(nativeTokens) != 0 {
+		for denom, amount := range nativeTokens {
+			m.keeper.SetTotalEscrowForDenom(ctx, denom, amount)
+		}
+	}
+
+	return nil
+}
+
 func equalTraces(dtA, dtB types.DenomTrace) bool {
 	return dtA.BaseDenom == dtB.BaseDenom && dtA.Path == dtB.Path
 }

modules/apps/transfer/keeper/migrations_test.go

@@ -3,8 +3,12 @@ package keeper_test
 import (
 	"fmt"
 
+	"cosmossdk.io/math"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktestutil "github.com/cosmos/cosmos-sdk/x/bank/testutil"
 	transferkeeper "github.com/cosmos/ibc-go/v7/modules/apps/transfer/keeper"
 	transfertypes "github.com/cosmos/ibc-go/v7/modules/apps/transfer/types"
+	ibctesting "github.com/cosmos/ibc-go/v7/testing"
 )
 
 func (suite *KeeperTestSuite) TestMigratorMigrateTraces() {
@@ -119,3 +123,79 @@ func (suite *KeeperTestSuite) TestMigratorMigrateTracesCorruptionDetection() {
 		migrator.MigrateTraces(suite.chainA.GetContext()) //nolint:errcheck // we shouldn't check the error here because we want to ensure that a panic occurs.
 	})
 }
+
+func (suite *KeeperTestSuite) TestMigrateTotalEscrowForDenom() {
+	var path *ibctesting.Path
+
+	testCases := []struct {
+		msg               string
+		malleate          func()
+		expectedEscrowAmt math.Int
+	}{
+		{
+			"success: one native denom escrowed in one channel",
+			func() {
+				escrowAddress := transfertypes.GetEscrowAddress(path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID)
+				coin := sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100))
+
+				// funds the escrow account to have balance
+				suite.Require().NoError(banktestutil.FundAccount(suite.chainA.GetSimApp().BankKeeper, suite.chainA.GetContext(), escrowAddress, sdk.NewCoins(coin)))
+			},
+			math.NewInt(100),
+		},
+		{
+			"success: one native denom escrowed in two channels",
+			func() {
+				extraPath := NewTransferPath(suite.chainA, suite.chainB)
+				suite.coordinator.Setup(extraPath)
+
+				escrowAddress1 := transfertypes.GetEscrowAddress(path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID)
+				escrowAddress2 := transfertypes.GetEscrowAddress(extraPath.EndpointA.ChannelConfig.PortID, extraPath.EndpointA.ChannelID)
+				coin1 := sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100))
+				coin2 := sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100))
+
+				// funds the escrow accounts to have balance
+				suite.Require().NoError(banktestutil.FundAccount(suite.chainA.GetSimApp().BankKeeper, suite.chainA.GetContext(), escrowAddress1, sdk.NewCoins(coin1)))
+				suite.Require().NoError(banktestutil.FundAccount(suite.chainA.GetSimApp().BankKeeper, suite.chainA.GetContext(), escrowAddress2, sdk.NewCoins(coin2)))
+			},
+			math.NewInt(200),
+		},
+		{
+			"success: valid ibc denom escrowed in one channel",
+			func() {
+				escrowAddress := transfertypes.GetEscrowAddress(path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID)
+				trace := transfertypes.ParseDenomTrace(transfertypes.GetPrefixedDenom(path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID, sdk.DefaultBondDenom))
+				coin := sdk.NewCoin(trace.IBCDenom(), sdk.NewInt(100))
+
+				suite.chainA.GetSimApp().TransferKeeper.SetDenomTrace(suite.chainA.GetContext(), trace)
+
+				// funds the escrow accounts to have balance
+				suite.Require().NoError(banktestutil.FundAccount(suite.chainA.GetSimApp().BankKeeper, suite.chainA.GetContext(), escrowAddress, sdk.NewCoins(coin)))
+			},
+			math.NewInt(0),
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.msg), func() {
+			suite.SetupTest() // reset
+
+			path = NewTransferPath(suite.chainA, suite.chainB)
+			suite.coordinator.Setup(path)
+
+			tc.malleate() // explicitly fund escrow account
+
+			migrator := transferkeeper.NewMigrator(suite.chainA.GetSimApp().TransferKeeper)
+			suite.Require().NoError(migrator.MigrateTotalEscrowForDenom(suite.chainA.GetContext()))
+
+			// check that the migration set the expected amount for the native tokens
+			amount := suite.chainA.GetSimApp().TransferKeeper.GetTotalEscrowForDenom(suite.chainA.GetContext(), sdk.DefaultBondDenom)
+			suite.Require().Equal(tc.expectedEscrowAmt, amount)
+
+			// check that the migration did not set amount for non-native tokens
+			trace := transfertypes.ParseDenomTrace(transfertypes.GetPrefixedDenom(path.EndpointA.ChannelConfig.PortID, path.EndpointA.ChannelID, sdk.DefaultBondDenom))
+			amount = suite.chainA.GetSimApp().TransferKeeper.GetTotalEscrowForDenom(suite.chainA.GetContext(), trace.IBCDenom())
+			suite.Require().Equal(math.ZeroInt(), amount)
+		})
+	}
+}

modules/apps/transfer/keeper/relay.go

@@ -110,6 +110,13 @@ func (k Keeper) sendTransfer(
 		); err != nil {
 			return 0, err
 		}
+
+		// get the existing total amount in escrow
+		currentTotalEscrow := k.GetTotalEscrowForDenom(ctx, token.GetDenom())
+		newTotalEscrow := currentTotalEscrow.Add(token.Amount)
+
+		// store the new total amount in escrow
+		k.SetTotalEscrowForDenom(ctx, token.GetDenom(), newTotalEscrow)
 	} else {
 		labels = append(labels, telemetry.NewLabel(coretypes.LabelSource, "false"))
 
@@ -197,7 +204,6 @@ func (k Keeper) OnRecvPacket(ctx sdk.Context, packet channeltypes.Packet, data t
 	// NOTE: We use SourcePort and SourceChannel here, because the counterparty
 	// chain would have prefixed with DestPort and DestChannel when originally
 	// receiving this coin as seen in the "sender chain is the source" condition.
-
 	if types.ReceiverChainIsSource(packet.GetSourcePort(), packet.GetSourceChannel(), data.Denom) {
 		// sender chain is not the source, unescrow tokens
 
@@ -211,7 +217,7 @@ func (k Keeper) OnRecvPacket(ctx sdk.Context, packet channeltypes.Packet, data t
 		// The denomination used to send the coins is either the native denom or the hash of the path
 		// if the denomination is not native.
 		denomTrace := types.ParseDenomTrace(unprefixedDenom)
-		if denomTrace.Path != "" {
+		if !denomTrace.IsNativeDenom() {
 			denom = denomTrace.IBCDenom()
 		}
 		token := sdk.NewCoin(denom, transferAmount)
@@ -230,6 +236,13 @@ func (k Keeper) OnRecvPacket(ctx sdk.Context, packet channeltypes.Packet, data t
 			return errorsmod.Wrap(err, "unable to unescrow tokens, this may be caused by a malicious counterparty module or a bug: please open an issue on counterparty module")
 		}
 
+		// get the existing total amount in escrow
+		currentTotalEscrow := k.GetTotalEscrowForDenom(ctx, token.GetDenom())
+		newTotalEscrow := currentTotalEscrow.Sub(token.Amount)
+
+		// store the new total amount in escrow
+		k.SetTotalEscrowForDenom(ctx, token.GetDenom(), newTotalEscrow)
+
 		defer func() {
 			if transferAmount.IsInt64() {
 				telemetry.SetGaugeWithLabels(
@@ -366,6 +379,13 @@ func (k Keeper) refundPacketToken(ctx sdk.Context, packet channeltypes.Packet, d
 			return errorsmod.Wrap(err, "unable to unescrow tokens, this may be caused by a malicious counterparty module or a bug: please open an issue on counterparty module")
 		}
 
+		// get the existing total amount in escrow
+		currentTotalEscrow := k.GetTotalEscrowForDenom(ctx, token.GetDenom())
+		newTotalEscrow := currentTotalEscrow.Sub(token.Amount)
+
+		// store the new total amount in escrow
+		k.SetTotalEscrowForDenom(ctx, token.GetDenom(), newTotalEscrow)
+
 		return nil
 	}