// SPDX-License-Identifier: GPL-2.0
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sys/ioctl.h>
#include <linux/userfaultfd.h>
#include <linux/fs.h>
#include <sys/syscall.h>
#include <unistd.h>
#include "kselftest.h"
#include "vm_util.h"

#define PMD_SIZE_FILE_PATH "/sys/kernel/mm/transparent_hugepage/hpage_pmd_size"
#define SMAP_FILE_PATH "/proc/self/smaps"
#define STATUS_FILE_PATH "/proc/self/status"
#define MAX_LINE_LENGTH 500
#define PAGEMAP_PATH "/proc/self/pagemap"
#define KPAGEFLAGS_PATH "/proc/kpageflags"
#define MAX_NR_ORDERS 20

unsigned int __page_size;
unsigned int __page_shift;

uint64_t pagemap_get_entry(int fd, char *start)
{
	const unsigned long pfn = (unsigned long)start / getpagesize();
	uint64_t entry;
	int ret;

	ret = pread(fd, &entry, sizeof(entry), pfn * sizeof(entry));
	if (ret != sizeof(entry))
		ksft_exit_fail_msg("reading pagemap failed\n");
	return entry;
}

static int __pagemap_scan_get_categories(int fd, char *start, struct page_region *r)
{
	struct pm_scan_arg arg;

	arg.start = (uintptr_t)start;
	arg.end = (uintptr_t)(start + psize());
	arg.vec = (uintptr_t)r;
	arg.vec_len = 1;
	arg.flags = 0;
	arg.size = sizeof(struct pm_scan_arg);
	arg.max_pages = 0;
	arg.category_inverted = 0;
	arg.category_mask = 0;
	arg.category_anyof_mask = PAGE_IS_WPALLOWED | PAGE_IS_WRITTEN | PAGE_IS_FILE |
				  PAGE_IS_PRESENT | PAGE_IS_SWAPPED | PAGE_IS_PFNZERO |
				  PAGE_IS_HUGE | PAGE_IS_SOFT_DIRTY;
	arg.return_mask = arg.category_anyof_mask;

	return ioctl(fd, PAGEMAP_SCAN, &arg);
}

static uint64_t pagemap_scan_get_categories(int fd, char *start)
{
	struct page_region r;
	int ret;

	ret = __pagemap_scan_get_categories(fd, start, &r);
	if (ret < 0)
		ksft_exit_fail_msg("PAGEMAP_SCAN failed: %s\n", strerror(errno));
	if (ret == 0)
		return 0;
	return r.categories;
}

static bool pagemap_scan_supported(int fd)
{
	const size_t pagesize = getpagesize();
	static int supported = -1;
	struct page_region r;
	void *test_area;

	if (supported != -1)
		return supported;

	test_area = mmap(0, pagesize, PROT_READ | PROT_WRITE,
			MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
	if (test_area == MAP_FAILED) {
		ksft_print_msg("WARN: mmap() failed: %s\n", strerror(errno));
		supported = 0;
	} else {
		supported = __pagemap_scan_get_categories(fd, test_area, &r) >= 0;
		munmap(test_area, pagesize);
	}

	return supported;
}

static bool page_entry_is(int fd, char *start, char *desc,
			  uint64_t pagemap_flags, uint64_t pagescan_flags)
{
	bool m = pagemap_get_entry(fd, start) & pagemap_flags;

	if (pagemap_scan_supported(fd)) {
		bool s = pagemap_scan_get_categories(fd, start) & pagescan_flags;

		if (m == s)
			return m;

		ksft_exit_fail_msg(
			"read and ioctl return unmatched results for %s: %d %d", desc, m, s);
	}
	return m;
}

bool pagemap_is_softdirty(int fd, char *start)
{
	return page_entry_is(fd, start, "soft-dirty",
				PM_SOFT_DIRTY, PAGE_IS_SOFT_DIRTY);
}

bool pagemap_is_swapped(int fd, char *start)
{
	return page_entry_is(fd, start, "swap", PM_SWAP, PAGE_IS_SWAPPED);
}

bool pagemap_is_populated(int fd, char *start)
{
	return page_entry_is(fd, start, "populated",
				PM_PRESENT | PM_SWAP,
				PAGE_IS_PRESENT | PAGE_IS_SWAPPED);
}

unsigned long pagemap_get_pfn(int fd, char *start)
{
	uint64_t entry = pagemap_get_entry(fd, start);

	/* If present (63th bit), PFN is at bit 0 -- 54. */
	if (entry & PM_PRESENT)
		return entry & 0x007fffffffffffffull;
	return -1ul;
}

void clear_softdirty(void)
{
	int ret;
	const char *ctrl = "4";
	int fd = open("/proc/self/clear_refs", O_WRONLY);

	if (fd < 0)
		ksft_exit_fail_msg("opening clear_refs failed\n");
	ret = write(fd, ctrl, strlen(ctrl));
	close(fd);
	if (ret != (signed int)strlen(ctrl))
		ksft_exit_fail_msg("writing clear_refs failed\n");
}

bool check_for_pattern(FILE *fp, const char *pattern, char *buf, size_t len)
{
	while (fgets(buf, len, fp)) {
		if (!strncmp(buf, pattern, strlen(pattern)))
			return true;
	}
	return false;
}

uint64_t read_pmd_pagesize(void)
{
	int fd;
	char buf[20];
	ssize_t num_read;

	fd = open(PMD_SIZE_FILE_PATH, O_RDONLY);
	if (fd == -1)
		return 0;

	num_read = read(fd, buf, 19);
	if (num_read < 1) {
		close(fd);
		return 0;
	}
	buf[num_read] = '\0';
	close(fd);

	return strtoul(buf, NULL, 10);
}

unsigned long rss_anon(void)
{
	unsigned long rss_anon = 0;
	FILE *fp;
	char buffer[MAX_LINE_LENGTH];

	fp = fopen(STATUS_FILE_PATH, "r");
	if (!fp)
		ksft_exit_fail_msg("%s: Failed to open file %s\n", __func__, STATUS_FILE_PATH);

	if (!check_for_pattern(fp, "RssAnon:", buffer, sizeof(buffer)))
		goto err_out;

	if (sscanf(buffer, "RssAnon:%10lu kB", &rss_anon) != 1)
		ksft_exit_fail_msg("Reading status error\n");

err_out:
	fclose(fp);
	return rss_anon;
}

static int vaddr_pageflags_get(char *vaddr, int pagemap_fd, int kpageflags_fd,
		uint64_t *flags)
{
	unsigned long pfn;

	pfn = pagemap_get_pfn(pagemap_fd, vaddr);

	/* non-present PFN */
	if (pfn == -1UL)
		return 1;

	if (pageflags_get(pfn, kpageflags_fd, flags))
		return -1;

	return 0;
}

/*
 * gather_folio_orders - scan through [vaddr_start, len) and record
 * folio orders
 *
 * @vaddr_start: start vaddr
 * @len: range length
 * @pagemap_fd: file descriptor to /proc/<pid>/pagemap
 * @kpageflags_fd: file descriptor to /proc/kpageflags
 * @orders: output folio order array
 * @nr_orders: folio order array size
 *
 * gather_folio_orders() scan through [vaddr_start, len) and check
 * all folios within the range and record their orders. All order-0 pages will
 * be recorded. Non-present vaddr is skipped.
 *
 * Return: 0 - no error, -1 - unhandled cases
 */
int gather_folio_orders(char *vaddr_start, size_t len,
		int pagemap_fd, int kpageflags_fd, int orders[], int nr_orders)
{
	uint64_t page_flags = 0;
	int cur_order = -1;
	char *vaddr;

	if (pagemap_fd == -1 || kpageflags_fd == -1)
		return -1;
	if (!orders)
		return -1;
	if (nr_orders <= 0)
		return -1;

	for (vaddr = vaddr_start; vaddr < vaddr_start + len;) {
		char *next_folio_vaddr;
		int status;

		status = vaddr_pageflags_get(vaddr, pagemap_fd, kpageflags_fd,
				&page_flags);
		if (status < 0)
			return -1;

		/* skip non present vaddr */
		if (status == 1) {
			vaddr += psize();
			continue;
		}

		/* all order-0 pages with possible false postive (non folio) */
		if (!(page_flags & (KPF_COMPOUND_HEAD | KPF_COMPOUND_TAIL))) {
			orders[0]++;
			vaddr += psize();
			continue;
		}

		/* skip non thp compound pages */
		if (!(page_flags & KPF_THP)) {
			vaddr += psize();
			continue;
		}

		/* vpn points to part of a THP at this point */
		if (page_flags & KPF_COMPOUND_HEAD)
			cur_order = 1;
		else {
			vaddr += psize();
			continue;
		}

		next_folio_vaddr = vaddr + (1UL << (cur_order + pshift()));

		if (next_folio_vaddr >= vaddr_start + len)
			break;

		while ((status = vaddr_pageflags_get(next_folio_vaddr,
						     pagemap_fd, kpageflags_fd,
						     &page_flags)) >= 0) {
			/*
			 * non present vaddr, next compound head page, or
			 * order-0 page
			 */
			if (status == 1 ||
			    (page_flags & KPF_COMPOUND_HEAD) ||
			    !(page_flags & (KPF_COMPOUND_HEAD | KPF_COMPOUND_TAIL))) {
				if (cur_order < nr_orders) {
					orders[cur_order]++;
					cur_order = -1;
					vaddr = next_folio_vaddr;
				}
				break;
			}

			cur_order++;
			next_folio_vaddr = vaddr + (1UL << (cur_order + pshift()));
		}

		if (status < 0)
			return status;
	}
	if (cur_order > 0 && cur_order < nr_orders)
		orders[cur_order]++;
	return 0;
}

char *__get_smap_entry(void *addr, const char *pattern, char *buf, size_t len)
{
	int ret;
	FILE *fp;
	char *entry = NULL;
	char addr_pattern[MAX_LINE_LENGTH];

	ret = snprintf(addr_pattern, MAX_LINE_LENGTH, "%08lx-",
		       (unsigned long) addr);
	if (ret >= MAX_LINE_LENGTH)
		ksft_exit_fail_msg("%s: Pattern is too long\n", __func__);

	fp = fopen(SMAP_FILE_PATH, "r");
	if (!fp)
		ksft_exit_fail_msg("%s: Failed to open file %s\n", __func__, SMAP_FILE_PATH);

	if (!check_for_pattern(fp, addr_pattern, buf, len))
		goto err_out;

	/* Fetch the pattern in the same block */
	if (!check_for_pattern(fp, pattern, buf, len))
		goto err_out;

	/* Trim trailing newline */
	entry = strchr(buf, '\n');
	if (entry)
		*entry = '\0';

	entry = buf + strlen(pattern);

err_out:
	fclose(fp);
	return entry;
}

static bool __check_pmd_huge(void *addr, char *pattern, int nr_hpages,
		  uint64_t hpage_size)
{
	char buffer[MAX_LINE_LENGTH];
	uint64_t thp = -1;
	char *entry;

	entry = __get_smap_entry(addr, pattern, buffer, sizeof(buffer));
	if (!entry)
		goto err_out;

	if (sscanf(entry, "%9" SCNu64 " kB", &thp) != 1)
		ksft_exit_fail_msg("Reading smap error\n");

err_out:
	return thp == (nr_hpages * (hpage_size >> 10));
}

static bool check_large_folios(void *addr, size_t len, int nr_hpages,
		uint64_t hpage_size)
{
	int order = 0, pagesize = getpagesize();
	unsigned int nr_pages = hpage_size / pagesize;
	int orders[MAX_NR_ORDERS], status;
	int pagemap_fd, kpageflags_fd;
	bool ret = false;

	if (!nr_pages)
		ksft_exit_fail_msg("invalid hugepage size\n");

	order = 31 - __builtin_clz(nr_pages);
	if (!order || order >= MAX_NR_ORDERS)
		ksft_exit_fail_msg("invalid order\n");

	memset(orders, 0, sizeof(int) * MAX_NR_ORDERS);
	pagemap_fd = open(PAGEMAP_PATH, O_RDONLY);
	if (pagemap_fd == -1)
		ksft_exit_fail_msg("read pagemap fail\n");

	kpageflags_fd = open(KPAGEFLAGS_PATH, O_RDONLY);
	if (kpageflags_fd == -1) {
		close(pagemap_fd);
		ksft_exit_fail_msg("read kpageflags fail\n");
	}

	status = gather_folio_orders(addr, len, pagemap_fd,
			kpageflags_fd, orders, MAX_NR_ORDERS);
	if (status)
		goto out;

	if (orders[order] == nr_hpages)
		ret = true;

out:
	close(pagemap_fd);
	close(kpageflags_fd);
	return ret;
}

bool check_huge_anon(void *addr, size_t len, int nr_hpages, uint64_t hpage_size)
{
	uint64_t pmd_pagesize = read_pmd_pagesize();

	if (!pmd_pagesize)
		ksft_exit_fail_msg("reading PMD pagesize failed\n");

	if (hpage_size == pmd_pagesize)
		return __check_pmd_huge(addr, "AnonHugePages: ", nr_hpages, hpage_size);

	return check_large_folios(addr, len, nr_hpages, hpage_size);
}

bool check_huge_file(void *addr, size_t len, int nr_hpages, uint64_t hpage_size)
{
	uint64_t pmd_pagesize = read_pmd_pagesize();

	if (!pmd_pagesize)
		ksft_exit_fail_msg("reading PMD pagesize failed\n");

	if (hpage_size == pmd_pagesize)
		return __check_pmd_huge(addr, "FilePmdMapped:", nr_hpages, hpage_size);

	return check_large_folios(addr, len, nr_hpages, hpage_size);
}

bool check_huge_shmem(void *addr, size_t len, int nr_hpages, uint64_t hpage_size)
{
	uint64_t pmd_pagesize = read_pmd_pagesize();

	if (!pmd_pagesize)
		ksft_exit_fail_msg("reading PMD pagesize failed\n");

	if (hpage_size == pmd_pagesize)
		return __check_pmd_huge(addr, "ShmemPmdMapped:", nr_hpages, hpage_size);

	return check_large_folios(addr, len, nr_hpages, hpage_size);
}

int64_t allocate_transhuge(void *ptr, int pagemap_fd)
{
	uint64_t ent[2];

	/* drop pmd */
	if (mmap(ptr, HPAGE_SIZE, PROT_READ | PROT_WRITE,
		 MAP_FIXED | MAP_ANONYMOUS |
		 MAP_NORESERVE | MAP_PRIVATE, -1, 0) != ptr)
		ksft_exit_fail_msg("mmap transhuge\n");

	if (madvise(ptr, HPAGE_SIZE, MADV_HUGEPAGE))
		ksft_exit_fail_msg("MADV_HUGEPAGE\n");

	/* allocate transparent huge page */
	*(volatile void **)ptr = ptr;

	if (pread(pagemap_fd, ent, sizeof(ent),
		  (uintptr_t)ptr >> (pshift() - 3)) != sizeof(ent))
		ksft_exit_fail_msg("read pagemap\n");

	if (PAGEMAP_PRESENT(ent[0]) && PAGEMAP_PRESENT(ent[1]) &&
	    PAGEMAP_PFN(ent[0]) + 1 == PAGEMAP_PFN(ent[1]) &&
	    !(PAGEMAP_PFN(ent[0]) & ((1 << (HPAGE_SHIFT - pshift())) - 1)))
		return PAGEMAP_PFN(ent[0]);

	return -1;
}

int pageflags_get(unsigned long pfn, int kpageflags_fd, uint64_t *flags)
{
	size_t count;

	count = pread(kpageflags_fd, flags, sizeof(*flags),
		      pfn * sizeof(*flags));

	if (count != sizeof(*flags))
		return -1;

	return 0;
}

/* If `ioctls' non-NULL, the allowed ioctls will be returned into the var */
int uffd_register_with_ioctls(int uffd, void *addr, uint64_t len,
			      bool miss, bool wp, bool minor, uint64_t *ioctls)
{
	struct uffdio_register uffdio_register = { 0 };
	uint64_t mode = 0;
	int ret = 0;

	if (miss)
		mode |= UFFDIO_REGISTER_MODE_MISSING;
	if (wp)
		mode |= UFFDIO_REGISTER_MODE_WP;
	if (minor)
		mode |= UFFDIO_REGISTER_MODE_MINOR;

	uffdio_register.range.start = (unsigned long)addr;
	uffdio_register.range.len = len;
	uffdio_register.mode = mode;

	if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
		ret = -errno;
	else if (ioctls)
		*ioctls = uffdio_register.ioctls;

	return ret;
}

int uffd_register(int uffd, void *addr, uint64_t len,
		  bool miss, bool wp, bool minor)
{
	return uffd_register_with_ioctls(uffd, addr, len,
					 miss, wp, minor, NULL);
}

int uffd_unregister(int uffd, void *addr, uint64_t len)
{
	struct uffdio_range range = { .start = (uintptr_t)addr, .len = len };
	int ret = 0;

	if (ioctl(uffd, UFFDIO_UNREGISTER, &range) == -1)
		ret = -errno;

	return ret;
}

static bool check_vmflag(void *addr, const char *flag)
{
	char buffer[MAX_LINE_LENGTH];
	const char *flags;
	size_t flaglen;

	flags = __get_smap_entry(addr, "VmFlags:", buffer, sizeof(buffer));
	if (!flags)
		ksft_exit_fail_msg("%s: No VmFlags for %p\n", __func__, addr);

	while (true) {
		flags += strspn(flags, " ");

		flaglen = strcspn(flags, " ");
		if (!flaglen)
			return false;

		if (flaglen == strlen(flag) && !memcmp(flags, flag, flaglen))
			return true;

		flags += flaglen;
	}
}

bool check_vmflag_io(void *addr)
{
	return check_vmflag(addr, "io");
}

bool check_vmflag_pfnmap(void *addr)
{
	return check_vmflag(addr, "pf");
}

bool check_vmflag_guard(void *addr)
{
	return check_vmflag(addr, "gu");
}

bool softdirty_supported(void)
{
	char *addr;
	bool supported = false;
	const size_t pagesize = getpagesize();

	/* New mappings are expected to be marked with VM_SOFTDIRTY (sd). */
	addr = mmap(0, pagesize, PROT_READ | PROT_WRITE,
		    MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
	if (addr == MAP_FAILED)
		ksft_exit_fail_msg("mmap failed\n");

	supported = check_vmflag(addr, "sd");
	munmap(addr, pagesize);
	return supported;
}

/*
 * Open an fd at /proc/$pid/maps and configure procmap_out ready for
 * PROCMAP_QUERY query. Returns 0 on success, or an error code otherwise.
 */
int open_procmap(pid_t pid, struct procmap_fd *procmap_out)
{
	char path[256];
	int ret = 0;

	memset(procmap_out, '\0', sizeof(*procmap_out));
	sprintf(path, "/proc/%d/maps", pid);
	procmap_out->query.size = sizeof(procmap_out->query);
	procmap_out->fd = open(path, O_RDONLY);
	if (procmap_out->fd < 0)
		ret = -errno;

	return ret;
}

/* Perform PROCMAP_QUERY. Returns 0 on success, or an error code otherwise. */
int query_procmap(struct procmap_fd *procmap)
{
	int ret = 0;

	if (ioctl(procmap->fd, PROCMAP_QUERY, &procmap->query) == -1)
		ret = -errno;

	return ret;
}

/*
 * Try to find the VMA at specified address, returns true if found, false if not
 * found, and the test is failed if any other error occurs.
 *
 * On success, procmap->query is populated with the results.
 */
bool find_vma_procmap(struct procmap_fd *procmap, void *address)
{
	int err;

	procmap->query.query_flags = 0;
	procmap->query.query_addr = (unsigned long)address;
	err = query_procmap(procmap);
	if (!err)
		return true;

	if (err != -ENOENT)
		ksft_exit_fail_msg("%s: Error %d on ioctl(PROCMAP_QUERY)\n",
				   __func__, err);
	return false;
}

/*
 * Close fd used by PROCMAP_QUERY mechanism. Returns 0 on success, or an error
 * code otherwise.
 */
int close_procmap(struct procmap_fd *procmap)
{
	return close(procmap->fd);
}

int write_sysfs(const char *file_path, unsigned long val)
{
	FILE *f = fopen(file_path, "w");

	if (!f) {
		fprintf(stderr, "f %s\n", file_path);
		perror("fopen");
		return 1;
	}
	if (fprintf(f, "%lu", val) < 0) {
		perror("fprintf");
		fclose(f);
		return 1;
	}
	fclose(f);

	return 0;
}

int read_sysfs(const char *file_path, unsigned long *val)
{
	FILE *f = fopen(file_path, "r");

	if (!f) {
		fprintf(stderr, "f %s\n", file_path);
		perror("fopen");
		return 1;
	}
	if (fscanf(f, "%lu", val) != 1) {
		perror("fscanf");
		fclose(f);
		return 1;
	}
	fclose(f);

	return 0;
}

void *sys_mremap(void *old_address, unsigned long old_size,
		 unsigned long new_size, int flags, void *new_address)
{
	return (void *)syscall(__NR_mremap, (unsigned long)old_address,
			       old_size, new_size, flags,
			       (unsigned long)new_address);
}

bool detect_huge_zeropage(void)
{
	int fd = open("/sys/kernel/mm/transparent_hugepage/use_zero_page",
		      O_RDONLY);
	bool enabled = 0;
	char buf[15];
	int ret;

	if (fd < 0)
		return 0;

	ret = pread(fd, buf, sizeof(buf), 0);
	if (ret > 0 && ret < sizeof(buf)) {
		buf[ret] = 0;

		if (strtoul(buf, NULL, 10) == 1)
			enabled = 1;
	}

	close(fd);
	return enabled;
}

long ksm_get_self_zero_pages(void)
{
	int proc_self_ksm_stat_fd;
	char buf[200];
	char *substr_ksm_zero;
	size_t value_pos;
	ssize_t read_size;

	proc_self_ksm_stat_fd = open("/proc/self/ksm_stat", O_RDONLY);
	if (proc_self_ksm_stat_fd < 0)
		return -errno;

	read_size = pread(proc_self_ksm_stat_fd, buf, sizeof(buf) - 1, 0);
	close(proc_self_ksm_stat_fd);
	if (read_size < 0)
		return -errno;

	buf[read_size] = 0;

	substr_ksm_zero = strstr(buf, "ksm_zero_pages");
	if (!substr_ksm_zero)
		return 0;

	value_pos = strcspn(substr_ksm_zero, "0123456789");
	return strtol(substr_ksm_zero + value_pos, NULL, 10);
}

long ksm_get_self_merging_pages(void)
{
	int proc_self_ksm_merging_pages_fd;
	char buf[10];
	ssize_t ret;

	proc_self_ksm_merging_pages_fd = open("/proc/self/ksm_merging_pages",
						O_RDONLY);
	if (proc_self_ksm_merging_pages_fd < 0)
		return -errno;

	ret = pread(proc_self_ksm_merging_pages_fd, buf, sizeof(buf) - 1, 0);
	close(proc_self_ksm_merging_pages_fd);
	if (ret <= 0)
		return -errno;
	buf[ret] = 0;

	return strtol(buf, NULL, 10);
}

long ksm_get_full_scans(void)
{
	int ksm_full_scans_fd;
	char buf[10];
	ssize_t ret;

	ksm_full_scans_fd = open("/sys/kernel/mm/ksm/full_scans", O_RDONLY);
	if (ksm_full_scans_fd < 0)
		return -errno;

	ret = pread(ksm_full_scans_fd, buf, sizeof(buf) - 1, 0);
	close(ksm_full_scans_fd);
	if (ret <= 0)
		return -errno;
	buf[ret] = 0;

	return strtol(buf, NULL, 10);
}

int ksm_use_zero_pages(void)
{
	int ksm_use_zero_pages_fd;
	ssize_t ret;

	ksm_use_zero_pages_fd = open("/sys/kernel/mm/ksm/use_zero_pages", O_RDWR);
	if (ksm_use_zero_pages_fd < 0)
		return -errno;

	ret = write(ksm_use_zero_pages_fd, "1", 1);
	close(ksm_use_zero_pages_fd);
	return ret == 1 ? 0 : -errno;
}

int ksm_start(void)
{
	int ksm_fd;
	ssize_t ret;
	long start_scans, end_scans;

	ksm_fd = open("/sys/kernel/mm/ksm/run", O_RDWR);
	if (ksm_fd < 0)
		return -errno;

	/* Wait for two full scans such that any possible merging happened. */
	start_scans = ksm_get_full_scans();
	if (start_scans < 0) {
		close(ksm_fd);
		return start_scans;
	}
	ret = write(ksm_fd, "1", 1);
	close(ksm_fd);
	if (ret != 1)
		return -errno;
	do {
		end_scans = ksm_get_full_scans();
		if (end_scans < 0)
			return end_scans;
	} while (end_scans < start_scans + 2);

	return 0;
}

int ksm_stop(void)
{
	int ksm_fd;
	ssize_t ret;

	ksm_fd = open("/sys/kernel/mm/ksm/run", O_RDWR);
	if (ksm_fd < 0)
		return -errno;

	ret = write(ksm_fd, "2", 1);
	close(ksm_fd);
	return ret == 1 ? 0 : -errno;
}

int get_hardware_corrupted_size(unsigned long *val)
{
	unsigned long size;
	char *line = NULL;
	size_t linelen = 0;
	FILE *f = fopen("/proc/meminfo", "r");
	int ret = -1;

	if (!f)
		return ret;

	while (getline(&line, &linelen, f) > 0) {
		if (sscanf(line, "HardwareCorrupted: %12lu kB", &size) == 1) {
			*val = size;
			ret = 0;
			break;
		}
	}

	free(line);
	fclose(f);
	return ret;
}

int unpoison_memory(unsigned long pfn)
{
	int unpoison_fd, len;
	char buf[32];
	ssize_t ret;

	unpoison_fd = open("/sys/kernel/debug/hwpoison/unpoison-pfn", O_WRONLY);
	if (unpoison_fd < 0)
		return -errno;

	len = sprintf(buf, "0x%lx\n", pfn);
	ret = write(unpoison_fd, buf, len);
	close(unpoison_fd);

	return ret > 0 ? 0 : -errno;
}

int read_file(const char *path, char *buf, size_t buflen)
{
	int fd;
	ssize_t numread;

	fd = open(path, O_RDONLY);
	if (fd == -1)
		return 0;

	numread = read(fd, buf, buflen - 1);
	if (numread < 1) {
		close(fd);
		return 0;
	}

	buf[numread] = '\0';
	close(fd);

	return (unsigned int) numread;
}

static void __write_file(const char *path, const char *buf, size_t buflen, bool ignore_einval)
{
	int fd, saved_errno;
	ssize_t numwritten;

	if (buflen < 2)
		ksft_exit_fail_msg("Incorrect buffer len: %zu\n", buflen);

	fd = open(path, O_WRONLY);
	if (fd == -1)
		ksft_exit_fail_msg("%s open failed: %s\n", path, strerror(errno));

	numwritten = write(fd, buf, buflen - 1);
	saved_errno = errno;
	close(fd);
	errno = saved_errno;
	if (numwritten < 0) {
		if (ignore_einval && errno == EINVAL)
			return;
		ksft_exit_fail_msg("%s write(%.*s) failed: %s\n", path, (int)(buflen - 1),
				buf, strerror(errno));
	}
	if (numwritten != buflen - 1)
		ksft_exit_fail_msg("%s write(%.*s) is truncated, expected %zu bytes, got %zd bytes\n",
				path, (int)(buflen - 1), buf, buflen - 1, numwritten);
}

void write_file(const char *path, const char *buf, size_t buflen)
{
	__write_file(path, buf, buflen, /* ignore_einval = */ false);
}

unsigned long read_num(const char *path)
{
	char buf[21];

	if (!read_file(path, buf, sizeof(buf)))
		ksft_exit_fail_perror("read_file()");

	return strtoul(buf, NULL, 10);
}

static void __write_num(const char *path, unsigned long num, bool ignore_einval)
{
	char buf[21];

	sprintf(buf, "%lu", num);
	__write_file(path, buf, strlen(buf) + 1, ignore_einval);
}

void write_num(const char *path, unsigned long num)
{
	return __write_num(path, num, /* ignore_einval = */ false);
}

void write_num_ignore_einval(const char *path, unsigned long num)
{
	return __write_num(path, num, /* ignore_einval = */ true);
}

static unsigned long shmall, shmmax;

void __shm_limits_restore(void)
{
	if (shmmax)
		write_num("/proc/sys/kernel/shmmax", shmmax);
	if (shmall)
		write_num("/proc/sys/kernel/shmall", shmall);
}

void shm_limits_prepare(unsigned long length)
{
	unsigned long nr = length / psize();
	unsigned long val;

	val = read_num("/proc/sys/kernel/shmmax");
	if (val < length) {
		write_num("/proc/sys/kernel/shmmax", length);
		shmmax = val;
	}

	val = read_num("/proc/sys/kernel/shmall");
	if (val < nr) {
		write_num("/proc/sys/kernel/shmall", nr);
		shmall = val;
	}
}