// SPDX-License-Identifier: GPL-2.0
/* Test selecting other page sizes for mmap/shmget. */

#define _GNU_SOURCE
#include <sys/mman.h>
#include <linux/mman.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <sys/stat.h>
#include <unistd.h>
#include <stdarg.h>
#include <string.h>
#include "vm_util.h"
#include "kselftest.h"
#include "hugepage_settings.h"

#if !defined(MAP_HUGETLB)
#define MAP_HUGETLB	0x40000
#endif

#define SHM_HUGETLB     04000   /* segment will use huge TLB pages */
#ifndef SHM_HUGE_SHIFT
#define SHM_HUGE_SHIFT  26
#endif

#define NUM_PAGESIZES   5
#define NUM_PAGES 4

unsigned long page_sizes[NUM_PAGESIZES];
int num_page_sizes;

int ilog2(unsigned long v)
{
	int l = 0;
	while ((1UL << l) < v)
		l++;
	return l;
}

void show(unsigned long ps)
{
	if (ps == getpagesize())
		return;

	ksft_print_msg("%luMB: %lu\n", ps >> 20, hugetlb_free_pages(ps));
}

void test_mmap(unsigned long size, unsigned flags)
{
	char *map;
	unsigned long before, after;

	before = hugetlb_free_pages(size);
	map = mmap(NULL, size*NUM_PAGES, PROT_READ|PROT_WRITE,
			MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB|flags, -1, 0);
	if (map == MAP_FAILED)
		ksft_exit_fail_msg("mmap: %s\n", strerror(errno));

	memset(map, 0xff, size*NUM_PAGES);
	after = hugetlb_free_pages(size);

	show(size);
	ksft_test_result(size == getpagesize() || (before - after) == NUM_PAGES,
			 "%s mmap %lu %x\n", __func__, size, flags);

	if (munmap(map, size * NUM_PAGES))
		ksft_exit_fail_msg("%s: unmap %s\n", __func__, strerror(errno));
}

void test_shmget(unsigned long size, unsigned flags)
{
	/* values for PAGE_SIZE test */
	unsigned long before = NUM_PAGES;
	unsigned long after = 0;
	struct shm_info i;
	char *map;
	int id;

	if (size != getpagesize())
		before = hugetlb_free_pages(size);

	id = shmget(IPC_PRIVATE, size * NUM_PAGES, IPC_CREAT|0600|flags);
	if (id < 0) {
		if (errno == EPERM) {
			ksft_test_result_skip("shmget requires root privileges: %s\n",
					      strerror(errno));
			return;
		}
		ksft_exit_fail_msg("shmget: %s\n", strerror(errno));
	}

	if (shmctl(id, SHM_INFO, (void *)&i) < 0)
		ksft_exit_fail_msg("shmctl: %s\n", strerror(errno));

	map = shmat(id, NULL, 0600);
	if (map == MAP_FAILED)
		ksft_exit_fail_msg("shmat: %s\n", strerror(errno));

	shmctl(id, IPC_RMID, NULL);

	memset(map, 0xff, size*NUM_PAGES);
	if (size != getpagesize())
		after = hugetlb_free_pages(size);

	show(size);
	ksft_test_result((before - after) == NUM_PAGES,
			 "%s: mmap %lu %x\n", __func__, size, flags);
	if (shmdt(map))
		ksft_exit_fail_msg("%s: shmdt: %s\n", __func__, strerror(errno));
}

void find_pagesizes(void)
{
	unsigned long largest = getpagesize();
	int i;

	num_page_sizes = hugetlb_setup(NUM_PAGES, page_sizes, ARRAY_SIZE(page_sizes));

	for (i = 0; i < num_page_sizes; i++)
		if (page_sizes[i] > largest)
			largest = page_sizes[i];

	shm_limits_prepare(NUM_PAGES * largest);
}

int main(void)
{
	unsigned default_hps = default_huge_page_size();
	int i;

	ksft_print_header();

	find_pagesizes();

	if (!num_page_sizes)
		ksft_finished();

	ksft_set_plan(2 * num_page_sizes + 3);

	for (i = 0; i < num_page_sizes; i++) {
		unsigned long ps = page_sizes[i];
		int arg = ilog2(ps) << MAP_HUGE_SHIFT;

		ksft_print_msg("Testing %luMB mmap with shift %x\n", ps >> 20, arg);
		test_mmap(ps, MAP_HUGETLB | arg);
	}

	ksft_print_msg("Testing default huge mmap\n");
	test_mmap(default_hps, MAP_HUGETLB);

	ksft_print_msg("Testing non-huge shmget\n");
	test_shmget(getpagesize(), 0);

	for (i = 0; i < num_page_sizes; i++) {
		unsigned long ps = page_sizes[i];
		int arg = ilog2(ps) << SHM_HUGE_SHIFT;
		ksft_print_msg("Testing %luMB shmget with shift %x\n", ps >> 20, arg);
		test_shmget(ps, SHM_HUGETLB | arg);
	}

	ksft_print_msg("default huge shmget\n");
	test_shmget(default_hps, SHM_HUGETLB);

	ksft_finished();
}

SHM_LIMITS_RESTORE()