// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <test_progs.h>
#include <string.h>
#include <stdio.h>
#include "rhash.skel.h"
#include "bpf_iter_bpf_rhash_map.skel.h"
#include <linux/bpf.h>
#include <linux/perf_event.h>
#include <sys/syscall.h>

static void rhash_run(const char *prog_name)
{
	struct rhash *skel;
	struct bpf_program *prog;
	LIBBPF_OPTS(bpf_test_run_opts, opts);
	int err;

	skel = rhash__open();
	if (!ASSERT_OK_PTR(skel, "rhash__open"))
		return;

	prog = bpf_object__find_program_by_name(skel->obj, prog_name);
	if (!ASSERT_OK_PTR(prog, "bpf_object__find_program_by_name"))
		goto cleanup;
	bpf_program__set_autoload(prog, true);

	err = rhash__load(skel);
	if (!ASSERT_OK(err, "skel_load"))
		goto cleanup;

	err = bpf_prog_test_run_opts(bpf_program__fd(prog), &opts);
	if (!ASSERT_OK(err, "prog run"))
		goto cleanup;

	if (!ASSERT_OK(opts.retval, "prog retval"))
		goto cleanup;

	if (!ASSERT_OK(skel->bss->err, "bss->err"))
		goto cleanup;

cleanup:
	rhash__destroy(skel);
}

static int rhash_map_create(__u32 max_entries, __u64 map_extra)
{
	LIBBPF_OPTS(bpf_map_create_opts, opts,
		    .map_flags = BPF_F_NO_PREALLOC,
		    .map_extra = map_extra);

	return bpf_map_create(BPF_MAP_TYPE_RHASH, "rhash_extra",
			      sizeof(__u32), sizeof(__u64), max_entries, &opts);
}

static void rhash_map_extra_presize(void)
{
	const __u32 max_entries = 1024;
	const __u32 nelem_hint = 256;
	struct bpf_map_info info = {};
	__u32 info_len = sizeof(info);
	__u64 val = 0;
	__u32 key;
	int fd, i;

	fd = rhash_map_create(max_entries, nelem_hint);
	if (!ASSERT_GE(fd, 0, "rhash_map_create presize"))
		return;

	if (!ASSERT_OK(bpf_map_get_info_by_fd(fd, &info, &info_len), "info"))
		goto close;
	ASSERT_EQ(info.map_extra, nelem_hint, "info.map_extra");

	for (i = 0; i < (int)nelem_hint; i++) {
		key = i;
		if (!ASSERT_OK(bpf_map_update_elem(fd, &key, &val, BPF_NOEXIST),
			       "update"))
			goto close;
	}
close:
	close(fd);
}

static void rhash_map_extra_too_big(void)
{
	int fd;

	fd = rhash_map_create(1U << 20, 0x10000);
	if (!ASSERT_LT(fd, 0, "rhash_map_create hint > U16_MAX"))
		close(fd);
}

static void rhash_iter_test(void)
{
	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
	struct bpf_iter_bpf_rhash_map *skel;
	int err, i, len, map_fd, iter_fd;
	union bpf_iter_link_info linfo;
	u32 expected_key_sum = 0, key;
	struct bpf_link *link;
	u64 val = 0;
	char buf[64];

	skel = bpf_iter_bpf_rhash_map__open();
	if (!ASSERT_OK_PTR(skel, "bpf_iter_bpf_rhash_map__open"))
		return;

	err = bpf_iter_bpf_rhash_map__load(skel);
	if (!ASSERT_OK(err, "bpf_iter_bpf_rhash_map__load"))
		goto out;

	map_fd = bpf_map__fd(skel->maps.rhashmap);

	/* Populate map with test data */
	for (i = 0; i < 64; i++) {
		key = i + 1;
		expected_key_sum += key;

		err = bpf_map_update_elem(map_fd, &key, &val, BPF_NOEXIST);
		if (!ASSERT_OK(err, "map_update"))
			goto out;
	}

	memset(&linfo, 0, sizeof(linfo));
	linfo.map.map_fd = map_fd;
	opts.link_info = &linfo;
	opts.link_info_len = sizeof(linfo);

	link = bpf_program__attach_iter(skel->progs.dump_bpf_rhash_map, &opts);
	if (!ASSERT_OK_PTR(link, "attach_iter"))
		goto out;

	iter_fd = bpf_iter_create(bpf_link__fd(link));
	if (!ASSERT_GE(iter_fd, 0, "create_iter"))
		goto free_link;

	do {
		len = read(iter_fd, buf, sizeof(buf));
	} while (len > 0);

	ASSERT_EQ(skel->bss->key_sum, expected_key_sum, "key_sum");
	ASSERT_EQ(skel->bss->elem_count, 64, "elem_count");

	close(iter_fd);

free_link:
	bpf_link__destroy(link);
out:
	bpf_iter_bpf_rhash_map__destroy(skel);
}

void test_rhash(void)
{
	if (test__start_subtest("test_rhash_lookup_update"))
		rhash_run("test_rhash_lookup_update");

	if (test__start_subtest("test_rhash_update_delete"))
		rhash_run("test_rhash_update_delete");

	if (test__start_subtest("test_rhash_update_elements"))
		rhash_run("test_rhash_update_elements");

	if (test__start_subtest("test_rhash_update_exist"))
		rhash_run("test_rhash_update_exist");

	if (test__start_subtest("test_rhash_update_any"))
		rhash_run("test_rhash_update_any");

	if (test__start_subtest("test_rhash_noexist_duplicate"))
		rhash_run("test_rhash_noexist_duplicate");

	if (test__start_subtest("test_rhash_delete_nonexistent"))
		rhash_run("test_rhash_delete_nonexistent");

	if (test__start_subtest("test_rhash_map_extra_presize"))
		rhash_map_extra_presize();

	if (test__start_subtest("test_rhash_map_extra_too_big"))
		rhash_map_extra_too_big();

	if (test__start_subtest("test_rhash_iter"))
		rhash_iter_test();
}