// 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(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");