[PATCH v2 02/14] hv: vmbus: validate chunk buffer allocation and cleanup

[Emerson Busson <emersonbusson@gmail.com>]

Validate rounded sizes before storing them and preserve allocated
chunks when cleanup or GPADL teardown has an uncertain outcome. Add
focused KUnit coverage for buffer sizing, ownership decisions, and the
ring fallback order descent.

vmbus_alloc_buffer() rounds the requested size up to PAGE_SIZE before it
computes nr_pages. A zero size, and a size within PAGE_SIZE of U32_MAX,
both produce a rounded value that is wrong: zero allocates nothing, and
the near-U32_MAX cases wrap to a small nr_pages, so a small chunk array
is built for a buffer the host is told is large. Reject both before the
rounded size is stored or used.

vmbus_buffer_should_free() gathers the leak decision in one place. A
cleanup whose GPADL create or teardown message may have reached the host
cannot prove the host has released the pages. Handing those pages back
to the allocator, or re-encrypting them, while the host may still map
them is a use-after-free from the host's side and a Confidential
Computing hazard. Such a buffer is marked leaked and its chunks are
retained instead of freed.

The new cases live in drivers/hv/vmbus_buffer_test.c, which is built
into the hv_vmbus object rather than a separate module. That keeps the
cases off the production sources while still letting them call the
internal sizing, order-descent and free-decision helpers. Those helpers
are therefore not static; nothing outside hv_vmbus links against them,
and no symbol is exported for test purposes.

The cases pin both rules, the partial-allocation rollback, and the
order-descent fallback.

Signed-off-by: Emerson Busson <emersonbusson@gmail.com>
---
 drivers/hv/Kconfig             |  11 ++++
 drivers/hv/Makefile            |   1 +
 drivers/hv/channel.c           |  67 ++++++++++++++++++----
 drivers/hv/hyperv_vmbus.h      |  12 ++++
 drivers/hv/vmbus_buffer_test.c | 102 +++++++++++++++++++++++++++++++++
 5 files changed, 181 insertions(+), 12 deletions(-)
 create mode 100644 drivers/hv/vmbus_buffer_test.c

diff --git a/drivers/hv/Kconfig b/drivers/hv/Kconfig
index aa11bcefddf2..d44dd60fbc23 100644
--- a/drivers/hv/Kconfig
+++ b/drivers/hv/Kconfig
@@ -65,6 +65,17 @@ config HYPERV_VMBUS
 	help
 	  Select this option to enable Hyper-V Vmbus driver.
 
+config HYPERV_VMBUS_KUNIT_TEST
+	bool "Build Hyper-V VMBus buffer KUnit tests"
+	depends on HYPERV_VMBUS && KUNIT
+	default KUNIT_ALL_TESTS
+	help
+	  Build the vmbus buffer sizing, GPADL lifetime and reclaim KUnit
+	  cases into the hv_vmbus object. The cases reach internal helpers
+	  that are not exported, so they cannot live in a separate module.
+
+	  If unsure, say N.
+
 config MSHV_ROOT
 	tristate "Microsoft Hyper-V root partition support"
 	depends on HYPERV && (X86_64 || ARM64)
diff --git a/drivers/hv/Makefile b/drivers/hv/Makefile
index 888a748cc7cb..80d33e9c0e57 100644
--- a/drivers/hv/Makefile
+++ b/drivers/hv/Makefile
@@ -12,6 +12,7 @@ hv_vmbus-y := vmbus_drv.o \
 		 hv.o connection.o channel.o \
 		 channel_mgmt.o ring_buffer.o hv_trace.o
 hv_vmbus-$(CONFIG_HYPERV_TESTING)	+= hv_debugfs.o
+hv_vmbus-$(CONFIG_HYPERV_VMBUS_KUNIT_TEST) += vmbus_buffer_test.o
 hv_utils-y := hv_util.o hv_kvp.o hv_snapshot.o hv_utils_transport.o
 mshv_root-y := mshv_root_main.o mshv_synic.o mshv_eventfd.o mshv_irq.o \
 	       mshv_root_hv_call.o mshv_portid_table.o mshv_regions.o
diff --git a/drivers/hv/channel.c b/drivers/hv/channel.c
index f8feb2a0ec0a..7d5b99281872 100644
--- a/drivers/hv/channel.c
+++ b/drivers/hv/channel.c
@@ -28,6 +28,44 @@
 
 #include "hyperv_vmbus.h"
 
+/*
+ * vmbus_buffer_round_size() and the order-descent helpers are also called
+ * from vmbus_buffer_test.c, which is built into this same object. They are
+ * therefore not static; nothing outside hv_vmbus links against them.
+ */
+int vmbus_buffer_round_size(u32 size, u32 *rounded_size)
+{
+	if (!size)
+		return -EINVAL;
+
+	if (size > U32_MAX - (u32)PAGE_SIZE + 1)
+		return -EOVERFLOW;
+
+	*rounded_size = round_up(size, (u32)PAGE_SIZE);
+	return 0;
+}
+
+unsigned int vmbus_buffer_order(unsigned long remaining,
+				unsigned int max_order)
+{
+	return min_t(unsigned int, max_order, ilog2(remaining));
+}
+
+bool vmbus_buffer_lower_order(unsigned int *order)
+{
+	if (!*order)
+		return false;
+
+	(*order)--;
+	return true;
+}
+
+bool vmbus_buffer_should_free(const struct vmbus_buffer *buffer)
+{
+	return !buffer->leak && !buffer->gpadl.leak &&
+	       !buffer->gpadl.gpadl_handle;
+}
+
 static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 				  u32 size,
 				  bool confidential,
@@ -661,7 +699,8 @@ void vmbus_free_buffer(void *addr, struct page **chunks, u32 chunk_cnt)
 		return;
 	}
 
-	vunmap(addr);
+	if (addr)
+		vunmap(addr);
 
 	for (i = 0; i < chunk_cnt; i++) {
 		unsigned long vaddr =
@@ -683,8 +722,7 @@ void vmbus_release_buffer(struct vmbus_buffer *buffer)
 		return;
 
 	kvfree(buffer->pages);
-	if (!buffer->leak && !buffer->gpadl.leak &&
-	    !buffer->gpadl.gpadl_handle)
+	if (vmbus_buffer_should_free(buffer))
 		vmbus_free_buffer(buffer->addr, buffer->chunks,
 				  buffer->chunk_cnt);
 	memset(buffer, 0, sizeof(*buffer));
@@ -717,12 +755,13 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 				  struct page ***chunks_out,
 				  u32 *chunk_cnt_out)
 {
-	unsigned long nr_pages = PFN_UP(size);
-	unsigned long remaining = nr_pages;
+	u32 rounded_size;
+	unsigned long nr_pages;
+	unsigned long remaining;
 	unsigned long page_idx = 0;
 	struct page **chunks = NULL;
 	struct page **pages = NULL;
-	int order = MAX_PAGE_ORDER;
+	unsigned int order = MAX_PAGE_ORDER;
 	u32 chunk_cnt = 0;
 	void *addr;
 	u32 i;
@@ -731,13 +770,15 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 	*chunks_out = NULL;
 	*chunk_cnt_out = 0;
 
-	if (!nr_pages)
+	if (vmbus_buffer_round_size(size, &rounded_size))
 		return NULL;
+	nr_pages = rounded_size >> PAGE_SHIFT;
+	remaining = nr_pages;
 
 	/* If the buffer does not need to be decrypted, just use vzalloc() */
 	if ((!hv_is_isolation_supported() &&
 	     !cc_platform_has(CC_ATTR_GUEST_MEM_ENCRYPT)) || confidential)
-		return vzalloc(nr_pages << PAGE_SHIFT);
+		return vzalloc(rounded_size);
 
 	/* Worst case: every chunk is a single page. */
 	chunks = kvmalloc_objs(*chunks, nr_pages, GFP_KERNEL | __GFP_ZERO);
@@ -752,7 +793,7 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 		struct page *page;
 		gfp_t gfp;
 
-		order = min(order, ilog2(remaining));
+		order = vmbus_buffer_order(remaining, order);
 
 		/*
 		 * Use __GFP_NORETRY | __GFP_NOWARN to avoid OOM-killing,
@@ -767,7 +808,7 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 		page = alloc_pages_node(cpu_to_node(channel->target_cpu),
 					gfp, order);
 		if (!page) {
-			if (!order--)
+			if (!vmbus_buffer_lower_order(&order))
 				goto err;
 			continue;
 		}
@@ -794,7 +835,7 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 	if (!addr)
 		goto err;
 
-	memset(addr, 0, nr_pages << PAGE_SHIFT);
+	memset(addr, 0, rounded_size);
 
 	kvfree(pages);
 	*chunks_out = chunks;
@@ -1067,8 +1108,10 @@ int vmbus_teardown_gpadl(struct vmbus_channel *channel, struct vmbus_gpadl *gpad
 
 	info = kzalloc(sizeof(*info) +
 		       sizeof(struct vmbus_channel_gpadl_teardown), GFP_KERNEL);
-	if (!info)
+	if (!info) {
+		gpadl->leak = true;
 		return -ENOMEM;
+	}
 
 	init_completion(&info->waitevent);
 	info->waiting_channel = channel;
diff --git a/drivers/hv/hyperv_vmbus.h b/drivers/hv/hyperv_vmbus.h
index 20d023c9735e..93df3b35cfd4 100644
--- a/drivers/hv/hyperv_vmbus.h
+++ b/drivers/hv/hyperv_vmbus.h
@@ -551,4 +551,16 @@ int hv_create_ring_sysfs(struct vmbus_channel *channel,
 							    struct vm_area_desc *desc));
 int hv_remove_ring_sysfs(struct vmbus_channel *channel);
 
+/*
+ * vmbus buffer sizing, order-descent and free-decision helpers.
+ *
+ * These are shared with vmbus_buffer_test.c, which is built into the
+ * hv_vmbus object alongside channel.c. They stay unexported.
+ */
+int vmbus_buffer_round_size(u32 size, u32 *rounded_size);
+unsigned int vmbus_buffer_order(unsigned long remaining,
+				unsigned int max_order);
+bool vmbus_buffer_lower_order(unsigned int *order);
+bool vmbus_buffer_should_free(const struct vmbus_buffer *buffer);
+
 #endif /* _HYPERV_VMBUS_H */
diff --git a/drivers/hv/vmbus_buffer_test.c b/drivers/hv/vmbus_buffer_test.c
new file mode 100644
index 000000000000..d0a927a46caf
--- /dev/null
+++ b/drivers/hv/vmbus_buffer_test.c
@@ -0,0 +1,102 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * KUnit tests for Hyper-V VMBus buffer allocation and GPADL lifetime.
+ *
+ * Built into the hv_vmbus object rather than a separate module so the
+ * cases can reach the internal helpers declared in hyperv_vmbus.h
+ * without exporting them.
+ */
+#include <kunit/test.h>
+#include <linux/hyperv.h>
+#include <linux/mm.h>
+#include <linux/slab.h>
+#include <linux/vmalloc.h>
+
+#include "hyperv_vmbus.h"
+
+static void vmbus_buffer_size_rounding_test(struct kunit *test)
+{
+	u32 rounded_size;
+
+	KUNIT_EXPECT_EQ(test, vmbus_buffer_round_size(1, &rounded_size), 0);
+	KUNIT_EXPECT_EQ(test, rounded_size, (u32)PAGE_SIZE);
+	KUNIT_EXPECT_EQ(test, vmbus_buffer_round_size(PAGE_SIZE, &rounded_size), 0);
+	KUNIT_EXPECT_EQ(test, rounded_size, (u32)PAGE_SIZE);
+}
+
+static void vmbus_buffer_size_overflow_test(struct kunit *test)
+{
+	u32 rounded_size = 0;
+	int ret;
+
+	KUNIT_EXPECT_EQ(test, vmbus_buffer_round_size(0, &rounded_size), -EINVAL);
+	ret = vmbus_buffer_round_size(U32_MAX, &rounded_size);
+	KUNIT_EXPECT_EQ(test, ret, -EOVERFLOW);
+	ret = vmbus_buffer_round_size(U32_MAX - PAGE_SIZE + 1, &rounded_size);
+	KUNIT_EXPECT_EQ(test, ret, 0);
+	KUNIT_EXPECT_EQ(test, rounded_size,
+			(u32)(U32_MAX - PAGE_SIZE + 1));
+}
+
+static void vmbus_ring_fallback_order_zero_test(struct kunit *test)
+{
+	unsigned int order;
+
+	order = vmbus_buffer_order(1UL << MAX_PAGE_ORDER, MAX_PAGE_ORDER);
+	KUNIT_EXPECT_EQ(test, order, (unsigned int)MAX_PAGE_ORDER);
+	while (order)
+		KUNIT_ASSERT_TRUE(test, vmbus_buffer_lower_order(&order));
+	KUNIT_EXPECT_FALSE(test, vmbus_buffer_lower_order(&order));
+	KUNIT_EXPECT_EQ(test, order, 0U);
+	KUNIT_EXPECT_EQ(test, vmbus_buffer_order(3, MAX_PAGE_ORDER), 1U);
+}
+
+static void vmbus_buffer_failed_teardown_leaks_test(struct kunit *test)
+{
+	struct vmbus_buffer buffer = {
+		.addr = (void *)1,
+		.gpadl.gpadl_handle = 1,
+	};
+
+	KUNIT_EXPECT_FALSE(test, vmbus_buffer_should_free(&buffer));
+	buffer.gpadl.gpadl_handle = 0;
+	buffer.gpadl.leak = true;
+	KUNIT_EXPECT_FALSE(test, vmbus_buffer_should_free(&buffer));
+	buffer.gpadl.leak = false;
+	buffer.leak = true;
+	KUNIT_EXPECT_FALSE(test, vmbus_buffer_should_free(&buffer));
+	buffer.leak = false;
+	KUNIT_EXPECT_TRUE(test, vmbus_buffer_should_free(&buffer));
+}
+
+static void vmbus_buffer_partial_allocation_cleanup_test(struct kunit *test)
+{
+	struct page **chunks;
+	struct vmbus_buffer buffer = {};
+
+	chunks = kmalloc_obj(*chunks, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, chunks);
+	vmbus_free_buffer(NULL, chunks, 0);
+
+	buffer.addr = vzalloc(PAGE_SIZE);
+	KUNIT_ASSERT_NOT_NULL(test, buffer.addr);
+	vmbus_release_buffer(&buffer);
+	KUNIT_EXPECT_PTR_EQ(test, buffer.addr, NULL);
+	vmbus_release_buffer(&buffer);
+}
+
+static struct kunit_case vmbus_buffer_test_cases[] = {
+	KUNIT_CASE(vmbus_buffer_size_rounding_test),
+	KUNIT_CASE(vmbus_buffer_size_overflow_test),
+	KUNIT_CASE(vmbus_ring_fallback_order_zero_test),
+	KUNIT_CASE(vmbus_buffer_failed_teardown_leaks_test),
+	KUNIT_CASE(vmbus_buffer_partial_allocation_cleanup_test),
+	{}
+};
+
+static struct kunit_suite vmbus_buffer_test_suite = {
+	.name = "hyperv-vmbus-buffer",
+	.test_cases = vmbus_buffer_test_cases,
+};
+
+kunit_test_suite(vmbus_buffer_test_suite);
-- 
2.43.0