Merge branch 'fix-garbage-data-in-task-local-data'

Amery Hung says:

====================
Fix garbage data in task local data

Hi,

The patchset fixes two scenarios where BPF side task local data API may
see garbage data and adds corresponding selftests.
====================

Link: https://patch.msgid.link/20260413190259.358442-1-ameryhung@gmail.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

3 files changed, 100 insertions, 14 deletions

diff --git a/tools/testing/selftests/bpf/prog_tests/task_local_data.h b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
index 1e5c67c78ffb..8ae4fb2027f7 100644
--- a/tools/testing/selftests/bpf/prog_tests/task_local_data.h
+++ b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
@@ -99,14 +99,20 @@ struct tld_meta_u {
 	struct tld_metadata metadata[];
 };
 
+/*
+ * The unused field ensures map_val.start > 0. On the BPF side, __tld_fetch_key()
+ * calculates off by summing map_val.start and tld_key_t.off and treats off == 0
+ * as key not cached.
+ */
 struct tld_data_u {
-	__u64 start; /* offset of tld_data_u->data in a page */
+	__u64 unused;
 	char data[] __attribute__((aligned(8)));
 };
 
 struct tld_map_value {
 	void *data;
 	struct tld_meta_u *meta;
+	__u16 start; /* offset of tld_data_u->data in a page */
 };
 
 struct tld_meta_u * _Atomic tld_meta_p __attribute__((weak));
@@ -182,7 +188,7 @@ static int __tld_init_data_p(int map_fd)
 	 * is a page in BTF.
 	 */
 	map_val.data = (void *)(TLD_PAGE_MASK & (intptr_t)data);
-	data->start = (~TLD_PAGE_MASK & (intptr_t)data) + sizeof(struct tld_data_u);
+	map_val.start = (~TLD_PAGE_MASK & (intptr_t)data) + sizeof(struct tld_data_u);
 	map_val.meta = tld_meta_p;
 
 	err = bpf_map_update_elem(map_fd, &tid_fd, &map_val, 0);
@@ -241,7 +247,8 @@ retry:
 		 * TLD_DYN_DATA_SIZE is allocated for tld_create_key()
 		 */
 		if (dyn_data) {
-			if (off + TLD_ROUND_UP(size, 8) > tld_meta_p->size)
+			if (off + TLD_ROUND_UP(size, 8) > tld_meta_p->size ||
+			    tld_meta_p->size > TLD_PAGE_SIZE - sizeof(struct tld_data_u))
 				return (tld_key_t){-E2BIG};
 		} else {
 			if (off + TLD_ROUND_UP(size, 8) > TLD_PAGE_SIZE - sizeof(struct tld_data_u))
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 e219ff506b56..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
@@ -3,8 +3,14 @@
 #include <bpf/btf.h>
 #include <test_progs.h>
 
+/*
+ * Only a page is pinned to kernel, so the maximum amount of dynamic data
+ * allowed is page_size - sizeof(struct tld_data_u) - static TLD fields.
+ */
+#define TLD_DYN_DATA_SIZE_MAX (getpagesize() - sizeof(struct tld_data_u) - 8)
+
 #define TLD_FREE_DATA_ON_THREAD_EXIT
-#define TLD_DYN_DATA_SIZE (getpagesize() - 8)
+#define TLD_DYN_DATA_SIZE TLD_DYN_DATA_SIZE_MAX
 #include "task_local_data.h"
 
 struct test_tld_struct {
@@ -24,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));
 	}
@@ -127,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");
 
@@ -147,11 +153,13 @@ static void test_task_local_data_basic(void)
 
 	/*
 	 * Shouldn't be able to store data exceed a page. Create a TLD just big
-	 * enough to exceed a page. TLDs already created are int value0, int
-	 * value1, and struct test_tld_struct value2.
+	 * enough to exceed a page. Data already contains struct tld_data_u,
+	 * value0 and value1 of int type, and value 2 of struct test_tld_struct.
 	 */
-	key = tld_create_key("value_not_exist",
-			     TLD_PAGE_SIZE - 2 * sizeof(int) - sizeof(struct test_tld_struct) + 1);
+	key = tld_create_key("value_not_exist", TLD_PAGE_SIZE + 1 -
+						sizeof(struct tld_data_u) -
+						TLD_ROUND_UP(sizeof(int), 8) * 2 -
+						TLD_ROUND_UP(sizeof(struct test_tld_struct), 8));
 	ASSERT_EQ(tld_key_err_or_zero(key), -E2BIG, "tld_create_key");
 
 	key = tld_create_key("value2", sizeof(struct test_tld_struct));
@@ -239,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++) {
 			/*
@@ -288,10 +296,80 @@ out:
 	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);
 }
diff --git a/tools/testing/selftests/bpf/progs/task_local_data.bpf.h b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
index 1f396711f487..0df8a12fd61e 100644
--- a/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
+++ b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
@@ -86,13 +86,14 @@ struct tld_meta_u {
 };
 
 struct tld_data_u {
-	__u64 start; /* offset of tld_data_u->data in a page */
+	__u64 unused;
 	char data[__PAGE_SIZE - sizeof(__u64)] __attribute__((aligned(8)));
 };
 
 struct tld_map_value {
 	struct tld_data_u __uptr *data;
 	struct tld_meta_u __uptr *meta;
+	__u16 start; /* offset of tld_data_u->data in a page */
 };
 
 typedef struct tld_uptr_dummy {
@@ -176,7 +177,7 @@ static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_s
 	if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
 		return 0;
 
-	start = tld_obj->data_map->data->start;
+	start = tld_obj->data_map->start;
 	cnt = tld_obj->data_map->meta->cnt;
 	metadata = tld_obj->data_map->meta->metadata;