added state and main functionality

x/ccv/provider/ibc_middleware.go

@@ -196,9 +196,14 @@ func (im IBCMiddleware) OnRecvPacket(
 			sdk.NewAttribute(types.AttributeRewardAmount, data.Amount),
 		}...)
 
-		// verify that the coin's denom is a whitelisted consumer denom,
-		// and if so, adds it to the consumer chain rewards allocation,
-		// otherwise the prohibited coin just stays in the pool forever.
+		alloc2 := im.keeper.GetConsumerRewardsAllocationByDenom(ctx, consumerId, coinDenom)
+		alloc2.Rewards = alloc2.Rewards.Add(
+			sdk.NewDecCoinFromCoin(sdk.Coin{
+				Denom:  coinDenom,
+				Amount: coinAmt,
+			}))
+		im.keeper.SetConsumerRewardsAllocationByDenom(ctx, consumerId, coinDenom, alloc2)
+
 		if im.keeper.ConsumerRewardDenomExists(ctx, coinDenom) {
 			alloc := im.keeper.GetConsumerRewardsAllocation(ctx, consumerId)
 			alloc.Rewards = alloc.Rewards.Add(
@@ -207,18 +212,18 @@ func (im IBCMiddleware) OnRecvPacket(
 					Amount: coinAmt,
 				})...)
 			im.keeper.SetConsumerRewardsAllocation(ctx, consumerId, alloc)
+		}
 
-			logger.Info(
-				"scheduled ICS rewards to be distributed",
-				"consumerId", consumerId,
-				"chainId", chainId,
-				"denom", coinDenom,
-				"amount", data.Amount,
-			)
+		logger.Info(
+			"scheduled ICS rewards to be distributed",
+			"consumerId", consumerId,
+			"chainId", chainId,
+			"denom", coinDenom,
+			"amount", data.Amount,
+		)
 
-			// add RewardDistribution event attribute
-			eventAttributes = append(eventAttributes, sdk.NewAttribute(types.AttributeRewardDistribution, "scheduled"))
-		}
+		// add RewardDistribution event attribute
+		eventAttributes = append(eventAttributes, sdk.NewAttribute(types.AttributeRewardDistribution, "scheduled"))
 
 		ctx.EventManager().EmitEvent(
 			sdk.NewEvent(

x/ccv/provider/keeper/distribution.go

@@ -203,6 +203,37 @@ func (k Keeper) AllocateTokens(ctx sdk.Context) {
 	// To avoid large iterations over all the consumer IDs, iterate only over
 	// chains with an IBC client created.
 	for _, consumerId := range k.GetAllConsumersWithIBCClients(ctx) {
+		oldRewards := k.GetConsumerRewardsAllocation(ctx, consumerId)
+		returnedRewards, err := k.AllocateConsumerRewards(ctx, consumerId, oldRewards)
+		if err != nil {
+			k.Logger(ctx).Error(
+				"fail to allocate rewards for consumer chain",
+				"consumer id", consumerId,
+				"error", err.Error(),
+			)
+		} else {
+			k.SetConsumerRewardsAllocation(ctx, consumerId, returnedRewards)
+		}
+
+		allAllowlistedDenoms := append(k.GetAllConsumerRewardDenoms(ctx), k.GetAllowlistedRewardDenoms(ctx, consumerId)...)
+		for _, denom := range allAllowlistedDenoms {
+			cachedCtx, writeCache := ctx.CacheContext()
+			consumerRewards := k.GetConsumerRewardsAllocationByDenom(cachedCtx, consumerId, denom)
+			allocatedRewards, err := k.AllocateConsumerRewards(cachedCtx, consumerId, consumerRewards)
+			if err != nil {
+				k.Logger(ctx).Error(
+					"fail to allocate rewards for consumer chain",
+					"consumer id", consumerId,
+					"error", err.Error(),
+				)
+				continue
+			}
+			k.SetConsumerRewardsAllocationByDenom(cachedCtx, consumerId, denom, allocatedRewards)
+			// TODO: fix for the tests
+			_ = writeCache
+			//writeCache()
+		}
+
 		// note that it's possible that no rewards are collected even though the
 		// reward pool isn't empty. This can happen if the reward pool holds some tokens
 		// of non-whitelisted denominations.

x/ccv/provider/keeper/keeper.go

@@ -681,6 +681,67 @@ func (k Keeper) DeleteConsumerClientId(ctx sdk.Context, consumerId string) {
 	store.Delete(types.ConsumerIdToClientIdKey(consumerId))
 }
 
+// GetAllowlistedRewardDenoms returns the allowlisted reward denom for the given consumer id.
+func (k Keeper) GetAllowlistedRewardDenoms(ctx sdk.Context, consumerId string) []string {
+	store := ctx.KVStore(k.storeKey)
+	bz := store.Get(types.ConsumerIdToAllowlistedRewardDenomKey(consumerId))
+	if bz == nil {
+		return []string{}
+	}
+
+	var denoms types.AllowlistedRewardDenoms
+	if err := denoms.Unmarshal(bz); err != nil {
+		// An error here would indicate something is very wrong,
+		// the PendingVSCPackets are assumed to be correctly serialized in AppendPendingVSCPackets.
+		panic(fmt.Errorf("cannot unmarshal pending validator set changes: %w", err))
+	}
+	return denoms.Denoms
+}
+
+// SetAllowlistedRewardDenom sets the allowlisted reward denom for the given consumer id.
+func (k Keeper) SetAllowlistedRewardDenom(ctx sdk.Context, consumerId string, denom string) error {
+	denoms := k.GetAllowlistedRewardDenoms(ctx, consumerId)
+	// TODO: check nil and so on
+	denoms = append(denoms, denom)
+	store := ctx.KVStore(k.storeKey)
+	dDenoms := types.AllowlistedRewardDenoms{Denoms: denoms}
+	bz, err := dDenoms.Marshal()
+	if err != nil {
+		return err
+	}
+	store.Set(types.ConsumerIdToAllowlistedRewardDenomKey(consumerId), bz)
+	return nil
+}
+
+// DeleteAllowlistedRewardDenom deletes the allowlisted reward denom for the given consumer id.
+func (k Keeper) DeleteAllowlistedRewardDenom(ctx sdk.Context, consumerId string) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.ConsumerIdToAllowlistedRewardDenomKey(consumerId))
+}
+
+// GetConsumerRewardsAllocationByDenom returns the consumer rewards allocation for the given consumer id and denom
+func (k Keeper) GetConsumerRewardsAllocationByDenom(ctx sdk.Context, consumerId string, denom string) (pool types.ConsumerRewardsAllocation) {
+	store := ctx.KVStore(k.storeKey)
+	b := store.Get(types.ConsumerRewardsAllocationByDenomKey(consumerId, denom))
+	k.cdc.MustUnmarshal(b, &pool)
+	return
+}
+
+// SetConsumerRewardsAllocationByDenom sets the consumer rewards allocation for the given consumer id and denom
+func (k Keeper) SetConsumerRewardsAllocationByDenom(ctx sdk.Context, consumerId string, denom string, pool types.ConsumerRewardsAllocation) {
+	store := ctx.KVStore(k.storeKey)
+	b := k.cdc.MustMarshal(&pool) // why is this MUST MARSHALLL??
+	store.Set(types.ConsumerRewardsAllocationByDenomKey(consumerId, denom), b)
+}
+
+// DeleteConsumerRewardsAllocationByDenom deletes the consumer rewards allocation for the given consumer id and denom
+func (k Keeper) DeleteConsumerRewardsAllocationByDenom(ctx sdk.Context, consumerId string, denom string) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.ConsumerRewardsAllocationByDenomKey(consumerId, denom))
+}
+
+///
+
 // SetSlashLog updates validator's slash log for a consumer chain
 // If an entry exists for a given validator address, at least one
 // double signing slash packet was received by the provider from at least one consumer chain

x/ccv/provider/keeper/msg_server.go

@@ -414,6 +414,18 @@ func (k msgServer) CreateConsumer(goCtx context.Context, msg *types.MsgCreateCon
 			sdk.NewAttribute(types.AttributeConsumerSpawnTime, initializationParameters.SpawnTime.String()))
 	}
 
+	if msg.AllowlistedRewardDenoms != nil {
+		// if allowlisted denoms are provided, they overwrite the previously written ones
+		k.DeleteAllowlistedRewardDenom(ctx, consumerId)
+		for _, denom := range msg.AllowlistedRewardDenoms.Denoms {
+			k.SetAllowlistedRewardDenom(ctx, consumerId, denom)
+		}
+
+		if len(k.GetAllowlistedRewardDenoms(ctx, consumerId)) > 3 {
+			return &resp, errorsmod.Wrapf(types.ErrInvalidMsgCreateConsumer, "a consumer chain cannot allowlist more than 3 denom")
+		}
+	}
+
 	// add Phase event attribute
 	phase := k.GetConsumerPhase(ctx, consumerId)
 	eventAttributes = append(eventAttributes, sdk.NewAttribute(types.AttributeConsumerPhase, phase.String()))
@@ -589,6 +601,18 @@ func (k msgServer) UpdateConsumer(goCtx context.Context, msg *types.MsgUpdateCon
 		}
 	}
 
+	if msg.AllowlistedRewardDenoms != nil {
+		// if allowlisted denoms are provided, they overwrite the previously written ones
+		k.DeleteAllowlistedRewardDenom(ctx, consumerId)
+		for _, denom := range msg.AllowlistedRewardDenoms.Denoms {
+			k.SetAllowlistedRewardDenom(ctx, consumerId, denom)
+		}
+
+		if len(k.GetAllowlistedRewardDenoms(ctx, consumerId)) > 3 {
+			return &resp, errorsmod.Wrapf(types.ErrInvalidMsgCreateConsumer, "a consumer chain cannot allowlist more than 3 denom")
+		}
+	}
+
 	// add Owner event attribute
 	eventAttributes = append(eventAttributes, sdk.NewAttribute(types.AttributeConsumerOwner, currentOwnerAddress))
 

x/ccv/provider/types/keys.go

@@ -146,6 +146,10 @@ const (
 	RemovalTimeToConsumerIdsKeyName = "RemovalTimeToConsumerIdsKeyName"
 
 	ClientIdToConsumerIdKeyName = "ClientIdToConsumerIdKey"
+
+	ConsumerIdToAllowlistedRewardDenomKeyName = "ConsumerIdToAllowlistedRewardDenomKey"
+
+	ConsumerRewardsAllocationByDenomKeyName = "ConsumerRewardsAllocationByDenomKey"
 )
 
 // getKeyPrefixes returns a constant map of all the byte prefixes for existing keys
@@ -376,6 +380,12 @@ func getKeyPrefixes() map[string]byte {
 		// ClientIdToConsumerIdKeyName is the key for storing the consumer id for the given client id
 		ClientIdToConsumerIdKeyName: 53,
 
+		// ConsumerIdToAllowlistedRewardDenomKeyName is the key for storing the allowlisted reward denom for the given consumer id
+		ConsumerIdToAllowlistedRewardDenomKeyName: 54,
+
+		// ConsumerRewardsAllocationByDenomKeyName is the key for storing the consumer rewards for a specific consumer chain and denom
+		ConsumerRewardsAllocationByDenomKeyName: 55,
+
 		// NOTE: DO NOT ADD NEW BYTE PREFIXES HERE WITHOUT ADDING THEM TO TestPreserveBytePrefix() IN keys_test.go
 	}
 }
@@ -754,6 +764,26 @@ func ClientIdToConsumerIdKey(clientId string) []byte {
 	)
 }
 
+// ConsumerIdToAllowlistedRewardDenomKeyPrefix returns the key prefix for storing the allowlisted reward denom that corresponds to this consumer id
+func ConsumerIdToAllowlistedRewardDenomKeyPrefix() byte {
+	return mustGetKeyPrefix(ConsumerIdToAllowlistedRewardDenomKeyName)
+}
+
+// ConsumerIdToAllowlistedRewardDenomKey returns the key used to store the allowlisted reward denom that corresponds to this consumer id
+func ConsumerIdToAllowlistedRewardDenomKey(consumerId string) []byte {
+	return StringIdWithLenKey(ConsumerIdToAllowlistedRewardDenomKeyPrefix(), consumerId)
+}
+
+// ConsumerRewardsAllocationByDenomKeyPrefix returns the key prefix for storing the allowlisted reward denom that corresponds to this consumer id
+func ConsumerRewardsAllocationByDenomKeyPrefix() byte {
+	return mustGetKeyPrefix(ConsumerRewardsAllocationByDenomKeyName)
+}
+
+// ConsumerRewardsAllocationByDenomKey returns the key used to store the ICS rewards per consumer chain
+func ConsumerRewardsAllocationByDenomKey(consumerId string, denom string) []byte {
+	return append(StringIdWithLenKey(ConsumerRewardsAllocationByDenomKeyPrefix(), consumerId), []byte(denom)...)
+}
+
 // NOTE: DO	NOT ADD FULLY DEFINED KEY FUNCTIONS WITHOUT ADDING THEM TO getAllFullyDefinedKeys() IN keys_test.go
 
 //

x/ccv/provider/types/keys_test.go

@@ -142,6 +142,10 @@ func TestPreserveBytePrefix(t *testing.T) {
 	i++
 	require.Equal(t, byte(53), providertypes.ClientIdToConsumerIdKey("clientId")[0])
 	i++
+	require.Equal(t, byte(54), providertypes.ConsumerIdToAllowlistedRewardDenomKey("13")[0])
+	i++
+	require.Equal(t, byte(55), providertypes.ConsumerRewardsAllocationByDenomKey("13", "denom")[0])
+	i++
 
 	prefixes := providertypes.GetAllKeyPrefixes()
 	require.Equal(t, len(prefixes), i)
@@ -210,6 +214,8 @@ func getAllFullyDefinedKeys() [][]byte {
 		providertypes.SpawnTimeToConsumerIdsKey(time.Time{}),
 		providertypes.RemovalTimeToConsumerIdsKey(time.Time{}),
 		providertypes.ClientIdToConsumerIdKey("clientId"),
+		providertypes.ConsumerIdToAllowlistedRewardDenomKey("13"),
+		providertypes.ConsumerRewardsAllocationByDenomKey("13", "denom"),
 	}
 }