Fix

x/lsm/keeper/distribution.go

@@ -144,3 +144,72 @@ func (k Keeper) WithdrawTokenizeShareRecordReward(ctx context.Context, ownerAddr
 
 	return rewards, nil
 }
+
+// withdraw reward for all owning TokenizeShareRecord
+func (k Keeper) WithdrawAllTokenizeShareRecordReward(ctx sdk.Context, ownerAddr sdk.AccAddress) (sdk.Coins, error) {
+	// This check is necessary to prevent sending rewards to a blacklisted address
+	if k.bankKeeper.BlockedAddr(ownerAddr) {
+		return nil, errorsmod.Wrapf(sdkerrors.ErrUnauthorized, "%s is not allowed to receive external funds", ownerAddr)
+	}
+
+	totalRewards := sdk.Coins{}
+
+	records := k.GetTokenizeShareRecordsByOwner(ctx, ownerAddr)
+
+	for _, record := range records {
+		valAddr, err := k.stakingKeeper.ValidatorAddressCodec().StringToBytes(record.Validator)
+		if err != nil {
+			return nil, err
+		}
+
+		_, err = k.stakingKeeper.Validator(ctx, valAddr)
+		if err != nil && !goerrors.Is(err, stakingtypes.ErrNoValidatorFound) {
+			return nil, err
+		}
+
+		// if the error is ErrNoValidatorFound
+		if err != nil {
+			continue
+		}
+
+		_, err = k.stakingKeeper.Delegation(ctx, record.GetModuleAddress(), valAddr)
+		if err != nil && !goerrors.Is(err, stakingtypes.ErrNoDelegation) {
+			return nil, err
+		}
+
+		// if the error is ErrNoDelegation
+		if err != nil {
+			continue
+		}
+
+		// withdraw rewards into reward module account and send it to reward owner
+		cacheCtx, write := ctx.CacheContext()
+		_, err = k.distKeeper.WithdrawDelegationRewards(cacheCtx, record.GetModuleAddress(), valAddr)
+		if err != nil {
+			k.Logger(ctx).Error(err.Error())
+			continue
+		}
+
+		// apply changes when the module account has positive balance
+		balances := k.bankKeeper.GetAllBalances(cacheCtx, record.GetModuleAddress())
+		if !balances.Empty() {
+			err = k.bankKeeper.SendCoins(cacheCtx, record.GetModuleAddress(), ownerAddr, balances)
+			if err != nil {
+				k.Logger(ctx).Error(err.Error())
+				continue
+			}
+			write()
+			totalRewards = totalRewards.Add(balances...)
+		}
+	}
+
+	ctx.EventManager().EmitEvent(
+		sdk.NewEvent(
+			types.EventTypeWithdrawTokenizeShareReward,
+			sdk.NewAttribute(types.AttributeKeyWithdrawAddress, ownerAddr.String()),
+			sdk.NewAttribute(sdk.AttributeKeyAmount, totalRewards.String()),
+		),
+	)
+
+	return totalRewards, nil
+}

x/lsm/keeper/msg_server.go

@@ -499,3 +499,37 @@ func (k msgServer) EnableTokenizeShares(ctx context.Context, msg *types.MsgEnabl
 
 	return &types.MsgEnableTokenizeSharesResponse{CompletionTime: completionTime}, nil
 }
+
+// WithdrawTokenizeShareRecordReward defines a method to withdraw reward for owning TokenizeShareRecord
+func (k msgServer) WithdrawTokenizeShareRecordReward(goCtx context.Context, msg *types.MsgWithdrawTokenizeShareRecordReward) (*types.MsgWithdrawTokenizeShareRecordRewardResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	ownerAddr, err := k.authKeeper.AddressCodec().StringToBytes(msg.OwnerAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = k.Keeper.WithdrawTokenizeShareRecordReward(ctx, ownerAddr, msg.RecordId)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.MsgWithdrawTokenizeShareRecordRewardResponse{}, nil
+}
+
+// WithdrawAllTokenizeShareRecordReward defines a method to withdraw reward for owning TokenizeShareRecord
+func (k msgServer) WithdrawAllTokenizeShareRecordReward(goCtx context.Context, msg *types.MsgWithdrawAllTokenizeShareRecordReward) (*types.MsgWithdrawAllTokenizeShareRecordRewardResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	ownerAddr, err := k.authKeeper.AddressCodec().StringToBytes(msg.OwnerAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = k.Keeper.WithdrawAllTokenizeShareRecordReward(ctx, ownerAddr)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.MsgWithdrawAllTokenizeShareRecordRewardResponse{}, nil
+}