/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2016 Joyent, Inc.
 * Copyright 2024 Oxide Computer Company
 */

/*
 * Basic tests for aligned_alloc(3C). Note that we test ENOMEM failure by
 * relying on the implementation of the current libc malloc. Specifically we go
 * through and add a mapping so we can't expand the heap and then use it up. If
 * the memory allocator is ever changed, this test will start failing, at which
 * point, it may not be worth the cost of keeping it around.
 */

#include <stdlib.h>
#include <errno.h>
#include <libproc.h>
#include <stdio.h>
#include <stdalign.h>
#include <err.h>
#include <sys/sysmacros.h>
#include <sys/mman.h>
#include <sys/debug.h>
#include <stdbool.h>
#include <string.h>

typedef struct {
        size_t at_size;
        size_t at_align;
        int at_errno;
        const char *at_desc;
} alloc_test_t;

static const alloc_test_t alloc_tests[] = {
        { 0, sizeof (long), EINVAL, "zero alignment fails with EINVAL" },
        { sizeof (long), 0, EINVAL, "zero size fails with EINVAL" },
        { 128, 3, EINVAL, "3-byte alignment fails with EINVAL" },
        { 128, 7777, EINVAL, "7777-byte alignment fails with EINVAL" },
        { 128, 23, EINVAL, "23-byte alignment fails with EINVAL" },
        { sizeof (char), alignof (char), 0, "alignof (char), 1 byte" },
        { 5, alignof (char), 0, "alignof (char), multiple bytes" },
        { 1, alignof (short), 0, "alignof (short), 1 byte" },
        { 16, alignof (short), 0, "alignof (short), 16 byte" },
        { 1, alignof (int), 0, "alignof (int), 1 byte" },
        { 4, alignof (int), 0, "alignof (int), 4 bytes" },
        { 22, alignof (int), 0, "alignof (int), 22 bytes" },
        /* We skip long here because it varies between ILP32/LP64 */
        { 7, alignof (long long), 0, "alignof (long long), 7 bytes" },
        { 128, alignof (long long), 0, "alignof (long long), 128 bytes" },
        { 511, alignof (long long), 0, "alignof (long long), 511 bytes" },
        { 16, 16, 0, "16-byte alignment), 16 bytes" },
        { 256, 16, 0, "16-byte alignment), 256 bytes" },
        { 256, 4096, 0, "4096-byte alignment), 256 bytes" },
        { 4096, 4096, 0, "4096-byte alignment), 4096 bytes" },
};

/*
 * Disable the per-thread caches and enable debugging if launched with umem.
 */
const char *
_umem_debug_init(void)
{
        return ("default,verbose");
}

const char *
_umem_options_init(void)
{
        return ("perthreadcache=0");
}

static bool
alloc_test_one(const alloc_test_t *test)
{
        bool ret = false;
        void *buf = aligned_alloc(test->at_align, test->at_size);

        if (buf == NULL) {
                if (test->at_errno == 0) {
                        warnx("TEST FAILED: %s: allocation failed with %s, but "
                            "expected success", test->at_desc,
                            strerrorname_np(errno));
                } else if (errno != test->at_errno) {
                        warnx("TEST FAILED: %s: allocation failed with %s, but "
                            "expected errno %s", test->at_desc,
                            strerrorname_np(errno),
                            strerrorname_np(test->at_errno));
                } else {
                        (void) printf("TEST PASSED: %s\n", test->at_desc);
                        ret = true;
                }
        } else if (test->at_errno != 0) {
                warnx("TEST FAILED: %s: allocation succeeded, but expected "
                    "errno %s", test->at_desc, strerrorname_np(test->at_errno));
        } else {
                (void) printf("TEST PASSED: %s\n", test->at_desc);
                ret = true;
        }

        free(buf);
        return (ret);
}

static bool
alloc_test_enomem(const alloc_test_t *test)
{
        bool ret = false;
        void *buf = aligned_alloc(test->at_align, test->at_size);

        if (buf != NULL) {
                warnx("TEST FAILED: %s (forced ENOMEM/EAGAIN): succeeded, but "
                    "expected ENOMEM", test->at_desc);
        } else if (errno != ENOMEM && errno != EAGAIN) {
                warnx("TEST FAILED: %s (forced ENOMEM/EAGAIN): failed with %s, "
                    "but expected ENOMEM", test->at_desc,
                    strerrordesc_np(errno));
        } else {
                (void) printf("TEST PASSED: %s: ENOMEM/EAGAIN forced\n",
                    test->at_desc);
                ret = true;
        }

        return (ret);
}

static void
libc_enomem(void)
{
        pstatus_t status;

        VERIFY0(proc_get_status(getpid(), &status));
        VERIFY3P(mmap((caddr_t)P2ROUNDUP(status.pr_brkbase +
            status.pr_brksize, 0x1000), 0x1000,
            PROT_READ, MAP_ANON | MAP_FIXED | MAP_PRIVATE, -1, 0),
            !=, (void *)-1);

        for (;;) {
                if (malloc(16) == NULL)
                        break;
        }

        for (;;) {
                if (aligned_alloc(sizeof (void *), 16) == NULL)
                        break;
        }
}

/*
 * Because this test is leveraging LD_PRELOAD in the test runner to test
 * different malloc libraries, we can't call umem_setmtbf() directly. Instead we
 * ask rtld to find it for us, which is a bit gross, but works.
 */
static void
libumem_enomem(void)
{
        void (*mtbf)(uint32_t) = dlsym(RTLD_DEFAULT, "umem_setmtbf");
        if (mtbf == NULL) {
                errx(EXIT_FAILURE, "failed to find umem_setmtbf: %s",
                    dlerror());
        }

        mtbf(1);
}

int
main(void)
{
        const char *preload;
        int ret = EXIT_SUCCESS;

        for (size_t i = 0; i < ARRAY_SIZE(alloc_tests); i++) {
                if (!alloc_test_one(&alloc_tests[i])) {
                        ret = EXIT_FAILURE;
                }
        }

        /*
         * To catch failure tests, we need to know what memory allocator we're
         * using which we expect to be indicated via an LD_PRELOAD environment
         * variable. If it's not set, assume libc.
         */
        preload = getenv("LD_PRELOAD");
        if (preload != NULL) {
                if (strstr(preload, "umem") != NULL) {
                        libumem_enomem();
                } else {
                        warnx("malloc(3C) library %s not supported, skipping "
                            "ENOMEM tests", preload);
                        goto skip;
                }
        } else {
                libc_enomem();
        }

        for (size_t i = 0; i < ARRAY_SIZE(alloc_tests); i++) {
                if (alloc_tests[i].at_errno != 0)
                        continue;
                if (!alloc_test_enomem(&alloc_tests[i])) {
                        ret = EXIT_FAILURE;
                }
        }

skip:
        if (ret == EXIT_SUCCESS) {
                (void) printf("All tests passed successfully\n");
        }

        return (ret);
}