// SPDX-License-Identifier: GPL-2.0
/*
 * Basic tests for PR_GET/SET_THP_DISABLE prctl calls
 *
 * Author(s): Usama Arif <usamaarif642@gmail.com>
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/mman.h>
#include <linux/mman.h>
#include <sys/prctl.h>
#include <sys/wait.h>

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

#ifndef PR_THP_DISABLE_EXCEPT_ADVISED
#define PR_THP_DISABLE_EXCEPT_ADVISED (1 << 1)
#endif

enum thp_collapse_type {
        THP_COLLAPSE_NONE,
        THP_COLLAPSE_MADV_NOHUGEPAGE,
        THP_COLLAPSE_MADV_HUGEPAGE,     /* MADV_HUGEPAGE before access */
        THP_COLLAPSE_MADV_COLLAPSE,     /* MADV_COLLAPSE after access */
};

/*
 * Function to mmap a buffer, fault it in, madvise it appropriately (before
 * page fault for MADV_HUGE, and after for MADV_COLLAPSE), and check if the
 * mmap region is huge.
 * Returns:
 * 0 if test doesn't give hugepage
 * 1 if test gives a hugepage
 * -errno if mmap fails
 */
static int test_mmap_thp(enum thp_collapse_type madvise_buf, size_t pmdsize)
{
        char *mem, *mmap_mem;
        size_t mmap_size;
        int ret;

        /* For alignment purposes, we need twice the THP size. */
        mmap_size = 2 * pmdsize;
        mmap_mem = (char *)mmap(NULL, mmap_size, PROT_READ | PROT_WRITE,
                                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (mmap_mem == MAP_FAILED)
                return -errno;

        /* We need a THP-aligned memory area. */
        mem = (char *)(((uintptr_t)mmap_mem + pmdsize) & ~(pmdsize - 1));

        if (madvise_buf == THP_COLLAPSE_MADV_HUGEPAGE)
                madvise(mem, pmdsize, MADV_HUGEPAGE);
        else if (madvise_buf == THP_COLLAPSE_MADV_NOHUGEPAGE)
                madvise(mem, pmdsize, MADV_NOHUGEPAGE);

        /* Ensure memory is allocated */
        memset(mem, 1, pmdsize);

        if (madvise_buf == THP_COLLAPSE_MADV_COLLAPSE)
                madvise(mem, pmdsize, MADV_COLLAPSE);

        /* HACK: make sure we have a separate VMA that we can check reliably. */
        mprotect(mem, pmdsize, PROT_READ);

        ret = check_huge_anon(mem, 1, pmdsize);
        munmap(mmap_mem, mmap_size);
        return ret;
}

static void prctl_thp_disable_completely_test(struct __test_metadata *const _metadata,
                                              size_t pmdsize,
                                              enum thp_enabled thp_policy)
{
        ASSERT_EQ(prctl(PR_GET_THP_DISABLE, NULL, NULL, NULL, NULL), 1);

        /* tests after prctl overrides global policy */
        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_NONE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_NOHUGEPAGE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_HUGEPAGE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_COLLAPSE, pmdsize), 0);

        /* Reset to global policy */
        ASSERT_EQ(prctl(PR_SET_THP_DISABLE, 0, NULL, NULL, NULL), 0);

        /* tests after prctl is cleared, and only global policy is effective */
        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_NONE, pmdsize),
                  thp_policy == THP_ALWAYS ? 1 : 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_NOHUGEPAGE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_HUGEPAGE, pmdsize),
                  thp_policy == THP_NEVER ? 0 : 1);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_COLLAPSE, pmdsize), 1);
}

FIXTURE(prctl_thp_disable_completely)
{
        struct thp_settings settings;
        size_t pmdsize;
};

FIXTURE_VARIANT(prctl_thp_disable_completely)
{
        enum thp_enabled thp_policy;
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_completely, never)
{
        .thp_policy = THP_NEVER,
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_completely, madvise)
{
        .thp_policy = THP_MADVISE,
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_completely, always)
{
        .thp_policy = THP_ALWAYS,
};

FIXTURE_SETUP(prctl_thp_disable_completely)
{
        if (!thp_available())
                SKIP(return, "Transparent Hugepages not available\n");

        self->pmdsize = read_pmd_pagesize();
        if (!self->pmdsize)
                SKIP(return, "Unable to read PMD size\n");

        if (prctl(PR_SET_THP_DISABLE, 1, NULL, NULL, NULL))
                SKIP(return, "Unable to disable THPs completely for the process\n");

        thp_save_settings();
        thp_read_settings(&self->settings);
        self->settings.thp_enabled = variant->thp_policy;
        self->settings.hugepages[sz2ord(self->pmdsize, getpagesize())].enabled = THP_INHERIT;
        thp_write_settings(&self->settings);
}

FIXTURE_TEARDOWN(prctl_thp_disable_completely)
{
        thp_restore_settings();
}

TEST_F(prctl_thp_disable_completely, nofork)
{
        prctl_thp_disable_completely_test(_metadata, self->pmdsize, variant->thp_policy);
}

TEST_F(prctl_thp_disable_completely, fork)
{
        int ret = 0;
        pid_t pid;

        /* Make sure prctl changes are carried across fork */
        pid = fork();
        ASSERT_GE(pid, 0);

        if (!pid) {
                prctl_thp_disable_completely_test(_metadata, self->pmdsize, variant->thp_policy);
                return;
        }

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

static void prctl_thp_disable_except_madvise_test(struct __test_metadata *const _metadata,
                                                  size_t pmdsize,
                                                  enum thp_enabled thp_policy)
{
        ASSERT_EQ(prctl(PR_GET_THP_DISABLE, NULL, NULL, NULL, NULL), 3);

        /* tests after prctl overrides global policy */
        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_NONE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_NOHUGEPAGE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_HUGEPAGE, pmdsize),
                  thp_policy == THP_NEVER ? 0 : 1);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_COLLAPSE, pmdsize), 1);

        /* Reset to global policy */
        ASSERT_EQ(prctl(PR_SET_THP_DISABLE, 0, NULL, NULL, NULL), 0);

        /* tests after prctl is cleared, and only global policy is effective */
        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_NONE, pmdsize),
                  thp_policy == THP_ALWAYS ? 1 : 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_NOHUGEPAGE, pmdsize), 0);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_HUGEPAGE, pmdsize),
                  thp_policy == THP_NEVER ? 0 : 1);

        ASSERT_EQ(test_mmap_thp(THP_COLLAPSE_MADV_COLLAPSE, pmdsize), 1);
}

FIXTURE(prctl_thp_disable_except_madvise)
{
        struct thp_settings settings;
        size_t pmdsize;
};

FIXTURE_VARIANT(prctl_thp_disable_except_madvise)
{
        enum thp_enabled thp_policy;
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_except_madvise, never)
{
        .thp_policy = THP_NEVER,
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_except_madvise, madvise)
{
        .thp_policy = THP_MADVISE,
};

FIXTURE_VARIANT_ADD(prctl_thp_disable_except_madvise, always)
{
        .thp_policy = THP_ALWAYS,
};

FIXTURE_SETUP(prctl_thp_disable_except_madvise)
{
        if (!thp_available())
                SKIP(return, "Transparent Hugepages not available\n");

        self->pmdsize = read_pmd_pagesize();
        if (!self->pmdsize)
                SKIP(return, "Unable to read PMD size\n");

        if (prctl(PR_SET_THP_DISABLE, 1, PR_THP_DISABLE_EXCEPT_ADVISED, NULL, NULL))
                SKIP(return, "Unable to set PR_THP_DISABLE_EXCEPT_ADVISED\n");

        thp_save_settings();
        thp_read_settings(&self->settings);
        self->settings.thp_enabled = variant->thp_policy;
        self->settings.hugepages[sz2ord(self->pmdsize, getpagesize())].enabled = THP_INHERIT;
        thp_write_settings(&self->settings);
}

FIXTURE_TEARDOWN(prctl_thp_disable_except_madvise)
{
        thp_restore_settings();
}

TEST_F(prctl_thp_disable_except_madvise, nofork)
{
        prctl_thp_disable_except_madvise_test(_metadata, self->pmdsize, variant->thp_policy);
}

TEST_F(prctl_thp_disable_except_madvise, fork)
{
        int ret = 0;
        pid_t pid;

        /* Make sure prctl changes are carried across fork */
        pid = fork();
        ASSERT_GE(pid, 0);

        if (!pid) {
                prctl_thp_disable_except_madvise_test(_metadata, self->pmdsize,
                                                      variant->thp_policy);
                return;
        }

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

TEST_HARNESS_MAIN