[PATCH v2 4/4] vfio: selftests: Allow a size for vfio_dma_mapping_perf_test

[Aaron Lewis <aaronlewis@google.com>]

Allow the user to specify a DMA region size via the command line for
vfio_dma_mapping_perf_test.

Because the selftest harness also parses command-line parameters, sharing
them directly is problematic. Adding options directly to the test could
create conflicts with harness-defined options. Even without conflicts, the
harness would need to be updated to recognize test-specific options to avoid
failing on unknown parameters.

Resolve this by isolating the two sets of parameters. The standard command-line
options are consumed by the test itself. To pass options through to the test
harness, introduce a new '-a' option.

For example, both the test size and the test harness options can be set
like this:

./vfio_dma_mapping_perf_test -b 16G -a "-v vfio_type1_iommu_memfd_hugetlb_1gb"

This invocation configures a 16G DMA region and restricts execution to the
specified test variant, which is useful when debugging DMA mapping latency
issues for a specific IOMMU type.

Signed-off-by: Aaron Lewis <aaronlewis@google.com>
---
 .../vfio/vfio_dma_mapping_perf_test.c         | 159 ++++++++++++++++--
 1 file changed, 149 insertions(+), 10 deletions(-)

diff --git a/tools/testing/selftests/vfio/vfio_dma_mapping_perf_test.c b/tools/testing/selftests/vfio/vfio_dma_mapping_perf_test.c
index 5ef85deba4ee..af2273a0c6f5 100644
--- a/tools/testing/selftests/vfio/vfio_dma_mapping_perf_test.c
+++ b/tools/testing/selftests/vfio/vfio_dma_mapping_perf_test.c
@@ -4,6 +4,7 @@
 #include <sys/mman.h>
 #include <time.h>
 #include <unistd.h>
+#include <wordexp.h>
 
 #include <linux/iommufd.h>
 #include <linux/limits.h>
@@ -17,7 +18,12 @@
 
 #include "kselftest_harness.h"
 
-static const char *device_bdf;
+static struct {
+	u64 size;
+	const char *device_bdf;
+} test_params = {
+	.size = SZ_1G,
+};
 
 FIXTURE(vfio_dma_mapping_perf_test) {
 	struct iommu *iommu;
@@ -45,7 +51,7 @@ FIXTURE_VARIANT_ADD_ALL_IOMMU_MODES(anonymous_hugetlb_1gb, MAP_HUGETLB | MAP_HUG
 FIXTURE_SETUP(vfio_dma_mapping_perf_test)
 {
 	self->iommu = iommu_init(variant->iommu_mode);
-	self->device = vfio_pci_device_init(device_bdf, self->iommu);
+	self->device = vfio_pci_device_init(test_params.device_bdf, self->iommu);
 	self->iova_allocator = iova_allocator_init(self->iommu);
 }
 
@@ -58,7 +64,7 @@ FIXTURE_TEARDOWN(vfio_dma_mapping_perf_test)
 
 TEST_F(vfio_dma_mapping_perf_test, dma_map_unmap)
 {
-	const u64 size = SZ_1G;
+	const u64 size = test_params.size;
 	const int flags = variant->mmap_flags;
 	struct dma_region region;
 
@@ -115,7 +121,7 @@ FIXTURE_VARIANT_ADD_MEMFD_MODE(memfd_hugetlb_1gb,
 FIXTURE_SETUP(vfio_dma_mapping_perf_memfd_test)
 {
 	self->iommu = iommu_init(variant->iommu_mode);
-	self->device = vfio_pci_device_init(device_bdf, self->iommu);
+	self->device = vfio_pci_device_init(test_params.device_bdf, self->iommu);
 	self->iova_allocator = iova_allocator_init(self->iommu);
 }
 
@@ -159,17 +165,17 @@ static void teardown_memfd(int fd, u64 size, void *vaddr)
 
 TEST_F(vfio_dma_mapping_perf_memfd_test, dma_map_unmap_from_file)
 {
-	const u64 size = SZ_1G;
-	const int flags = variant->mmap_flags;
+	const u64 size = test_params.size;
+	const int mmap_flags = variant->mmap_flags;
 	struct dma_region region;
 	int fd;
 
 	printf("mmap size = %lluG\n", (unsigned long long)(size / SZ_1G));
 
-	region.vaddr = setup_memfd(&fd, size, variant->mmap_flags, variant->memfd_flags);
+	region.vaddr = setup_memfd(&fd, size, mmap_flags, variant->memfd_flags);
 
 	/* Skip the test if there aren't enough HugeTLB pages available. */
-	if (flags & MAP_HUGETLB && region.vaddr == MAP_FAILED)
+	if (mmap_flags & MAP_HUGETLB && region.vaddr == MAP_FAILED)
 		SKIP(return, "setup_memfd() failed: %s (%d)\n", strerror(errno), errno);
 	else
 		ASSERT_NE(region.vaddr, MAP_FAILED);
@@ -185,8 +191,141 @@ TEST_F(vfio_dma_mapping_perf_memfd_test, dma_map_unmap_from_file)
 	teardown_memfd(fd, size, region.vaddr);
 }
 
+/*
+ * Parses "[0-9]+[kmgt]?".
+ */
+u64 parse_size(const char *size)
+{
+	int shift = 0;
+	char *scale;
+	u64 base;
+
+	VFIO_ASSERT_TRUE(size && isdigit(size[0]),
+			 "Need at least one digit in '%s'.", size);
+
+	base = strtoull(size, &scale, 0);
+
+	VFIO_ASSERT_TRUE(base != ULLONG_MAX, "Overflow parsing size!");
+
+	switch (tolower(*scale)) {
+	case 't':
+		shift = 40;
+		break;
+	case 'g':
+		shift = 30;
+		break;
+	case 'm':
+		shift = 20;
+		break;
+	case 'k':
+		shift = 10;
+		break;
+	case 'b':
+	case '\0':
+		shift = 0;
+		break;
+	default:
+		VFIO_FAIL("Unknown size letter '%c'.", *scale);
+	}
+
+	VFIO_ASSERT_TRUE((base << shift) >> shift == base,
+			 "Overflow scaling size!");
+
+	return base << shift;
+}
+
+static void help(char *name)
+{
+	puts("");
+	printf("usage: %s [-h] [-b bytes] [-a \"test harness args\"]\n", name);
+	puts("");
+	printf(" -h: Display this help message.\n"
+	       " -b: Specify the size of the DMA region to be mapped\n"
+	       "     and unmapped. e.g. 16M or 8G, (default: 1G)\n"
+	       " -a: Args that are forwarded to the test harness,\n"
+	       "     e.g. -a \"-t dma_map_unmap_from_file\"\n");
+}
+
+struct harness_args {
+	int argc;
+	char **argv;
+	wordexp_t exp;
+};
+
+static void populate_harness_args(struct harness_args *args, const char *argv_0,
+				  const char *cmdlne)
+{
+	int flags = WRDE_NOCMD;
+
+	if (!args->argv) {
+		/*
+		 * Initialize the argument list with the program name (argv[0]).
+		 * WRDE_NOCMD disables command substitution for safety.
+		 */
+		if (wordexp(argv_0, &args->exp, flags) != 0)
+			VFIO_FAIL("Failed to evaluate test harness argv_0 args!");
+	}
+
+	flags |= WRDE_APPEND;
+
+	/*
+	 * Use wordexp() to reliably parse the user-supplied command line string
+	 * into individual arguments, respecting shell quoting and escaping rules.
+	 * WRDE_APPEND merges these new arguments with the earlier argv[0].
+	 */
+	if (wordexp(cmdlne, &args->exp, flags) != 0)
+		VFIO_FAIL("Failed to evaluate test harness cmdlne args!");
+
+	args->argc = args->exp.we_wordc;
+	args->argv = args->exp.we_wordv;
+}
+
+static void setup_test(struct harness_args *args, int *argc, char *argv[])
+{
+	char *h_argv[] = { argv[0], "-h" };
+	int opt;
+
+	test_params.device_bdf = vfio_selftests_get_bdf(argc, argv);
+
+	while ((opt = getopt(*argc, argv, "a:b:h")) != -1) {
+		switch (opt) {
+		case 'a':
+			populate_harness_args(args, argv[0], optarg);
+			break;
+		case 'b':
+			test_params.size = parse_size(optarg);
+			break;
+		case 'h':
+		default:
+			help(argv[0]);
+			exit(test_harness_run(2, h_argv));
+		}
+	}
+
+	// Reset getopt() state to allow the test harness to use it.
+	optind = 1;
+}
+
+static void teardown_test(struct harness_args *args)
+{
+	if (args->argv) {
+		args->argc = 0;
+		args->argv = NULL;
+		wordfree(&args->exp);
+	}
+}
+
 int main(int argc, char *argv[])
 {
-	device_bdf = vfio_selftests_get_bdf(&argc, argv);
-	return test_harness_run(argc, argv);
+	char *default_hargs[] = { argv[0], NULL };
+	struct harness_args args = {};
+	int r;
+
+	setup_test(&args, &argc, argv);
+
+	r = test_harness_run(args.argc ?: 1, args.argv ?: default_hargs);
+
+	teardown_test(&args);
+
+	return r;
 }
-- 
2.55.0.654.g21b8a5bc05-goog