Make sure task local data is working correctly for different allocation
sizes. Existing task local data selftests allocate the maximum amount of
data possible but miss the garbage data issue when only small amount of
data is allocated. Therefore, test small data allocations as well.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/test_task_local_data.c     | 78 ++++++++++++++++++-
 1 file changed, 74 insertions(+), 4 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
index 8b99b4880d24..6a5806b36113 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
@@ -30,12 +30,12 @@ TLD_DEFINE_KEY(value0_key, "value0", sizeof(int));
  * sequentially. Users of task local data library should not touch
  * library internal.
  */
-static void reset_tld(void)
+static void reset_tld(__u16 dyn_data_size)
 {
 	if (tld_meta_p) {
 		/* Remove TLDs created by tld_create_key() */
 		tld_meta_p->cnt = 1;
-		tld_meta_p->size = TLD_DYN_DATA_SIZE;
+		tld_meta_p->size = dyn_data_size + 8;
 		memset(&tld_meta_p->metadata[1], 0,
 		       (TLD_MAX_DATA_CNT - 1) * sizeof(struct tld_metadata));
 	}
@@ -133,7 +133,7 @@ static void test_task_local_data_basic(void)
 	tld_key_t key;
 	int i, err;
 
-	reset_tld();
+	reset_tld(TLD_DYN_DATA_SIZE_MAX);
 
 	ASSERT_OK(pthread_mutex_init(&global_mutex, NULL), "pthread_mutex_init");
 
@@ -247,7 +247,7 @@ static void test_task_local_data_race(void)
 	tld_keys[0] = value0_key;
 
 	for (j = 0; j < 100; j++) {
-		reset_tld();
+		reset_tld(TLD_DYN_DATA_SIZE_MAX);
 
 		for (i = 0; i < TEST_RACE_THREAD_NUM; i++) {
 			/*
@@ -296,10 +296,80 @@ static void test_task_local_data_race(void)
 	test_task_local_data__destroy(skel);
 }
 
+static void test_task_local_data_dyn_size(__u16 dyn_data_size)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	struct test_task_local_data *skel;
+	int max_keys, i, err, fd, *data;
+	char name[TLD_NAME_LEN];
+	tld_key_t key;
+
+	reset_tld(dyn_data_size);
+
+	skel = test_task_local_data__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
+		return;
+
+	tld_keys = calloc(TLD_MAX_DATA_CNT, sizeof(tld_key_t));
+	if (!ASSERT_OK_PTR(tld_keys, "calloc tld_keys"))
+		goto out;
+
+	fd = bpf_map__fd(skel->maps.tld_data_map);
+
+	/* Create as many int-sized TLDs as the dynamic data size allows */
+	max_keys = dyn_data_size / TLD_ROUND_UP(sizeof(int), 8);
+	for (i = 0; i < max_keys; i++) {
+		snprintf(name, TLD_NAME_LEN, "value_%d", i);
+		tld_keys[i] = tld_create_key(name, sizeof(int));
+		if (!ASSERT_FALSE(tld_key_is_err(tld_keys[i]), "tld_create_key"))
+			goto out;
+
+		data = tld_get_data(fd, tld_keys[i]);
+		if (!ASSERT_OK_PTR(data, "tld_get_data"))
+			goto out;
+		*data = i;
+	}
+
+	/* The next key should fail with E2BIG */
+	key = tld_create_key("overflow", sizeof(int));
+	ASSERT_EQ(tld_key_err_or_zero(key), -E2BIG, "tld_create_key overflow");
+
+	/* Verify data for value_i do not overlap */
+	for (i = 0; i < max_keys; i++) {
+		data = tld_get_data(fd, tld_keys[i]);
+		if (!ASSERT_OK_PTR(data, "tld_get_data"))
+			goto out;
+
+		ASSERT_EQ(*data, i, "tld_get_data value_i");
+	}
+
+	/* Verify BPF side can still read the static key */
+	data = tld_get_data(fd, value0_key);
+	if (!ASSERT_OK_PTR(data, "tld_get_data value0"))
+		goto out;
+	*data = 0xdeadbeef;
+
+	err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
+	ASSERT_OK(err, "run task_main");
+	ASSERT_EQ(skel->bss->test_value0, 0xdeadbeef, "tld_get_data value0");
+
+out:
+	if (tld_keys) {
+		free(tld_keys);
+		tld_keys = NULL;
+	}
+	tld_free();
+	test_task_local_data__destroy(skel);
+}
+
 void test_task_local_data(void)
 {
 	if (test__start_subtest("task_local_data_basic"))
 		test_task_local_data_basic();
 	if (test__start_subtest("task_local_data_race"))
 		test_task_local_data_race();
+	if (test__start_subtest("task_local_data_dyn_size_small"))
+		test_task_local_data_dyn_size(64);
+	if (test__start_subtest("task_local_data_dyn_size_zero"))
+		test_task_local_data_dyn_size(0);
 }
-- 
2.52.0