// SPDX-License-Identifier: GPL-2.0
/*
 * Test cases for binder allocator code.
 *
 * Copyright 2025 Google LLC.
 * Author: Tiffany Yang <ynaffit@google.com>
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <kunit/test.h>
#include <linux/anon_inodes.h>
#include <linux/err.h>
#include <linux/file.h>
#include <linux/fs.h>
#include <linux/mm.h>
#include <linux/mman.h>
#include <linux/seq_buf.h>
#include <linux/sizes.h>

#include "../binder_alloc.h"
#include "../binder_internal.h"

MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");

#define BINDER_MMAP_SIZE SZ_128K

#define BUFFER_NUM 5
#define BUFFER_MIN_SIZE (PAGE_SIZE / 8)

#define FREESEQ_BUFLEN ((3 * BUFFER_NUM) + 1)

#define ALIGN_TYPE_STRLEN (12)

#define ALIGNMENTS_BUFLEN (((ALIGN_TYPE_STRLEN + 6) * BUFFER_NUM) + 1)

#define PRINT_ALL_CASES (0)

/* 5^5 alignment combinations * 2 places to share pages * 5! free sequences */
#define TOTAL_EXHAUSTIVE_CASES (3125 * 2 * 120)

/**
 * enum buf_end_align_type - Page alignment of a buffer
 * end with regard to the end of the previous buffer.
 *
 * In the pictures below, buf2 refers to the buffer we
 * are aligning. buf1 refers to previous buffer by addr.
 * Symbol [ means the start of a buffer, ] means the end
 * of a buffer, and | means page boundaries.
 */
enum buf_end_align_type {
        /**
         * @SAME_PAGE_UNALIGNED: The end of this buffer is on
         * the same page as the end of the previous buffer and
         * is not page aligned. Examples:
         * buf1 ][ buf2 ][ ...
         * buf1 ]|[ buf2 ][ ...
         */
        SAME_PAGE_UNALIGNED = 0,
        /**
         * @SAME_PAGE_ALIGNED: When the end of the previous buffer
         * is not page aligned, the end of this buffer is on the
         * same page as the end of the previous buffer and is page
         * aligned. When the previous buffer is page aligned, the
         * end of this buffer is aligned to the next page boundary.
         * Examples:
         * buf1 ][ buf2 ]| ...
         * buf1 ]|[ buf2 ]| ...
         */
        SAME_PAGE_ALIGNED,
        /**
         * @NEXT_PAGE_UNALIGNED: The end of this buffer is on
         * the page next to the end of the previous buffer and
         * is not page aligned. Examples:
         * buf1 ][ buf2 | buf2 ][ ...
         * buf1 ]|[ buf2 | buf2 ][ ...
         */
        NEXT_PAGE_UNALIGNED,
        /**
         * @NEXT_PAGE_ALIGNED: The end of this buffer is on
         * the page next to the end of the previous buffer and
         * is page aligned. Examples:
         * buf1 ][ buf2 | buf2 ]| ...
         * buf1 ]|[ buf2 | buf2 ]| ...
         */
        NEXT_PAGE_ALIGNED,
        /**
         * @NEXT_NEXT_UNALIGNED: The end of this buffer is on
         * the page that follows the page after the end of the
         * previous buffer and is not page aligned. Examples:
         * buf1 ][ buf2 | buf2 | buf2 ][ ...
         * buf1 ]|[ buf2 | buf2 | buf2 ][ ...
         */
        NEXT_NEXT_UNALIGNED,
        /**
         * @LOOP_END: The number of enum values in &buf_end_align_type.
         * It is used for controlling loop termination.
         */
        LOOP_END,
};

static const char *const buf_end_align_type_strs[LOOP_END] = {
        [SAME_PAGE_UNALIGNED] = "SP_UNALIGNED",
        [SAME_PAGE_ALIGNED]   = " SP_ALIGNED ",
        [NEXT_PAGE_UNALIGNED] = "NP_UNALIGNED",
        [NEXT_PAGE_ALIGNED]   = " NP_ALIGNED ",
        [NEXT_NEXT_UNALIGNED] = "NN_UNALIGNED",
};

struct binder_alloc_test_case_info {
        char alignments[ALIGNMENTS_BUFLEN];
        struct seq_buf alignments_sb;
        size_t *buffer_sizes;
        int *free_sequence;
        bool front_pages;
};

static void stringify_free_seq(struct kunit *test, int *seq, struct seq_buf *sb)
{
        int i;

        for (i = 0; i < BUFFER_NUM; i++)
                seq_buf_printf(sb, "[%d]", seq[i]);

        KUNIT_EXPECT_FALSE(test, seq_buf_has_overflowed(sb));
}

static void stringify_alignments(struct kunit *test, int *alignments,
                                 struct seq_buf *sb)
{
        int i;

        for (i = 0; i < BUFFER_NUM; i++)
                seq_buf_printf(sb, "[ %d:%s ]", i,
                               buf_end_align_type_strs[alignments[i]]);

        KUNIT_EXPECT_FALSE(test, seq_buf_has_overflowed(sb));
}

static bool check_buffer_pages_allocated(struct kunit *test,
                                         struct binder_alloc *alloc,
                                         struct binder_buffer *buffer,
                                         size_t size)
{
        unsigned long page_addr;
        unsigned long end;
        int page_index;

        end = PAGE_ALIGN(buffer->user_data + size);
        page_addr = buffer->user_data;
        for (; page_addr < end; page_addr += PAGE_SIZE) {
                page_index = (page_addr - alloc->vm_start) / PAGE_SIZE;
                if (!alloc->pages[page_index] ||
                    !list_empty(page_to_lru(alloc->pages[page_index]))) {
                        kunit_err(test, "expect alloc but is %s at page index %d\n",
                                  alloc->pages[page_index] ?
                                  "lru" : "free", page_index);
                        return false;
                }
        }
        return true;
}

static unsigned long binder_alloc_test_alloc_buf(struct kunit *test,
                                                 struct binder_alloc *alloc,
                                                 struct binder_buffer *buffers[],
                                                 size_t *sizes, int *seq)
{
        unsigned long failures = 0;
        int i;

        for (i = 0; i < BUFFER_NUM; i++) {
                buffers[i] = binder_alloc_new_buf(alloc, sizes[i], 0, 0, 0);
                if (IS_ERR(buffers[i]) ||
                    !check_buffer_pages_allocated(test, alloc, buffers[i], sizes[i]))
                        failures++;
        }

        return failures;
}

static unsigned long binder_alloc_test_free_buf(struct kunit *test,
                                                struct binder_alloc *alloc,
                                                struct binder_buffer *buffers[],
                                                size_t *sizes, int *seq, size_t end)
{
        unsigned long failures = 0;
        int i;

        for (i = 0; i < BUFFER_NUM; i++)
                binder_alloc_free_buf(alloc, buffers[seq[i]]);

        for (i = 0; i <= (end - 1) / PAGE_SIZE; i++) {
                if (list_empty(page_to_lru(alloc->pages[i]))) {
                        kunit_err(test, "expect lru but is %s at page index %d\n",
                                  alloc->pages[i] ? "alloc" : "free", i);
                        failures++;
                }
        }

        return failures;
}

static unsigned long binder_alloc_test_free_page(struct kunit *test,
                                                 struct binder_alloc *alloc)
{
        unsigned long failures = 0;
        unsigned long count;
        int i;

        while ((count = list_lru_count(alloc->freelist))) {
                list_lru_walk(alloc->freelist, binder_alloc_free_page,
                              NULL, count);
        }

        for (i = 0; i < (alloc->buffer_size / PAGE_SIZE); i++) {
                if (alloc->pages[i]) {
                        kunit_err(test, "expect free but is %s at page index %d\n",
                                  list_empty(page_to_lru(alloc->pages[i])) ?
                                  "alloc" : "lru", i);
                        failures++;
                }
        }

        return failures;
}

/* Executes one full test run for the given test case. */
static bool binder_alloc_test_alloc_free(struct kunit *test,
                                         struct binder_alloc *alloc,
                                         struct binder_alloc_test_case_info *tc,
                                         size_t end)
{
        unsigned long pages = PAGE_ALIGN(end) / PAGE_SIZE;
        struct binder_buffer *buffers[BUFFER_NUM];
        unsigned long failures;
        bool failed = false;

        failures = binder_alloc_test_alloc_buf(test, alloc, buffers,
                                               tc->buffer_sizes,
                                               tc->free_sequence);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "Initial allocation failed: %lu/%u buffers with errors",
                            failures, BUFFER_NUM);

        failures = binder_alloc_test_free_buf(test, alloc, buffers,
                                              tc->buffer_sizes,
                                              tc->free_sequence, end);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "Initial buffers not freed correctly: %lu/%lu pages not on lru list",
                            failures, pages);

        /* Allocate from lru. */
        failures = binder_alloc_test_alloc_buf(test, alloc, buffers,
                                               tc->buffer_sizes,
                                               tc->free_sequence);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "Reallocation failed: %lu/%u buffers with errors",
                            failures, BUFFER_NUM);

        failures = list_lru_count(alloc->freelist);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "lru list should be empty after reallocation but still has %lu pages",
                            failures);

        failures = binder_alloc_test_free_buf(test, alloc, buffers,
                                              tc->buffer_sizes,
                                              tc->free_sequence, end);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "Reallocated buffers not freed correctly: %lu/%lu pages not on lru list",
                            failures, pages);

        failures = binder_alloc_test_free_page(test, alloc);
        failed = failed || failures;
        KUNIT_EXPECT_EQ_MSG(test, failures, 0,
                            "Failed to clean up allocated pages: %lu/%lu pages still installed",
                            failures, (alloc->buffer_size / PAGE_SIZE));

        return failed;
}

static bool is_dup(int *seq, int index, int val)
{
        int i;

        for (i = 0; i < index; i++) {
                if (seq[i] == val)
                        return true;
        }
        return false;
}

/* Generate BUFFER_NUM factorial free orders. */
static void permute_frees(struct kunit *test, struct binder_alloc *alloc,
                          struct binder_alloc_test_case_info *tc,
                          unsigned long *runs, unsigned long *failures,
                          int index, size_t end)
{
        bool case_failed;
        int i;

        if (index == BUFFER_NUM) {
                DECLARE_SEQ_BUF(freeseq_sb, FREESEQ_BUFLEN);

                case_failed = binder_alloc_test_alloc_free(test, alloc, tc, end);
                *runs += 1;
                *failures += case_failed;

                if (case_failed || PRINT_ALL_CASES) {
                        stringify_free_seq(test, tc->free_sequence,
                                           &freeseq_sb);
                        kunit_err(test, "case %lu: [%s] | %s - %s - %s", *runs,
                                  case_failed ? "FAILED" : "PASSED",
                                  tc->front_pages ? "front" : "back ",
                                  seq_buf_str(&tc->alignments_sb),
                                  seq_buf_str(&freeseq_sb));
                }

                return;
        }
        for (i = 0; i < BUFFER_NUM; i++) {
                if (is_dup(tc->free_sequence, index, i))
                        continue;
                tc->free_sequence[index] = i;
                permute_frees(test, alloc, tc, runs, failures, index + 1, end);
        }
}

static void gen_buf_sizes(struct kunit *test,
                          struct binder_alloc *alloc,
                          struct binder_alloc_test_case_info *tc,
                          size_t *end_offset, unsigned long *runs,
                          unsigned long *failures)
{
        size_t last_offset, offset = 0;
        size_t front_sizes[BUFFER_NUM];
        size_t back_sizes[BUFFER_NUM];
        int seq[BUFFER_NUM] = {0};
        int i;

        tc->free_sequence = seq;
        for (i = 0; i < BUFFER_NUM; i++) {
                last_offset = offset;
                offset = end_offset[i];
                front_sizes[i] = offset - last_offset;
                back_sizes[BUFFER_NUM - i - 1] = front_sizes[i];
        }
        back_sizes[0] += alloc->buffer_size - end_offset[BUFFER_NUM - 1];

        /*
         * Buffers share the first or last few pages.
         * Only BUFFER_NUM - 1 buffer sizes are adjustable since
         * we need one giant buffer before getting to the last page.
         */
        tc->front_pages = true;
        tc->buffer_sizes = front_sizes;
        permute_frees(test, alloc, tc, runs, failures, 0,
                      end_offset[BUFFER_NUM - 1]);

        tc->front_pages = false;
        tc->buffer_sizes = back_sizes;
        permute_frees(test, alloc, tc, runs, failures, 0, alloc->buffer_size);
}

static void gen_buf_offsets(struct kunit *test, struct binder_alloc *alloc,
                            size_t *end_offset, int *alignments,
                            unsigned long *runs, unsigned long *failures,
                            int index)
{
        size_t end, prev;
        int align;

        if (index == BUFFER_NUM) {
                struct binder_alloc_test_case_info tc = {0};

                seq_buf_init(&tc.alignments_sb, tc.alignments,
                             ALIGNMENTS_BUFLEN);
                stringify_alignments(test, alignments, &tc.alignments_sb);

                gen_buf_sizes(test, alloc, &tc, end_offset, runs, failures);
                return;
        }
        prev = index == 0 ? 0 : end_offset[index - 1];
        end = prev;

        BUILD_BUG_ON(BUFFER_MIN_SIZE * BUFFER_NUM >= PAGE_SIZE);

        for (align = SAME_PAGE_UNALIGNED; align < LOOP_END; align++) {
                if (align % 2)
                        end = ALIGN(end, PAGE_SIZE);
                else
                        end += BUFFER_MIN_SIZE;
                end_offset[index] = end;
                alignments[index] = align;
                gen_buf_offsets(test, alloc, end_offset, alignments, runs,
                                failures, index + 1);
        }
}

struct binder_alloc_test {
        struct binder_alloc alloc;
        struct list_lru binder_test_freelist;
        struct file *filp;
        unsigned long mmap_uaddr;
};

static void binder_alloc_test_init_freelist(struct kunit *test)
{
        struct binder_alloc_test *priv = test->priv;

        KUNIT_EXPECT_PTR_EQ(test, priv->alloc.freelist,
                            &priv->binder_test_freelist);
}

static void binder_alloc_test_mmap(struct kunit *test)
{
        struct binder_alloc_test *priv = test->priv;
        struct binder_alloc *alloc = &priv->alloc;
        struct binder_buffer *buf;
        struct rb_node *n;

        KUNIT_EXPECT_EQ(test, alloc->mapped, true);
        KUNIT_EXPECT_EQ(test, alloc->buffer_size, BINDER_MMAP_SIZE);

        n = rb_first(&alloc->allocated_buffers);
        KUNIT_EXPECT_PTR_EQ(test, n, NULL);

        n = rb_first(&alloc->free_buffers);
        buf = rb_entry(n, struct binder_buffer, rb_node);
        KUNIT_EXPECT_EQ(test, binder_alloc_buffer_size(alloc, buf),
                        BINDER_MMAP_SIZE);
        KUNIT_EXPECT_TRUE(test, list_is_last(&buf->entry, &alloc->buffers));
}

/**
 * binder_alloc_exhaustive_test() - Exhaustively test alloc and free of buffer pages.
 * @test: The test context object.
 *
 * Allocate BUFFER_NUM buffers to cover all page alignment cases,
 * then free them in all orders possible. Check that pages are
 * correctly allocated, put onto lru when buffers are freed, and
 * are freed when binder_alloc_free_page() is called.
 */
static void binder_alloc_exhaustive_test(struct kunit *test)
{
        struct binder_alloc_test *priv = test->priv;
        size_t end_offset[BUFFER_NUM];
        int alignments[BUFFER_NUM];
        unsigned long failures = 0;
        unsigned long runs = 0;

        gen_buf_offsets(test, &priv->alloc, end_offset, alignments, &runs,
                        &failures, 0);

        KUNIT_EXPECT_EQ(test, runs, TOTAL_EXHAUSTIVE_CASES);
        KUNIT_EXPECT_EQ(test, failures, 0);
}

/* ===== End test cases ===== */

static void binder_alloc_test_vma_close(struct vm_area_struct *vma)
{
        struct binder_alloc *alloc = vma->vm_private_data;

        binder_alloc_vma_close(alloc);
}

static const struct vm_operations_struct binder_alloc_test_vm_ops = {
        .close = binder_alloc_test_vma_close,
        .fault = binder_vm_fault,
};

static int binder_alloc_test_mmap_handler(struct file *filp,
                                          struct vm_area_struct *vma)
{
        struct binder_alloc *alloc = filp->private_data;

        vm_flags_mod(vma, VM_DONTCOPY | VM_MIXEDMAP, VM_MAYWRITE);

        vma->vm_ops = &binder_alloc_test_vm_ops;
        vma->vm_private_data = alloc;

        return binder_alloc_mmap_handler(alloc, vma);
}

static const struct file_operations binder_alloc_test_fops = {
        .mmap = binder_alloc_test_mmap_handler,
};

static int binder_alloc_test_init(struct kunit *test)
{
        struct binder_alloc_test *priv;
        int ret;

        priv = kunit_kzalloc(test, sizeof(*priv), GFP_KERNEL);
        if (!priv)
                return -ENOMEM;
        test->priv = priv;

        ret = list_lru_init(&priv->binder_test_freelist);
        if (ret) {
                kunit_err(test, "Failed to initialize test freelist\n");
                return ret;
        }

        /* __binder_alloc_init requires mm to be attached */
        ret = kunit_attach_mm();
        if (ret) {
                kunit_err(test, "Failed to attach mm\n");
                return ret;
        }
        __binder_alloc_init(&priv->alloc, &priv->binder_test_freelist);

        priv->filp = anon_inode_getfile("binder_alloc_kunit",
                                        &binder_alloc_test_fops, &priv->alloc,
                                        O_RDWR | O_CLOEXEC);
        if (IS_ERR_OR_NULL(priv->filp)) {
                kunit_err(test, "Failed to open binder alloc test driver file\n");
                return priv->filp ? PTR_ERR(priv->filp) : -ENOMEM;
        }

        priv->mmap_uaddr = kunit_vm_mmap(test, priv->filp, 0, BINDER_MMAP_SIZE,
                                         PROT_READ, MAP_PRIVATE | MAP_NORESERVE,
                                         0);
        if (!priv->mmap_uaddr) {
                kunit_err(test, "Could not map the test's transaction memory\n");
                return -ENOMEM;
        }

        return 0;
}

static void binder_alloc_test_exit(struct kunit *test)
{
        struct binder_alloc_test *priv = test->priv;

        /* Close the backing file to make sure binder_alloc_vma_close runs */
        if (!IS_ERR_OR_NULL(priv->filp))
                fput(priv->filp);

        if (priv->alloc.mm)
                binder_alloc_deferred_release(&priv->alloc);

        /* Make sure freelist is empty */
        KUNIT_EXPECT_EQ(test, list_lru_count(&priv->binder_test_freelist), 0);
        list_lru_destroy(&priv->binder_test_freelist);
}

static struct kunit_case binder_alloc_test_cases[] = {
        KUNIT_CASE(binder_alloc_test_init_freelist),
        KUNIT_CASE(binder_alloc_test_mmap),
        KUNIT_CASE_SLOW(binder_alloc_exhaustive_test),
        {}
};

static struct kunit_suite binder_alloc_test_suite = {
        .name = "binder_alloc",
        .test_cases = binder_alloc_test_cases,
        .init = binder_alloc_test_init,
        .exit = binder_alloc_test_exit,
};

kunit_test_suite(binder_alloc_test_suite);

MODULE_AUTHOR("Tiffany Yang <ynaffit@google.com>");
MODULE_DESCRIPTION("Binder Alloc KUnit tests");
MODULE_LICENSE("GPL");