// SPDX-License-Identifier: GPL-2.0-only
/*
 * Basic VM_PFNMAP tests relying on mmap() of input file provided.
 * Use '/dev/mem' as default.
 *
 * Copyright 2025, Red Hat, Inc.
 *
 * Author(s): David Hildenbrand <david@redhat.com>
 */
#define _GNU_SOURCE
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <errno.h>
#include <stdio.h>
#include <ctype.h>
#include <fcntl.h>
#include <signal.h>
#include <setjmp.h>
#include <linux/mman.h>
#include <sys/mman.h>
#include <sys/wait.h>

#include "kselftest_harness.h"
#include "vm_util.h"

#define DEV_MEM_NPAGES	2

static sigjmp_buf sigjmp_buf_env;
static char *file = "/dev/mem";
static off_t file_offset;
static int fd;

static void signal_handler(int sig)
{
	siglongjmp(sigjmp_buf_env, -EFAULT);
}

static int test_read_access(char *addr, size_t size, size_t pagesize)
{
	int ret;

	if (signal(SIGSEGV, signal_handler) == SIG_ERR)
		return -EINVAL;

	ret = sigsetjmp(sigjmp_buf_env, 1);
	if (!ret)
		force_read_pages(addr, size/pagesize, pagesize);

	if (signal(SIGSEGV, SIG_DFL) == SIG_ERR)
		return -EINVAL;

	return ret;
}

static int find_ram_target(off_t *offset,
		unsigned long long pagesize)
{
	unsigned long long start, end;
	char line[80], *end_ptr;
	FILE *file;

	/* Search /proc/iomem for the first suitable "System RAM" range. */
	file = fopen("/proc/iomem", "r");
	if (!file)
		return -errno;

	while (fgets(line, sizeof(line), file)) {
		/* Ignore any child nodes. */
		if (!isalnum(line[0]))
			continue;

		if (!strstr(line, "System RAM\n"))
			continue;

		start = strtoull(line, &end_ptr, 16);
		/* Skip over the "-" */
		end_ptr++;
		/* Make end "exclusive". */
		end = strtoull(end_ptr, NULL, 16) + 1;

		/* Actual addresses are not exported */
		if (!start && !end)
			break;

		/* We need full pages. */
		start = (start + pagesize - 1) & ~(pagesize - 1);
		end &= ~(pagesize - 1);

		if (start != (off_t)start)
			break;

		/* We need two pages. */
		if (end > start + DEV_MEM_NPAGES * pagesize) {
			fclose(file);
			*offset = start;
			return 0;
		}
	}
	return -ENOENT;
}

static void pfnmap_init(void)
{
	size_t pagesize = getpagesize();
	size_t size = DEV_MEM_NPAGES * pagesize;
	void *addr;

	if (strncmp(file, "/dev/mem", strlen("/dev/mem")) == 0) {
		int err = find_ram_target(&file_offset, pagesize);

		if (err)
			ksft_exit_skip("Cannot find ram target in '/proc/iomem': %s\n",
				       strerror(-err));
	} else {
		file_offset = 0;
	}

	fd = open(file, O_RDONLY);
	if (fd < 0)
		ksft_exit_skip("Cannot open '%s': %s\n", file, strerror(errno));

	/*
	 * Make sure we can map the file, and perform some basic checks; skip
	 * the whole suite if anything goes wrong.
	 * A fresh mapping is then created for every test case by
	 * FIXTURE_SETUP(pfnmap).
	 */
	addr = mmap(NULL, size, PROT_READ, MAP_SHARED, fd, file_offset);
	if (addr == MAP_FAILED)
		ksft_exit_skip("Cannot mmap '%s': %s\n", file, strerror(errno));

	if (!check_vmflag_pfnmap(addr))
		ksft_exit_skip("Invalid file: '%s'. Not pfnmap'ed\n", file);

	if (test_read_access(addr, size, pagesize))
		ksft_exit_skip("Cannot read-access mmap'ed '%s'\n", file);

	munmap(addr, size);
}

FIXTURE(pfnmap)
{
	size_t pagesize;
	char *addr1;
	size_t size1;
	char *addr2;
	size_t size2;
};

FIXTURE_SETUP(pfnmap)
{
	self->pagesize = getpagesize();

	self->size1 = DEV_MEM_NPAGES * self->pagesize;
	self->addr1 = mmap(NULL, self->size1, PROT_READ, MAP_SHARED,
			   fd, file_offset);
	ASSERT_NE(self->addr1, MAP_FAILED);

	self->size2 = 0;
	self->addr2 = MAP_FAILED;
}

FIXTURE_TEARDOWN(pfnmap)
{
	if (self->addr2 != MAP_FAILED)
		munmap(self->addr2, self->size2);
	if (self->addr1 != MAP_FAILED)
		munmap(self->addr1, self->size1);
}

TEST_F(pfnmap, madvise_disallowed)
{
	int advices[] = {
		MADV_DONTNEED,
		MADV_DONTNEED_LOCKED,
		MADV_FREE,
		MADV_WIPEONFORK,
		MADV_COLD,
		MADV_PAGEOUT,
		MADV_POPULATE_READ,
		MADV_POPULATE_WRITE,
	};
	int i;

	/* All these advices must be rejected. */
	for (i = 0; i < ARRAY_SIZE(advices); i++) {
		EXPECT_LT(madvise(self->addr1, self->pagesize, advices[i]), 0);
		EXPECT_EQ(errno, EINVAL);
	}
}

TEST_F(pfnmap, munmap_split)
{
	/*
	 * Unmap the first page. This munmap() call is not really expected to
	 * fail, but we might be able to trigger other internal issues.
	 */
	ASSERT_EQ(munmap(self->addr1, self->pagesize), 0);

	/*
	 * Remap the first page while the second page is still mapped. This
	 * makes sure that any PAT tracking on x86 will allow for mmap()'ing
	 * a page again while some parts of the first mmap() are still
	 * around.
	 */
	self->size2 = self->pagesize;
	self->addr2 = mmap(NULL, self->pagesize, PROT_READ, MAP_SHARED,
			   fd, file_offset);
	ASSERT_NE(self->addr2, MAP_FAILED);
}

TEST_F(pfnmap, mremap_fixed)
{
	char *ret;

	/* Reserve a destination area. */
	self->size2 = self->size1;
	self->addr2 = mmap(NULL, self->size2, PROT_READ, MAP_ANON | MAP_PRIVATE,
			   -1, 0);
	ASSERT_NE(self->addr2, MAP_FAILED);

	/* mremap() over our destination. */
	ret = mremap(self->addr1, self->size1, self->size2,
		     MREMAP_FIXED | MREMAP_MAYMOVE, self->addr2);
	ASSERT_NE(ret, MAP_FAILED);
}

TEST_F(pfnmap, mremap_shrink)
{
	char *ret;

	/* Shrinking is expected to work. */
	ret = mremap(self->addr1, self->size1, self->size1 - self->pagesize, 0);
	ASSERT_NE(ret, MAP_FAILED);
}

TEST_F(pfnmap, mremap_expand)
{
	/*
	 * Growing is not expected to work, and getting it right would
	 * be challenging. So this test primarily serves as an early warning
	 * that something that probably should never work suddenly works.
	 */
	self->size2 = self->size1 + self->pagesize;
	self->addr2 = mremap(self->addr1, self->size1, self->size2, MREMAP_MAYMOVE);
	ASSERT_EQ(self->addr2, MAP_FAILED);
}

TEST_F(pfnmap, fork)
{
	pid_t pid;
	int ret;

	/* fork() a child and test if the child can access the pages. */
	pid = fork();
	ASSERT_GE(pid, 0);

	if (!pid) {
		EXPECT_EQ(test_read_access(self->addr1, self->size1,
					   self->pagesize), 0);
		exit(0);
	}

	wait(&ret);
	if (WIFEXITED(ret))
		ret = WEXITSTATUS(ret);
	else
		ret = -EINVAL;
	ASSERT_EQ(ret, 0);
}

int main(int argc, char **argv)
{
	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--") == 0) {
			if (i + 1 < argc && strlen(argv[i + 1]) > 0)
				file = argv[i + 1];
			argc = i;
			break;
		}
	}

	pfnmap_init();

	return test_harness_run(argc, argv);
}