selftests/bpf: Prevent allocating data larger than a page

Fix a bug in the task local data library that may allocate more than a
a page for tld_data_u. This may happen when users set a too large
TLD_DYN_DATA_SIZE, so check it when creating dynamic TLD fields and fix
the corresponding selftest.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
Link: https://lore.kernel.org/r/20260413190259.358442-2-ameryhung@gmail.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

2 files changed, 15 insertions, 6 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..489f07045c9f 100644
--- a/tools/testing/selftests/bpf/prog_tests/task_local_data.h
+++ b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
@@ -241,7 +241,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..8b99b4880d24 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 {
@@ -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));