[PATCH 2/2] selftests/mm: Add userfaultfd unit test

[Peter Xu]

Add a new test for userfaultfd unit test.

Signed-off-by: Peter Xu <peterx@xxxxxxxxxx>
---
 tools/testing/selftests/mm/.gitignore         |   1 +
 tools/testing/selftests/mm/Makefile           |   2 +
 .../selftests/mm/userfaultfd-unit-test.c      | 407 ++++++++++++++++++
 3 files changed, 410 insertions(+)
 create mode 100644 tools/testing/selftests/mm/userfaultfd-unit-test.c

diff --git a/tools/testing/selftests/mm/.gitignore b/tools/testing/selftests/mm/.gitignore
index 1f8c36a9fa10..7404f27cba8d 100644
--- a/tools/testing/selftests/mm/.gitignore
+++ b/tools/testing/selftests/mm/.gitignore
@@ -22,6 +22,7 @@ protection_keys_32
 protection_keys_64
 madv_populate
 userfaultfd
+userfaultfd-unit-test
 mlock-intersect-test
 mlock-random-test
 virtual_address_range
diff --git a/tools/testing/selftests/mm/Makefile b/tools/testing/selftests/mm/Makefile
index fbf5646b1072..a540ab7c84ec 100644
--- a/tools/testing/selftests/mm/Makefile
+++ b/tools/testing/selftests/mm/Makefile
@@ -56,6 +56,7 @@ TEST_GEN_FILES += on-fault-limit
 TEST_GEN_FILES += thuge-gen
 TEST_GEN_FILES += transhuge-stress
 TEST_GEN_FILES += userfaultfd
+TEST_GEN_FILES += userfaultfd-unit-test
 TEST_GEN_PROGS += soft-dirty
 TEST_GEN_PROGS += split_huge_page_test
 TEST_GEN_FILES += ksm_tests
@@ -111,6 +112,7 @@ $(OUTPUT)/madv_populate: vm_util.c
 $(OUTPUT)/soft-dirty: vm_util.c
 $(OUTPUT)/split_huge_page_test: vm_util.c
 $(OUTPUT)/userfaultfd: vm_util.c
+$(OUTPUT)/userfaultfd-unit-test: vm_util.c
 
 ifeq ($(MACHINE),x86_64)
 BINARIES_32 := $(patsubst %,$(OUTPUT)/%,$(BINARIES_32))
diff --git a/tools/testing/selftests/mm/userfaultfd-unit-test.c b/tools/testing/selftests/mm/userfaultfd-unit-test.c
new file mode 100644
index 000000000000..cd7b4f564f3b
--- /dev/null
+++ b/tools/testing/selftests/mm/userfaultfd-unit-test.c
@@ -0,0 +1,407 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Userfaultfd unit tests.
+ *
+ *  Copyright (C) 2023  Red Hat, Inc.
+ */
+#define _GNU_SOURCE
+#include <stdio.h>
+#include <errno.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <time.h>
+#include <signal.h>
+#include <poll.h>
+#include <string.h>
+#include <linux/mman.h>
+#include <sys/mman.h>
+#include <sys/syscall.h>
+#include <sys/ioctl.h>
+#include <sys/wait.h>
+#include <pthread.h>
+#include <linux/userfaultfd.h>
+#include <setjmp.h>
+#include <stdbool.h>
+#include <assert.h>
+#include <inttypes.h>
+#include <stdint.h>
+#include <sys/random.h>
+#include <semaphore.h>
+
+#include "../kselftest.h"
+#include "vm_util.h"
+
+/* 128MB */
+#define  ASYNC_MEM_SIZE  (128UL << 20)
+
+#define ALIGN_DOWN(x,a) ((x) & ~((a) - 1))
+#define PAGE_SIZE 4096
+
+int pagemap_fd;
+/* TODO: detect this */
+int thp_size = 2UL << 20;
+
+typedef unsigned char *(*mem_map_fn)(size_t size);
+
+unsigned char *anon_mem_map(size_t size);
+unsigned char *shmem_mem_map(size_t size);
+unsigned char *hugetlb_mem_map(size_t size);
+unsigned char *localfs_mem_map(size_t size);
+
+typedef enum {
+	TEST_ASYNC_ANON = 0,
+	TEST_ASYNC_SHMEM,
+	TEST_ASYNC_HUGETLB,
+	/* Direct test on `pwd`, with the hope that it's a local file system. */
+	TEST_ASYNC_LOCAL_FS,
+	TEST_ASYNC_NUM,
+} test_async_type;
+
+typedef struct {
+	const char *name;
+	/* madvise() used to zap a pte (only, but keep the data) */
+	int zap_madvise;
+	int page_size;
+	mem_map_fn mem_map;
+} async_ops_t;
+
+async_ops_t async_ops[TEST_ASYNC_NUM] = {
+	{
+		.name = "anonymous",
+		.zap_madvise = MADV_PAGEOUT,
+		.mem_map = anon_mem_map,
+		.page_size = PAGE_SIZE,
+	},
+	{
+		.name = "shmem",
+		.zap_madvise = MADV_DONTNEED,
+		.mem_map = shmem_mem_map,
+		.page_size = PAGE_SIZE,
+	},
+	{
+		.name = "hugetlb",
+		.zap_madvise = MADV_DONTNEED,
+		.mem_map = hugetlb_mem_map,
+		.page_size = (2UL << 20),
+	},
+	{
+		.name = "local-fs",
+		.zap_madvise = MADV_DONTNEED,
+		.mem_map = localfs_mem_map,
+		.page_size = PAGE_SIZE,
+	},
+};
+
+//#define debug(...) printf(__VA_ARGS__)
+#define debug(...)
+
+#define _err(fmt, ...)						\
+	do {							\
+		int ret = errno;				\
+		fprintf(stderr, "ERROR: " fmt, ##__VA_ARGS__);	\
+		fprintf(stderr, " (errno=%d, line=%d)\n",	\
+			ret, __LINE__);				\
+	} while (0)
+
+#define errexit(exitcode, fmt, ...)		\
+	do {					\
+		_err(fmt, ##__VA_ARGS__);	\
+		exit(exitcode);			\
+	} while (0)
+
+#define err(fmt, ...) errexit(1, fmt, ##__VA_ARGS__)
+
+#define PM_UFFD_WP                    (1UL << 57)
+
+#define  pagemap_check_wp(addr, wp) do {				\
+		if (!!(pagemap_get_entry(pagemap_fd, \
+					 (char *)addr) & PM_UFFD_WP) != wp) \
+			err("pagemap uffd-wp bit error: addr=0x%lx", \
+			    (unsigned long)addr);		     \
+	} while (0)
+
+static uint64_t random_uint64(void)
+{
+	uint64_t value;
+	int ret;
+
+	ret = getrandom(&value, sizeof(value), 0);
+	assert(ret == sizeof(value));
+
+	return value;
+}
+
+static int __userfaultfd_open_dev(void)
+{
+	int fd, _uffd;
+
+	fd = open("/dev/userfaultfd", O_RDWR | O_CLOEXEC);
+	if (fd < 0) {
+		perror("open(/dev/userfaultfd)");
+		return -1;
+	}
+
+	_uffd = ioctl(fd, USERFAULTFD_IOC_NEW, O_CLOEXEC);
+	if (_uffd < 0) {
+		perror("USERFAULTFD_IOC_NEW");
+		close(fd);
+		return -1;
+	}
+	close(fd);
+	return _uffd;
+}
+
+static void wp_range(int ufd, __u64 start, __u64 len, bool wp)
+{
+	struct uffdio_writeprotect prms;
+
+	/* Write protection page faults */
+	prms.range.start = start;
+	prms.range.len = len;
+	/* Undo write-protect, do wakeup after that */
+	prms.mode = wp ? UFFDIO_WRITEPROTECT_MODE_WP : 0;
+
+	if (ioctl(ufd, UFFDIO_WRITEPROTECT, &prms))
+		err("clear WP failed: address=0x%"PRIx64, (uint64_t)start);
+}
+
+unsigned char *anon_mem_map(size_t size)
+{
+	unsigned char *buffer;
+
+	buffer = mmap(NULL, size, PROT_READ|PROT_WRITE,
+		      MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+
+	if (buffer == MAP_FAILED)
+		err("mmap(MAP_HUGETLB)");
+
+	return buffer;
+}
+
+unsigned char *shmem_mem_map(size_t size)
+{
+	unsigned char *buffer;
+
+	buffer = mmap(NULL, size, PROT_READ|PROT_WRITE,
+		      MAP_SHARED|MAP_ANONYMOUS, -1, 0);
+
+	if (buffer == MAP_FAILED)
+		err("mmap(MAP_HUGETLB)");
+
+	return buffer;
+}
+
+unsigned char *hugetlb_mem_map(size_t size)
+{
+	unsigned char *buffer;
+
+	buffer = mmap(NULL, size, PROT_READ|PROT_WRITE,
+		      MAP_SHARED|MAP_ANONYMOUS|MAP_HUGETLB|MAP_HUGE_2MB,
+		      -1, 0);
+
+	if (buffer == MAP_FAILED)
+		err("mmap(MAP_HUGETLB)");
+
+	return buffer;
+}
+
+unsigned char *localfs_mem_map(size_t size)
+{
+#define  TMP_LOCAL_PATH  "./test-async"
+	int ret, fd = open(TMP_LOCAL_PATH, O_RDWR | O_CREAT, 0644);
+	unsigned char *buffer = NULL;
+
+	if (fd < 0) {
+		perror("open()");
+		return NULL;
+	}
+
+	ret = ftruncate(fd, size);
+	if (ret) {
+		perror("ftruncate()");
+		close(fd);
+		return NULL;
+	}
+
+	buffer = mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
+	if (buffer == MAP_FAILED) {
+		perror("mmap()");
+		close(fd);
+		return NULL;
+	}
+	close(fd);
+	unlink(TMP_LOCAL_PATH);
+
+	return buffer;
+}
+
+typedef struct {
+	async_ops_t *ops;
+	unsigned char *buffer;
+	size_t size;
+	volatile bool quit;
+} async_args_t;
+
+static void *uffd_zap_thread(void *data)
+{
+	async_args_t *args = data;
+	unsigned char *start = args->buffer;
+	size_t size = args->size;
+	async_ops_t *ops = args->ops;
+	unsigned int psize = ops->page_size;
+	int zap_madvise = ops->zap_madvise;
+	uint64_t offset;
+
+	debug("zap thread started\n");
+
+	while (!args->quit) {
+		offset = ALIGN_DOWN(random_uint64() % size, psize);
+		madvise(start + offset, psize, zap_madvise);
+		usleep(100);
+	}
+
+	debug("zap thread ended\n");
+
+	return NULL;
+}
+
+static bool bitmap_test(unsigned long *bitmap, unsigned long bit)
+{
+	unsigned long nbits = sizeof(unsigned long) * 8;
+
+	return bitmap[bit / nbits] & (1UL << (bit % nbits));
+}
+
+static bool test_async(async_ops_t *ops, int uffd,
+		       unsigned char *buffer, size_t size)
+{
+#define  LOOP_N  5
+	unsigned long *bitmap, npages, i;
+	pthread_t zap_tid;
+	async_args_t args = {
+		.buffer = buffer,
+		.size = size,
+		.ops = ops,
+		.quit = false,
+	};
+	int loops;
+
+	npages = size / ops->page_size;
+	bitmap = calloc(1, npages / 8);
+	assert(bitmap);
+
+	/* Random prefaults */
+	getrandom(bitmap, npages / 8, 0);
+	for (i = 0; i < npages; i++) {
+		if (bitmap_test(bitmap, i)) {
+			debug("prefault page %ld as write\n", i);
+			*(buffer + i * ops->page_size) = 1;
+		}
+	}
+
+	/* Create zapper to randomly zap pgtables */
+	pthread_create(&zap_tid, NULL, uffd_zap_thread, &args);
+
+	for (loops = 0; loops < LOOP_N; loops++) {
+		/* Start tracking, or reset, on all pages */
+		wp_range(uffd, (uintptr_t)buffer, size, true);
+
+		/* Random writes */
+		getrandom(bitmap, npages / 8, 0);
+		for (i = 0; i < npages; i++) {
+			if (bitmap_test(bitmap, i)) {
+				debug("update page %ld\n", i);
+				*(buffer + i * ops->page_size) = 2;
+			}
+		}
+
+		/* Verify pagemap tracked all the writes */
+		for (i = 0; i < npages; i++) {
+			debug("check page %ld\n", i);
+			pagemap_check_wp(buffer + i * ops->page_size,
+					 !bitmap_test(bitmap, i));
+		}
+	}
+
+	free(bitmap);
+	args.quit = true;
+	pthread_join(zap_tid, NULL);
+
+	return true;
+}
+
+static void test_uffd_wp_async_one(async_ops_t *ops)
+{
+	struct uffdio_register uffdio_register = {0};
+	struct uffdio_api uffdio_api;
+	unsigned char *buffer = NULL;
+	bool succeed = false;
+	int uffd = -1;
+
+	if (!ops->mem_map) {
+		ksft_test_result_skip("Userfaultfd-wp async (%s)\n",
+				      ops->name);
+		return;
+	}
+
+	buffer = ops->mem_map(ASYNC_MEM_SIZE);
+	if (!buffer)
+		goto out;
+
+	uffd = __userfaultfd_open_dev();
+	if (uffd < 0)
+		goto out;
+
+	uffdio_api.api = UFFD_API;
+	uffdio_api.features = UFFD_FEATURE_WP_ASYNC;
+	if (ioctl(uffd, UFFDIO_API, &uffdio_api)) {
+		perror("UFFDIO_API");
+		goto out;
+	}
+
+	uffdio_register.range.start = (unsigned long) buffer;
+	uffdio_register.range.len = ASYNC_MEM_SIZE;
+	uffdio_register.mode = UFFDIO_REGISTER_MODE_WP;
+	if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register)) {
+		perror("UFFDIO_REGISTER");
+		goto out;
+	}
+
+	succeed = test_async(ops, uffd, buffer, ASYNC_MEM_SIZE);
+
+	if (ioctl(uffd, UFFDIO_UNREGISTER, &uffdio_register.range)) {
+		perror("UFFDIO_UNREGISTER");
+		goto out;
+	}
+out:
+	if (uffd > 0)
+		close(uffd);
+	if (buffer)
+		munmap(buffer, ASYNC_MEM_SIZE);
+
+	if (succeed)
+		ksft_test_result_pass("Userfaultfd-wp async (%s)\n", ops->name);
+	else
+		ksft_test_result_fail("Userfaultfd-wp async (%s)\n", ops->name);
+}
+
+static void test_uffd_wp_async_all(void)
+{
+	test_async_type type;
+
+	for (type = 0; type < TEST_ASYNC_NUM; type++)
+		test_uffd_wp_async_one(&async_ops[type]);
+}
+
+int main(void)
+{
+	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
+
+	ksft_print_header();
+	ksft_set_plan(TEST_ASYNC_NUM);
+	test_uffd_wp_async_all();
+	return ksft_exit_pass();
+}
-- 
2.39.1


--wmm/kV5Ei1/X4v9u--