bpf: Check freeing sk->sk_local_storage with sk_local_storage->smap is NULLlocal-stg-mem-alloc

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>

2 files changed, 31 insertions, 10 deletions

diff --git a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
index 7399292ad08fd..a75ac6ceb90e2 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
@@ -156,7 +156,7 @@ void test_test_local_storage(void)
 	if (!ASSERT_GE(serv_sk, 0, "start_server"))
 		goto close_prog_rmdir;
 
-	ASSERT_EQ(skel->data->sk_storage_result, 0, "sk_storage_result");
+	ASSERT_EQ(skel->bss->sk_storage_errs, 0, "sk_storage_errs");
 
 	if (!check_syscall_operations(bpf_map__fd(skel->maps.sk_storage_map),
 				      serv_sk))
diff --git a/tools/testing/selftests/bpf/progs/local_storage.c b/tools/testing/selftests/bpf/progs/local_storage.c
index 19423ed862e35..c6eaf5d71d1be 100644
--- a/tools/testing/selftests/bpf/progs/local_storage.c
+++ b/tools/testing/selftests/bpf/progs/local_storage.c
@@ -15,7 +15,7 @@ char _license[] SEC("license") = "GPL";
 
 int monitored_pid = 0;
 int inode_storage_result = -1;
-int sk_storage_result = -1;
+int sk_storage_errs = 0;
 
 struct local_storage {
 	struct inode *exec_inode;
@@ -119,24 +119,45 @@ int BPF_PROG(socket_bind, struct socket *sock, struct sockaddr *address,
 	if (!storage)
 		return 0;
 
-	if (storage->value != DUMMY_STORAGE_VALUE)
-		sk_storage_result = -1;
+	if (storage->value != DUMMY_STORAGE_VALUE) {
+		sk_storage_errs++;
+		return 0;
+	}
 
 	/* This tests that we can associate multiple elements
 	 * with the local storage.
 	 */
 	storage = bpf_sk_storage_get(&sk_storage_map2, sock->sk, 0,
 				     BPF_LOCAL_STORAGE_GET_F_CREATE);
-	if (!storage)
+	if (!storage) {
+		sk_storage_errs++;
 		return 0;
+	}
 
-	err = bpf_sk_storage_delete(&sk_storage_map, sock->sk);
-	if (err)
+	if (bpf_sk_storage_delete(&sk_storage_map2, sock->sk)) {
+		sk_storage_errs++;
 		return 0;
+	}
 
-	err = bpf_sk_storage_delete(&sk_storage_map2, sock->sk);
-	if (!err)
-		sk_storage_result = err;
+	storage = bpf_sk_storage_get(&sk_storage_map2, sock->sk, 0,
+				     BPF_LOCAL_STORAGE_GET_F_CREATE);
+	if (!storage) {
+		sk_storage_errs++;
+		return 0;
+	}
+
+	if (bpf_sk_storage_delete(&sk_storage_map, sock->sk)) {
+		sk_storage_errs++;
+		return 0;
+	}
+
+	/* Ensure that the sk_storage_map is disconnected from the storage.
+	 * The storage memory should not be freed back to the
+	 * bpf_mem_cache ub the sk_bpf_storage_map because
+	 * sk_bpf_storage_map may have been gone.
+	 */
+	if (!sock->sk->sk_bpf_storage || sock->sk->sk_bpf_storage->smap)
+		sk_storage_errs++;
 
 	return 0;
 }