// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
#include <test_progs.h>
#include "cgroup_helpers.h"
#include "cgroup_mprog.skel.h"

static void assert_mprog_count(int cg, int atype, int expected)
{
	__u32 count = 0, attach_flags = 0;
	int err;

	err = bpf_prog_query(cg, atype, 0, &attach_flags,
			     NULL, &count);
	ASSERT_EQ(count, expected, "count");
	ASSERT_EQ(err, 0, "prog_query");
}

static void test_prog_attach_detach(int atype)
{
	LIBBPF_OPTS(bpf_prog_attach_opts, opta);
	LIBBPF_OPTS(bpf_prog_detach_opts, optd);
	LIBBPF_OPTS(bpf_prog_query_opts, optq);
	__u32 fd1, fd2, fd3, fd4, id1, id2, id3, id4;
	struct cgroup_mprog *skel;
	__u32 prog_ids[10];
	int cg, err;

	cg = test__join_cgroup("/prog_attach_detach");
	if (!ASSERT_GE(cg, 0, "join_cgroup /prog_attach_detach"))
		return;

	skel = cgroup_mprog__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel_load"))
		goto cleanup;

	fd1 = bpf_program__fd(skel->progs.getsockopt_1);
	fd2 = bpf_program__fd(skel->progs.getsockopt_2);
	fd3 = bpf_program__fd(skel->progs.getsockopt_3);
	fd4 = bpf_program__fd(skel->progs.getsockopt_4);

	id1 = id_from_prog_fd(fd1);
	id2 = id_from_prog_fd(fd2);
	id3 = id_from_prog_fd(fd3);
	id4 = id_from_prog_fd(fd4);

	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE | BPF_F_AFTER,
		.expected_revision = 1,
	);

	/* ordering: [fd1] */
	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup;

	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE,
		.expected_revision = 2,
	);

	/* ordering: [fd2, fd1] */
	err = bpf_prog_attach_opts(fd2, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup1;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER,
		.relative_fd = fd2,
		.expected_revision = 3,
	);

	/* ordering: [fd2, fd3, fd1] */
	err = bpf_prog_attach_opts(fd3, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 3);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI,
		.expected_revision = 4,
	);

	/* ordering: [fd2, fd3, fd1, fd4] */
	err = bpf_prog_attach_opts(fd4, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup3;

	assert_mprog_count(cg, atype, 4);

	/* retrieve optq.prog_cnt */
	err = bpf_prog_query_opts(cg, atype, &optq);
	if (!ASSERT_OK(err, "prog_query"))
		goto cleanup4;

	/* optq.prog_cnt will be used in below query */
	memset(prog_ids, 0, sizeof(prog_ids));
	optq.prog_ids = prog_ids;
	err = bpf_prog_query_opts(cg, atype, &optq);
	if (!ASSERT_OK(err, "prog_query"))
		goto cleanup4;

	ASSERT_EQ(optq.count, 4, "count");
	ASSERT_EQ(optq.revision, 5, "revision");
	ASSERT_EQ(optq.prog_ids[0], id2, "prog_ids[0]");
	ASSERT_EQ(optq.prog_ids[1], id3, "prog_ids[1]");
	ASSERT_EQ(optq.prog_ids[2], id1, "prog_ids[2]");
	ASSERT_EQ(optq.prog_ids[3], id4, "prog_ids[3]");
	ASSERT_EQ(optq.prog_ids[4], 0, "prog_ids[4]");
	ASSERT_EQ(optq.link_ids, NULL, "link_ids");

cleanup4:
	optd.expected_revision = 5;
	err = bpf_prog_detach_opts(fd4, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 3);

cleanup3:
	LIBBPF_OPTS_RESET(optd);
	err = bpf_prog_detach_opts(fd3, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 2);

	/* Check revision after two detach operations */
	err = bpf_prog_query_opts(cg, atype, &optq);
	ASSERT_OK(err, "prog_query");
	ASSERT_EQ(optq.revision, 7, "revision");

cleanup2:
	err = bpf_prog_detach_opts(fd2, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 1);

cleanup1:
	err = bpf_prog_detach_opts(fd1, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 0);

cleanup:
	cgroup_mprog__destroy(skel);
	close(cg);
}

static void test_link_attach_detach(int atype)
{
	LIBBPF_OPTS(bpf_cgroup_opts, opta);
	LIBBPF_OPTS(bpf_cgroup_opts, optd);
	LIBBPF_OPTS(bpf_prog_query_opts, optq);
	struct bpf_link *link1, *link2, *link3, *link4;
	__u32 fd1, fd2, fd3, fd4, id1, id2, id3, id4;
	struct cgroup_mprog *skel;
	__u32 prog_ids[10];
	int cg, err;

	cg = test__join_cgroup("/link_attach_detach");
	if (!ASSERT_GE(cg, 0, "join_cgroup /link_attach_detach"))
		return;

	skel = cgroup_mprog__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel_load"))
		goto cleanup;

	fd1 = bpf_program__fd(skel->progs.getsockopt_1);
	fd2 = bpf_program__fd(skel->progs.getsockopt_2);
	fd3 = bpf_program__fd(skel->progs.getsockopt_3);
	fd4 = bpf_program__fd(skel->progs.getsockopt_4);

	id1 = id_from_prog_fd(fd1);
	id2 = id_from_prog_fd(fd2);
	id3 = id_from_prog_fd(fd3);
	id4 = id_from_prog_fd(fd4);

	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.expected_revision = 1,
	);

	/* ordering: [fd1] */
	link1 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_1, cg, &opta);
	if (!ASSERT_OK_PTR(link1, "link_attach"))
		goto cleanup;

	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_BEFORE | BPF_F_LINK,
		.relative_id = id_from_link_fd(bpf_link__fd(link1)),
		.expected_revision = 2,
	);

	/* ordering: [fd2, fd1] */
	link2 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_2, cg, &opta);
	if (!ASSERT_OK_PTR(link2, "link_attach"))
		goto cleanup1;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_AFTER | BPF_F_LINK,
		.relative_fd = bpf_link__fd(link2),
		.expected_revision = 3,
	);

	/* ordering: [fd2, fd3, fd1] */
	link3 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_3, cg, &opta);
	if (!ASSERT_OK_PTR(link3, "link_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 3);

	LIBBPF_OPTS_RESET(opta,
		.expected_revision = 4,
	);

	/* ordering: [fd2, fd3, fd1, fd4] */
	link4 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_4, cg, &opta);
	if (!ASSERT_OK_PTR(link4, "link_attach"))
		goto cleanup3;

	assert_mprog_count(cg, atype, 4);

	/* retrieve optq.prog_cnt */
	err = bpf_prog_query_opts(cg, atype, &optq);
	if (!ASSERT_OK(err, "prog_query"))
		goto cleanup4;

	/* optq.prog_cnt will be used in below query */
	memset(prog_ids, 0, sizeof(prog_ids));
	optq.prog_ids = prog_ids;
	err = bpf_prog_query_opts(cg, atype, &optq);
	if (!ASSERT_OK(err, "prog_query"))
		goto cleanup4;

	ASSERT_EQ(optq.count, 4, "count");
	ASSERT_EQ(optq.revision, 5, "revision");
	ASSERT_EQ(optq.prog_ids[0], id2, "prog_ids[0]");
	ASSERT_EQ(optq.prog_ids[1], id3, "prog_ids[1]");
	ASSERT_EQ(optq.prog_ids[2], id1, "prog_ids[2]");
	ASSERT_EQ(optq.prog_ids[3], id4, "prog_ids[3]");
	ASSERT_EQ(optq.prog_ids[4], 0, "prog_ids[4]");
	ASSERT_EQ(optq.link_ids, NULL, "link_ids");

cleanup4:
	bpf_link__destroy(link4);
	assert_mprog_count(cg, atype, 3);

cleanup3:
	bpf_link__destroy(link3);
	assert_mprog_count(cg, atype, 2);

	/* Check revision after two detach operations */
	err = bpf_prog_query_opts(cg, atype, &optq);
	ASSERT_OK(err, "prog_query");
	ASSERT_EQ(optq.revision, 7, "revision");

cleanup2:
	bpf_link__destroy(link2);
	assert_mprog_count(cg, atype, 1);

cleanup1:
	bpf_link__destroy(link1);
	assert_mprog_count(cg, atype, 0);

cleanup:
	cgroup_mprog__destroy(skel);
	close(cg);
}

static void test_preorder_prog_attach_detach(int atype)
{
	LIBBPF_OPTS(bpf_prog_attach_opts, opta);
	LIBBPF_OPTS(bpf_prog_detach_opts, optd);
	__u32 fd1, fd2, fd3, fd4;
	struct cgroup_mprog *skel;
	int cg, err;

	cg = test__join_cgroup("/preorder_prog_attach_detach");
	if (!ASSERT_GE(cg, 0, "join_cgroup /preorder_prog_attach_detach"))
		return;

	skel = cgroup_mprog__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel_load"))
		goto cleanup;

	fd1 = bpf_program__fd(skel->progs.getsockopt_1);
	fd2 = bpf_program__fd(skel->progs.getsockopt_2);
	fd3 = bpf_program__fd(skel->progs.getsockopt_3);
	fd4 = bpf_program__fd(skel->progs.getsockopt_4);

	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI,
		.expected_revision = 1,
	);

	/* ordering: [fd1] */
	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup;

	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_PREORDER,
		.expected_revision = 2,
	);

	/* ordering: [fd1, fd2] */
	err = bpf_prog_attach_opts(fd2, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup1;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER,
		.relative_fd = fd2,
		.expected_revision = 3,
	);

	err = bpf_prog_attach_opts(fd3, cg, atype, &opta);
	if (!ASSERT_EQ(err, -EINVAL, "prog_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER | BPF_F_PREORDER,
		.relative_fd = fd2,
		.expected_revision = 3,
	);

	/* ordering: [fd1, fd2, fd3] */
	err = bpf_prog_attach_opts(fd3, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 3);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI,
		.expected_revision = 4,
	);

	/* ordering: [fd2, fd3, fd1, fd4] */
	err = bpf_prog_attach_opts(fd4, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup3;

	assert_mprog_count(cg, atype, 4);

	err = bpf_prog_detach_opts(fd4, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 3);

cleanup3:
	err = bpf_prog_detach_opts(fd3, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 2);

cleanup2:
	err = bpf_prog_detach_opts(fd2, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 1);

cleanup1:
	err = bpf_prog_detach_opts(fd1, cg, atype, &optd);
	ASSERT_OK(err, "prog_detach");
	assert_mprog_count(cg, atype, 0);

cleanup:
	cgroup_mprog__destroy(skel);
	close(cg);
}

static void test_preorder_link_attach_detach(int atype)
{
	LIBBPF_OPTS(bpf_cgroup_opts, opta);
	struct bpf_link *link1, *link2, *link3, *link4;
	struct cgroup_mprog *skel;
	__u32 fd2;
	int cg;

	cg = test__join_cgroup("/preorder_link_attach_detach");
	if (!ASSERT_GE(cg, 0, "join_cgroup /preorder_link_attach_detach"))
		return;

	skel = cgroup_mprog__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel_load"))
		goto cleanup;

	fd2 = bpf_program__fd(skel->progs.getsockopt_2);

	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.expected_revision = 1,
	);

	/* ordering: [fd1] */
	link1 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_1, cg, &opta);
	if (!ASSERT_OK_PTR(link1, "link_attach"))
		goto cleanup;

	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_PREORDER,
		.expected_revision = 2,
	);

	/* ordering: [fd1, fd2] */
	link2 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_2, cg, &opta);
	if (!ASSERT_OK_PTR(link2, "link_attach"))
		goto cleanup1;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_AFTER,
		.relative_fd = fd2,
		.expected_revision = 3,
	);

	link3 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_3, cg, &opta);
	if (!ASSERT_ERR_PTR(link3, "link_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 2);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_AFTER | BPF_F_PREORDER | BPF_F_LINK,
		.relative_fd = bpf_link__fd(link2),
		.expected_revision = 3,
	);

	/* ordering: [fd1, fd2, fd3] */
	link3 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_3, cg, &opta);
	if (!ASSERT_OK_PTR(link3, "link_attach"))
		goto cleanup2;

	assert_mprog_count(cg, atype, 3);

	LIBBPF_OPTS_RESET(opta,
		.expected_revision = 4,
	);

	/* ordering: [fd2, fd3, fd1, fd4] */
	link4 = bpf_program__attach_cgroup_opts(skel->progs.getsockopt_4, cg, &opta);
	if (!ASSERT_OK_PTR(link4, "prog_attach"))
		goto cleanup3;

	assert_mprog_count(cg, atype, 4);

	bpf_link__destroy(link4);
	assert_mprog_count(cg, atype, 3);

cleanup3:
	bpf_link__destroy(link3);
	assert_mprog_count(cg, atype, 2);

cleanup2:
	bpf_link__destroy(link2);
	assert_mprog_count(cg, atype, 1);

cleanup1:
	bpf_link__destroy(link1);
	assert_mprog_count(cg, atype, 0);

cleanup:
	cgroup_mprog__destroy(skel);
	close(cg);
}

static void test_invalid_attach_detach(int atype)
{
	LIBBPF_OPTS(bpf_prog_attach_opts, opta);
	__u32 fd1, fd2, id2;
	struct cgroup_mprog *skel;
	int cg, err;

	cg = test__join_cgroup("/invalid_attach_detach");
	if (!ASSERT_GE(cg, 0, "join_cgroup /invalid_attach_detach"))
		return;

	skel = cgroup_mprog__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel_load"))
		goto cleanup;

	fd1 = bpf_program__fd(skel->progs.getsockopt_1);
	fd2 = bpf_program__fd(skel->progs.getsockopt_2);

	id2 = id_from_prog_fd(fd2);

	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE | BPF_F_AFTER,
		.relative_id = id2,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE | BPF_F_ID,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -ENOENT, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER | BPF_F_ID,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -ENOENT, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE | BPF_F_AFTER,
		.relative_id = id2,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_LINK,
		.relative_id = id2,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI,
		.relative_id = id2,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_BEFORE,
		.relative_fd = fd1,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -ENOENT, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER,
		.relative_fd = fd1,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -ENOENT, "prog_attach");
	assert_mprog_count(cg, atype, 0);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	if (!ASSERT_EQ(err, 0, "prog_attach"))
		goto cleanup;
	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_AFTER,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 1);

	LIBBPF_OPTS_RESET(opta,
		.flags = BPF_F_ALLOW_MULTI | BPF_F_REPLACE | BPF_F_AFTER,
		.replace_prog_fd = fd1,
	);

	err = bpf_prog_attach_opts(fd1, cg, atype, &opta);
	ASSERT_EQ(err, -EINVAL, "prog_attach");
	assert_mprog_count(cg, atype, 1);
cleanup:
	cgroup_mprog__destroy(skel);
	close(cg);
}

void test_cgroup_mprog_opts(void)
{
	if (test__start_subtest("prog_attach_detach"))
		test_prog_attach_detach(BPF_CGROUP_GETSOCKOPT);
	if (test__start_subtest("link_attach_detach"))
		test_link_attach_detach(BPF_CGROUP_GETSOCKOPT);
	if (test__start_subtest("preorder_prog_attach_detach"))
		test_preorder_prog_attach_detach(BPF_CGROUP_GETSOCKOPT);
	if (test__start_subtest("preorder_link_attach_detach"))
		test_preorder_link_attach_detach(BPF_CGROUP_GETSOCKOPT);
	if (test__start_subtest("invalid_attach_detach"))
		test_invalid_attach_detach(BPF_CGROUP_GETSOCKOPT);
}