// SPDX-License-Identifier: GPL-2.0
/*
 * hugetlb-mmap:
 *
 * Example of using huge page memory in a user application using the mmap
 * system call.  Before running this application, make sure that the
 * administrator has mounted the hugetlbfs filesystem (on some directory
 * like /mnt) using the command mount -t hugetlbfs nodev /mnt. In this
 * example, the app is requesting memory of size 256MB that is backed by
 * huge pages.
 */
#define _GNU_SOURCE
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <linux/memfd.h>
#include "vm_util.h"
#include "kselftest.h"
#include "hugepage_settings.h"

#define LENGTH (256UL*1024*1024)
#define PROTECTION (PROT_READ | PROT_WRITE)

static void check_bytes(char *addr)
{
	ksft_print_msg("First hex is %x\n", *((unsigned int *)addr));
}

static void write_bytes(char *addr, size_t length)
{
	unsigned long i;

	for (i = 0; i < length; i++)
		*(addr + i) = (char)i;
}

static bool verify_bytes(char *addr, size_t length)
{
	unsigned long i;

	check_bytes(addr);
	for (i = 0; i < length; i++)
		if (*(addr + i) != (char)i) {
			ksft_print_msg("Error: Mismatch at %lu(%p)\n", i, addr);
			return false;
		}

	return true;
}

static void test_mmap(size_t length, int mmap_flags, int fd,
		      const char *test_name)
{
	bool passed = true;
	void *addr;

	addr = mmap(NULL, length, PROTECTION, mmap_flags, fd, 0);
	if (addr == MAP_FAILED)
		ksft_exit_fail_perror("mmap");

	ksft_print_msg("Returned address is %p\n", addr);
	check_bytes(addr);
	write_bytes(addr, length);
	if (!verify_bytes(addr, length))
		passed = false;

	/* munmap() length of MAP_HUGETLB memory must be hugepage aligned */
	if (munmap(addr, length))
		ksft_exit_fail_perror("munmap");

	ksft_test_result(passed, "%s\n", test_name);
}

static void test_anon_mmap(size_t length, int shift)
{
	const char *test_name = "hugetlb anonymous mmap";
	int mmap_flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB;

	if (shift)
		mmap_flags |= (shift & MAP_HUGE_MASK) << MAP_HUGE_SHIFT;

	test_mmap(length, mmap_flags, -1, test_name);
}

static void test_file_mmap(size_t length, int shift)
{
	const char *test_name = "hugetlb file mmap";
	int mfd_flags = MFD_HUGETLB;
	int fd;

	if (shift)
		mfd_flags |= (shift & MFD_HUGE_MASK) << MFD_HUGE_SHIFT;

	fd = memfd_create("hugetlb-mmap", mfd_flags);
	if (fd < 0)
		ksft_exit_fail_perror("memfd_create");

	test_mmap(length, MAP_SHARED, fd, test_name);
	close(fd);
}

int main(int argc, char **argv)
{
	size_t hugepage_size;
	size_t length = LENGTH;
	int shift = 0, nr;

	ksft_print_header();

	if (argc > 1)
		length = atol(argv[1]) << 20;
	if (argc > 2)
		shift = atoi(argv[2]);

	hugetlb_save_settings();
	if (shift) {
		hugepage_size = (1UL << shift);
		ksft_print_msg("%lu kB hugepages\n", 1UL << (shift - 10));
	} else {
		hugepage_size = default_huge_page_size();
		if (!hugepage_size)
			ksft_exit_skip("Could not detect default hugetlb page size.");
		ksft_print_msg("Default size hugepages (%lu kB)\n", hugepage_size >> 10);
	}

	/* munmap will fail if the length is not page aligned */
	length = (length + hugepage_size - 1) & ~(hugepage_size - 1);
	nr = length / hugepage_size;

	hugetlb_set_nr_pages(hugepage_size, nr);
	if (hugetlb_free_pages(hugepage_size) < nr)
		ksft_exit_skip("Not enough %lu Kb pages\n", hugepage_size >> 10);

	ksft_set_plan(2);
	ksft_print_msg("Mapping %lu Mbytes\n", (unsigned long)length >> 20);

	test_anon_mmap(length, shift);
	test_file_mmap(length, shift);

	ksft_finished();
}