// SPDX-License-Identifier: GPL-2.0-only
/*
 * KUnit tests for element fragmentation
 *
 * Copyright (C) 2023-2024 Intel Corporation
 */
#include <linux/ieee80211.h>
#include <net/cfg80211.h>
#include <kunit/test.h>

static void defragment_0(struct kunit *test)
{
 ssize_t ret;
 static const u8 input[] = {
  [0] = WLAN_EID_EXTENSION,
  [1] = 254,
  [2] = WLAN_EID_EXT_EHT_MULTI_LINK,
  [27] = 27,
  [123] = 123,
  [254 + 2] = WLAN_EID_FRAGMENT,
  [254 + 3] = 7,
  [254 + 3 + 7] = 0, /* for size */
 };
 u8 *data = kunit_kzalloc(test, sizeof(input), GFP_KERNEL);

 KUNIT_ASSERT_NOT_NULL(test, data);

 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       NULL, 0,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 253);
 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       data, ret,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 253);
 KUNIT_EXPECT_MEMEQ(test, data, input + 3, 253);
}

static void defragment_1(struct kunit *test)
{
 ssize_t ret;
 static const u8 input[] = {
  [0] = WLAN_EID_EXTENSION,
  [1] = 255,
  [2] = WLAN_EID_EXT_EHT_MULTI_LINK,
  [27] = 27,
  [123] = 123,
  [255 + 2] = WLAN_EID_FRAGMENT,
  [255 + 3] = 7,
  [255 + 3 + 1] = 0xaa,
  [255 + 3 + 8] = WLAN_EID_FRAGMENT, /* not used */
  [255 + 3 + 9] = 1,
  [255 + 3 + 10] = 0, /* for size */
 };
 u8 *data = kunit_kzalloc(test, sizeof(input), GFP_KERNEL);
 const struct element *elem;
 int count = 0;

 KUNIT_ASSERT_NOT_NULL(test, data);

 for_each_element(elem, input, sizeof(input))
  count++;

 /* check the elements are right */
 KUNIT_ASSERT_EQ(test, count, 3);

 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       NULL, 0,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 254 + 7);
 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       data, ret,
       WLAN_EID_FRAGMENT);
 /* this means the last fragment was not used */
 KUNIT_EXPECT_EQ(test, ret, 254 + 7);
 KUNIT_EXPECT_MEMEQ(test, data, input + 3, 254);
 KUNIT_EXPECT_MEMEQ(test, data + 254, input + 255 + 4, 7);
}

static void defragment_2(struct kunit *test)
{
 ssize_t ret;
 static const u8 input[] = {
  [0] = WLAN_EID_EXTENSION,
  [1] = 255,
  [2] = WLAN_EID_EXT_EHT_MULTI_LINK,
  [27] = 27,
  [123] = 123,

  [257 + 0] = WLAN_EID_FRAGMENT,
  [257 + 1] = 255,
  [257 + 20] = 0xaa,

  [2 * 257 + 0] = WLAN_EID_FRAGMENT,
  [2 * 257 + 1] = 1,
  [2 * 257 + 2] = 0xcc,
  [2 * 257 + 3] = WLAN_EID_FRAGMENT, /* not used */
  [2 * 257 + 4] = 1,
  [2 * 257 + 5] = 0, /* for size */
 };
 u8 *data = kunit_kzalloc(test, sizeof(input), GFP_KERNEL);
 const struct element *elem;
 int count = 0;

 KUNIT_ASSERT_NOT_NULL(test, data);

 for_each_element(elem, input, sizeof(input))
  count++;

 /* check the elements are right */
 KUNIT_ASSERT_EQ(test, count, 4);

 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       NULL, 0,
       WLAN_EID_FRAGMENT);
 /* this means the last fragment was not used */
 KUNIT_EXPECT_EQ(test, ret, 254 + 255 + 1);
 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       data, ret,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 254 + 255 + 1);
 KUNIT_EXPECT_MEMEQ(test, data, input + 3, 254);
 KUNIT_EXPECT_MEMEQ(test, data + 254, input + 257 + 2, 255);
 KUNIT_EXPECT_MEMEQ(test, data + 254 + 255, input + 2 * 257 + 2, 1);
}

static void defragment_at_end(struct kunit *test)
{
 ssize_t ret;
 static const u8 input[] = {
  [0] = WLAN_EID_EXTENSION,
  [1] = 255,
  [2] = WLAN_EID_EXT_EHT_MULTI_LINK,
  [27] = 27,
  [123] = 123,
  [255 + 2] = WLAN_EID_FRAGMENT,
  [255 + 3] = 7,
  [255 + 3 + 7] = 0, /* for size */
 };
 u8 *data = kunit_kzalloc(test, sizeof(input), GFP_KERNEL);

 KUNIT_ASSERT_NOT_NULL(test, data);

 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       NULL, 0,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 254 + 7);
 ret = cfg80211_defragment_element((void *)input,
       input, sizeof(input),
       data, ret,
       WLAN_EID_FRAGMENT);
 KUNIT_EXPECT_EQ(test, ret, 254 + 7);
 KUNIT_EXPECT_MEMEQ(test, data, input + 3, 254);
 KUNIT_EXPECT_MEMEQ(test, data + 254, input + 255 + 4, 7);
}

static struct kunit_case element_fragmentation_test_cases[] = {
 KUNIT_CASE(defragment_0),
 KUNIT_CASE(defragment_1),
 KUNIT_CASE(defragment_2),
 KUNIT_CASE(defragment_at_end),
 {}
};

static struct kunit_suite element_fragmentation = {
 .name = "cfg80211-element-defragmentation",
 .test_cases = element_fragmentation_test_cases,
};

kunit_test_suite(element_fragmentation);