// SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause
/*
 * KUnit tests for channel helper functions
 *
 * Copyright (C) 2024-2025 Intel Corporation
 */
#include <kunit/test.h>
#include <kunit/static_stub.h>
#include <kunit/skbuff.h>

#include "utils.h"
#include "mld.h"
#include "sta.h"
#include "agg.h"
#include "rx.h"

#define FC_QOS_DATA (IEEE80211_FTYPE_DATA | IEEE80211_STYPE_QOS_DATA)
#define BA_WINDOW_SIZE 64
#define QUEUE 0

static const struct reorder_buffer_case {
	const char *desc;
	struct {
		/* ieee80211_hdr fields */
		u16 fc;
		u8 tid;
		bool multicast;
		/* iwl_rx_mpdu_desc fields */
		u16 nssn;
		/* used also for setting hdr::seq_ctrl */
		u16 sn;
		u8 baid;
		bool amsdu;
		bool last_subframe;
		bool old_sn;
		bool dup;
	} rx_pkt;
	struct {
		bool valid;
		u16 head_sn;
		u8 baid;
		u16 num_entries;
		/* The test prepares the reorder buffer with fake skbs based
		 * on the sequence numbers provided in @entries array.
		 */
		struct {
			u16 sn;
			/* Set add_subframes > 0 to simulate an A-MSDU by
			 * queueing additional @add_subframes skbs in the
			 * appropriate reorder buffer entry (based on the @sn)
			 */
			u8 add_subframes;
		} entries[BA_WINDOW_SIZE];
	} reorder_buf_state;
	struct {
		enum iwl_mld_reorder_result reorder_res;
		u16 head_sn;
		u16 num_stored;
		u16 skb_release_order[BA_WINDOW_SIZE];
		u16 skb_release_order_count;
	} expected;
} reorder_buffer_cases[] = {
	{
		.desc = "RX packet with invalid BAID",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.baid = IWL_RX_REORDER_DATA_INVALID_BAID,
		},
		.reorder_buf_state = {
			.valid = true,
		},
		.expected = {
			/* Invalid BAID should not be buffered. The frame is
			 * passed to the network stack immediately.
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
			.num_stored = 0,
		},
	},
	{
		.desc = "RX Multicast packet",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.multicast = true,
		},
		.reorder_buf_state = {
			.valid = true,
		},
		.expected = {
			/* Multicast packets are not buffered. The packet is
			 * passed to the network stack immediately.
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
			.num_stored = 0,
		},
	},
	{
		.desc = "RX non-QoS data",
		.rx_pkt = {
			.fc = IEEE80211_FTYPE_DATA,
		},
		.reorder_buf_state = {
			.valid = true,
		},
		.expected = {
			/* non-QoS data frames do not require reordering.
			 * The packet is passed to the network stack
			 * immediately.
			 */
		.reorder_res = IWL_MLD_PASS_SKB,
		},
	},
	{
		.desc = "RX old SN packet, reorder buffer is not yet valid",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.old_sn = true,
		},
		.reorder_buf_state = {
			.valid = false,
		},
		.expected = {
			/* The buffer is invalid and the RX packet has an old
			 * SN. The packet is passed to the network stack
			 * immediately.
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
		},
	},
	{
		.desc = "RX old SN packet, reorder buffer valid",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.old_sn = true,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
		},
		.expected = {
			/* The buffer is valid and the RX packet has an old SN.
			 * The packet should be dropped.
			 */
			.reorder_res = IWL_MLD_DROP_SKB,
			.num_stored = 0,
			.head_sn = 100,
		},
	},
	{
		.desc = "RX duplicate packet",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.dup = true,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
		},
		.expected = {
			/* Duplicate packets should be dropped */
			.reorder_res = IWL_MLD_DROP_SKB,
			.num_stored = 0,
			.head_sn = 100,
		},
	},
	{
		.desc = "RX In-order packet, sn < nssn",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 100,
			.nssn = 101,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
		},
		.expected = {
			/* 1. Reorder buffer is empty.
			 * 2. RX packet SN is in order and less than NSSN.
			 * Packet is released to the network stack immediately
			 * and buffer->head_sn is updated to NSSN.
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
			.num_stored = 0,
			.head_sn = 101,
		},
	},
	{
		.desc = "RX In-order packet, sn == head_sn",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 101,
			.nssn = 100,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 101,
		},
		.expected = {
			/* 1. Reorder buffer is empty.
			 * 2. RX packet SN is equal to buffer->head_sn.
			 * Packet is released to the network stack immediately
			 * and buffer->head_sn is incremented.
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
			.num_stored = 0,
			.head_sn = 102,
		},
	},
	{
		.desc = "RX In-order packet, IEEE80211_MAX_SN wrap around",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = IEEE80211_MAX_SN,
			.nssn = IEEE80211_MAX_SN - 1,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = IEEE80211_MAX_SN,
		},
		.expected = {
			/* 1. Reorder buffer is empty.
			 * 2. RX SN == buffer->head_sn == IEEE80211_MAX_SN
			 * Packet is released to the network stack immediately
			 * and buffer->head_sn is incremented correctly (wraps
			 * around to 0).
			 */
			.reorder_res = IWL_MLD_PASS_SKB,
			.num_stored = 0,
			.head_sn = 0,
		},
	},
	{
		.desc = "RX Out-of-order packet, pending packet in buffer",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 100,
			.nssn = 102,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 1,
			.entries[0].sn = 101,
		},
		.expected = {
			/* 1. Reorder buffer contains one packet with SN=101.
			 * 2. RX packet SN = buffer->head_sn.
			 * Both packets are released (in order) to the network
			 * stack as there are no gaps.
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 0,
			.head_sn = 102,
			.skb_release_order = {100, 101},
			.skb_release_order_count = 2,
		},
	},
	{
		.desc = "RX Out-of-order packet, pending packet in buffer (wrap around)",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 0,
			.nssn = 1,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = IEEE80211_MAX_SN - 1,
			.num_entries = 1,
			.entries[0].sn = IEEE80211_MAX_SN,
		},
		.expected = {
			/* 1. Reorder buffer contains one packet with
			 *    SN=IEEE80211_MAX_SN.
			 * 2. RX Packet SN = 0 (after wrap around)
			 * Both packets are released (in order) to the network
			 * stack as there are no gaps.
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 0,
			.head_sn = 1,
			.skb_release_order = { 4095, 0 },
			.skb_release_order_count = 2,
		},
	},
	{
		.desc = "RX Out-of-order packet, filling 1/2 holes in buffer, release RX packet",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 100,
			.nssn = 101,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 1,
			.entries[0].sn = 102,
		},
		.expected = {
			/* 1. Reorder buffer contains one packet with SN=102.
			 * 2. There are 2 holes at SN={100, 101}.
			 * Only the RX packet (SN=100) is released, there is
			 * still a hole at 101.
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 1,
			.head_sn = 101,
			.skb_release_order = {100},
			.skb_release_order_count = 1,
		},
	},
	{
		.desc = "RX Out-of-order packet, filling 1/2 holes, release 2 packets",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 102,
			.nssn = 103,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 3,
			.entries[0].sn = 101,
			.entries[1].sn = 104,
			.entries[2].sn = 105,
		},
		.expected = {
			/* 1. Reorder buffer contains three packets.
			 * 2. RX packet fills one of two holes (at SN=102).
			 * Two packets are released (until the next hole at
			 * SN=103).
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 2,
			.head_sn = 103,
			.skb_release_order = {101, 102},
			.skb_release_order_count = 2,
		},
	},
		{
		.desc = "RX Out-of-order packet, filling 1/1 holes, no packets released",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 102,
			.nssn = 100,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 3,
			.entries[0].sn = 101,
			.entries[1].sn = 103,
			.entries[2].sn = 104,
		},
		.expected = {
			/* 1. Reorder buffer contains three packets:
			 *    SN={101, 103, 104}.
			 * 2. RX packet fills a hole (SN=102), but NSSN is
			 *    smaller than buffered frames.
			 * No packets can be released yet and buffer->head_sn
			 * is not updated.
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 4,
			.head_sn = 100,
		},
	},
	{
		.desc = "RX In-order A-MSDU, last subframe",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 100,
			.nssn = 101,
			.amsdu = true,
			.last_subframe = true,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 1,
			.entries[0] = {
				.sn = 100,
				.add_subframes = 1,
			},
		},
		.expected = {
			/* 1. Reorder buffer contains a 2-sub frames A-MSDU
			 *    at SN=100.
			 * 2. RX packet is the last SN=100 A-MSDU subframe
			 * All packets are released in order (3 x SN=100).
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 0,
			.head_sn = 101,
			.skb_release_order = {100, 100, 100},
			.skb_release_order_count = 3,
		},
	},
	{
		.desc = "RX In-order A-MSDU, not the last subframe",
		.rx_pkt = {
			.fc = FC_QOS_DATA,
			.sn = 100,
			.nssn = 101,
			.amsdu = true,
			.last_subframe = false,
		},
		.reorder_buf_state = {
			.valid = true,
			.head_sn = 100,
			.num_entries = 1,
			.entries[0] = {
				.sn = 100,
				.add_subframes = 1,
			},
		},
		.expected = {
			/* 1. Reorder buffer contains a 2-sub frames A-MSDU
			 *    at SN=100.
			 * 2. RX packet additional SN=100 A-MSDU subframe,
			 *    but not the last one
			 * No packets are released and head_sn is not updated.
			 */
			.reorder_res = IWL_MLD_BUFFERED_SKB,
			.num_stored = 3,
			.head_sn = 100,
		},
	},
};

KUNIT_ARRAY_PARAM_DESC(test_reorder_buffer, reorder_buffer_cases, desc);

static struct sk_buff_head g_released_skbs;
static u16 g_num_released_skbs;

/* Add released skbs from reorder buffer to a global list; This allows us
 * to verify the correct release order of packets after they pass through the
 * simulated reorder logic.
 */
static void
fake_iwl_mld_pass_packet_to_mac80211(struct iwl_mld *mld,
				     struct napi_struct *napi,
				     struct sk_buff *skb, int queue,
				     struct ieee80211_sta *sta)
{
	__skb_queue_tail(&g_released_skbs, skb);
	g_num_released_skbs++;
}

static u32
fake_iwl_mld_fw_sta_id_mask(struct iwl_mld *mld, struct ieee80211_sta *sta)
{
	struct iwl_mld_sta *mld_sta = iwl_mld_sta_from_mac80211(sta);
	struct iwl_mld_link_sta *mld_link_sta;
	u8 link_id;
	u32 sta_mask = 0;

	/* This is the expectation in the real function */
	lockdep_assert_wiphy(mld->wiphy);

	/* We can't use for_each_sta_active_link */
	for_each_mld_link_sta(mld_sta, mld_link_sta, link_id)
		sta_mask |= BIT(mld_link_sta->fw_id);
	return sta_mask;
}

static struct iwl_rx_mpdu_desc *setup_mpdu_desc(void)
{
	struct kunit *test = kunit_get_current_test();
	const struct reorder_buffer_case *param =
		(const void *)(test->param_value);
	struct iwl_rx_mpdu_desc *mpdu_desc;

	KUNIT_ALLOC_AND_ASSERT(test, mpdu_desc);

	mpdu_desc->reorder_data |=
		le32_encode_bits(param->rx_pkt.baid,
				 IWL_RX_MPDU_REORDER_BAID_MASK);
	mpdu_desc->reorder_data |=
		le32_encode_bits(param->rx_pkt.sn,
				 IWL_RX_MPDU_REORDER_SN_MASK);
	mpdu_desc->reorder_data |=
		le32_encode_bits(param->rx_pkt.nssn,
				 IWL_RX_MPDU_REORDER_NSSN_MASK);
	if (param->rx_pkt.old_sn)
		mpdu_desc->reorder_data |=
			cpu_to_le32(IWL_RX_MPDU_REORDER_BA_OLD_SN);

	if (param->rx_pkt.dup)
		mpdu_desc->status |= cpu_to_le32(IWL_RX_MPDU_STATUS_DUPLICATE);

	if (param->rx_pkt.amsdu) {
		mpdu_desc->mac_flags2 |= IWL_RX_MPDU_MFLG2_AMSDU;
		if (param->rx_pkt.last_subframe)
			mpdu_desc->amsdu_info |=
				IWL_RX_MPDU_AMSDU_LAST_SUBFRAME;
	}

	return mpdu_desc;
}

static struct sk_buff *alloc_and_setup_skb(u16 fc, u16 seq_ctrl, u8 tid,
					   bool mcast)
{
	struct kunit *test = kunit_get_current_test();
	struct ieee80211_hdr hdr = {
		.frame_control = cpu_to_le16(fc),
		.seq_ctrl = cpu_to_le16(seq_ctrl),
	};
	struct sk_buff *skb;

	skb = kunit_zalloc_skb(test, 128, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, skb);

	if (ieee80211_is_data_qos(hdr.frame_control)) {
		u8 *qc = ieee80211_get_qos_ctl(&hdr);

		qc[0] = tid & IEEE80211_QOS_CTL_TID_MASK;
	}

	if (mcast)
		hdr.addr1[0] = 0x1;

	skb_set_mac_header(skb, skb->len);
	skb_put_data(skb, &hdr, ieee80211_hdrlen(hdr.frame_control));

	return skb;
}

static struct iwl_mld_reorder_buffer *
setup_reorder_buffer(struct iwl_mld_baid_data *baid_data)
{
	struct kunit *test = kunit_get_current_test();
	const struct reorder_buffer_case *param =
		(const void *)(test->param_value);
	struct iwl_mld_reorder_buffer *buffer = baid_data->reorder_buf;
	struct iwl_mld_reorder_buf_entry *entries = baid_data->entries;
	struct sk_buff *fake_skb;

	buffer->valid = param->reorder_buf_state.valid;
	buffer->head_sn = param->reorder_buf_state.head_sn;
	buffer->queue = QUEUE;

	for (int i = 0; i < baid_data->buf_size; i++)
		__skb_queue_head_init(&entries[i].frames);

	for (int i = 0; i < param->reorder_buf_state.num_entries; i++) {
		u16 sn = param->reorder_buf_state.entries[i].sn;
		int index = sn % baid_data->buf_size;
		u8 add_subframes =
			param->reorder_buf_state.entries[i].add_subframes;
		/* create 1 skb per entry + additional skbs per num of
		 * requested subframes
		 */
		u8 num_skbs = 1 + add_subframes;

		for (int j = 0; j < num_skbs; j++) {
			fake_skb = alloc_and_setup_skb(FC_QOS_DATA, sn, 0,
						       false);
			__skb_queue_tail(&entries[index].frames, fake_skb);
			buffer->num_stored++;
		}
	}

	return buffer;
}

static struct iwl_mld_reorder_buffer *setup_ba_data(struct ieee80211_sta *sta)
{
	struct kunit *test = kunit_get_current_test();
	struct iwl_mld *mld = test->priv;
	const struct reorder_buffer_case *param =
		(const void *)(test->param_value);
	struct iwl_mld_baid_data *baid_data = NULL;
	struct iwl_mld_reorder_buffer *buffer;
	u32 reorder_buf_size = BA_WINDOW_SIZE * sizeof(baid_data->entries[0]);
	u8 baid = param->reorder_buf_state.baid;

	/* Assuming only 1 RXQ */
	KUNIT_ALLOC_AND_ASSERT_SIZE(test, baid_data,
				    sizeof(*baid_data) + reorder_buf_size);

	baid_data->baid = baid;
	baid_data->tid = param->rx_pkt.tid;
	baid_data->buf_size = BA_WINDOW_SIZE;

	wiphy_lock(mld->wiphy);
	baid_data->sta_mask = iwl_mld_fw_sta_id_mask(mld, sta);
	wiphy_unlock(mld->wiphy);

	baid_data->entries_per_queue = BA_WINDOW_SIZE;

	buffer = setup_reorder_buffer(baid_data);

	KUNIT_EXPECT_NULL(test, rcu_access_pointer(mld->fw_id_to_ba[baid]));
	rcu_assign_pointer(mld->fw_id_to_ba[baid], baid_data);

	return buffer;
}

static void test_reorder_buffer(struct kunit *test)
{
	struct iwl_mld *mld = test->priv;
	const struct reorder_buffer_case *param =
		(const void *)(test->param_value);
	struct iwl_rx_mpdu_desc *mpdu_desc;
	struct ieee80211_vif *vif;
	struct ieee80211_sta *sta;
	struct sk_buff *skb;
	struct iwl_mld_reorder_buffer *buffer;
	enum iwl_mld_reorder_result reorder_res;
	u16 skb_release_order_count = param->expected.skb_release_order_count;
	u16 skb_idx = 0;

	/* Init globals and activate stubs */
	__skb_queue_head_init(&g_released_skbs);
	g_num_released_skbs = 0;
	kunit_activate_static_stub(test, iwl_mld_fw_sta_id_mask,
				   fake_iwl_mld_fw_sta_id_mask);
	kunit_activate_static_stub(test, iwl_mld_pass_packet_to_mac80211,
				   fake_iwl_mld_pass_packet_to_mac80211);

	vif = iwlmld_kunit_add_vif(false, NL80211_IFTYPE_STATION);
	sta = iwlmld_kunit_setup_sta(vif, IEEE80211_STA_AUTHORIZED, -1);

	/* Prepare skb, mpdu_desc, BA data and the reorder buffer */
	skb = alloc_and_setup_skb(param->rx_pkt.fc, param->rx_pkt.sn,
				  param->rx_pkt.tid, param->rx_pkt.multicast);
	buffer = setup_ba_data(sta);
	mpdu_desc = setup_mpdu_desc();

	rcu_read_lock();
	reorder_res = iwl_mld_reorder(mld, NULL, QUEUE, sta, skb, mpdu_desc);
	rcu_read_unlock();

	KUNIT_ASSERT_EQ(test, reorder_res, param->expected.reorder_res);
	KUNIT_ASSERT_EQ(test, buffer->num_stored, param->expected.num_stored);
	KUNIT_ASSERT_EQ(test, buffer->head_sn, param->expected.head_sn);

	/* Verify skbs release order */
	KUNIT_ASSERT_EQ(test, skb_release_order_count, g_num_released_skbs);
	while ((skb = __skb_dequeue(&g_released_skbs))) {
		struct ieee80211_hdr *hdr = (void *)skb_mac_header(skb);

		KUNIT_ASSERT_EQ(test, le16_to_cpu(hdr->seq_ctrl),
				param->expected.skb_release_order[skb_idx]);
		skb_idx++;
	}
	KUNIT_ASSERT_EQ(test, skb_idx, skb_release_order_count);
}

static struct kunit_case reorder_buffer_test_cases[] = {
	KUNIT_CASE_PARAM(test_reorder_buffer, test_reorder_buffer_gen_params),
	{},
};

static struct kunit_suite reorder_buffer = {
	.name = "iwlmld-reorder-buffer",
	.test_cases = reorder_buffer_test_cases,
	.init = iwlmld_kunit_test_init,
};

kunit_test_suite(reorder_buffer);