// SPDX-License-Identifier: GPL-2.0
#include <linux/bpf.h>
#include <test_progs.h>
#include "cgroup_helpers.h"

static char bpf_log_buf[4096];
static bool verbose;

enum sock_create_test_error {
	OK = 0,
	DENY_CREATE,
};

static struct sock_create_test {
	const char			*descr;
	const struct bpf_insn		insns[64];
	enum bpf_attach_type		attach_type;
	enum bpf_attach_type		expected_attach_type;

	int				domain;
	int				type;
	int				protocol;

	int				optname;
	int				optval;
	enum sock_create_test_error	error;
} tests[] = {
	{
		.descr = "AF_INET set priority",
		.insns = {
			/* r3 = 123 (priority) */
			BPF_MOV64_IMM(BPF_REG_3, 123),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, priority)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_DGRAM,

		.optname = SO_PRIORITY,
		.optval = 123,
	},
	{
		.descr = "AF_INET6 set priority",
		.insns = {
			/* r3 = 123 (priority) */
			BPF_MOV64_IMM(BPF_REG_3, 123),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, priority)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET6,
		.type = SOCK_DGRAM,

		.optname = SO_PRIORITY,
		.optval = 123,
	},
	{
		.descr = "AF_INET set mark",
		.insns = {
			BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),

			/* get uid of process */
			BPF_EMIT_CALL(BPF_FUNC_get_current_uid_gid),
			BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 0xffffffff),

			/* if uid is 0, use given mark(666), else use uid as the mark */
			BPF_MOV64_REG(BPF_REG_3, BPF_REG_0),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),
			BPF_MOV64_IMM(BPF_REG_3, 666),

			BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, mark)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_DGRAM,

		.optname = SO_MARK,
		.optval = 666,
	},
	{
		.descr = "AF_INET6 set mark",
		.insns = {
			BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),

			/* get uid of process */
			BPF_EMIT_CALL(BPF_FUNC_get_current_uid_gid),
			BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 0xffffffff),

			/* if uid is 0, use given mark(666), else use uid as the mark */
			BPF_MOV64_REG(BPF_REG_3, BPF_REG_0),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),
			BPF_MOV64_IMM(BPF_REG_3, 666),

			BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, mark)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET6,
		.type = SOCK_DGRAM,

		.optname = SO_MARK,
		.optval = 666,
	},
	{
		.descr = "AF_INET bound to iface",
		.insns = {
			/* r3 = 1 (lo interface) */
			BPF_MOV64_IMM(BPF_REG_3, 1),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, bound_dev_if)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_DGRAM,

		.optname = SO_BINDTOIFINDEX,
		.optval = 1,
	},
	{
		.descr = "AF_INET6 bound to iface",
		.insns = {
			/* r3 = 1 (lo interface) */
			BPF_MOV64_IMM(BPF_REG_3, 1),
			BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3,
				    offsetof(struct bpf_sock, bound_dev_if)),

			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET6,
		.type = SOCK_DGRAM,

		.optname = SO_BINDTOIFINDEX,
		.optval = 1,
	},
	{
		.descr = "block AF_INET, SOCK_DGRAM, IPPROTO_ICMP socket",
		.insns = {
			BPF_MOV64_IMM(BPF_REG_0, 1),	/* r0 = verdict */

			/* sock->family == AF_INET */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, family)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, AF_INET, 5),

			/* sock->type == SOCK_DGRAM */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, type)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, SOCK_DGRAM, 3),

			/* sock->protocol == IPPROTO_ICMP */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, protocol)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, IPPROTO_ICMP, 1),

			/* return 0 (block) */
			BPF_MOV64_IMM(BPF_REG_0, 0),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_DGRAM,
		.protocol = IPPROTO_ICMP,

		.error = DENY_CREATE,
	},
	{
		.descr = "block AF_INET6, SOCK_DGRAM, IPPROTO_ICMPV6 socket",
		.insns = {
			BPF_MOV64_IMM(BPF_REG_0, 1),	/* r0 = verdict */

			/* sock->family == AF_INET6 */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, family)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, AF_INET6, 5),

			/* sock->type == SOCK_DGRAM */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, type)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, SOCK_DGRAM, 3),

			/* sock->protocol == IPPROTO_ICMPV6 */
			BPF_LDX_MEM(BPF_H, BPF_REG_2, BPF_REG_1,
				    offsetof(struct bpf_sock, protocol)),
			BPF_JMP_IMM(BPF_JNE, BPF_REG_2, IPPROTO_ICMPV6, 1),

			/* return 0 (block) */
			BPF_MOV64_IMM(BPF_REG_0, 0),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = BPF_CGROUP_INET_SOCK_CREATE,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_DGRAM,
		.protocol = IPPROTO_ICMPV6,

		.error = DENY_CREATE,
	},
	{
		.descr = "load w/o expected_attach_type (compat mode)",
		.insns = {
			/* return 1 */
			BPF_MOV64_IMM(BPF_REG_0, 1),
			BPF_EXIT_INSN(),
		},
		.expected_attach_type = 0,
		.attach_type = BPF_CGROUP_INET_SOCK_CREATE,

		.domain = AF_INET,
		.type = SOCK_STREAM,
	},
};

static int load_prog(const struct bpf_insn *insns,
		     enum bpf_attach_type expected_attach_type)
{
	LIBBPF_OPTS(bpf_prog_load_opts, opts,
		    .expected_attach_type = expected_attach_type,
		    .log_level = 2,
		    .log_buf = bpf_log_buf,
		    .log_size = sizeof(bpf_log_buf),
	);
	int fd, insns_cnt = 0;

	for (;
	     insns[insns_cnt].code != (BPF_JMP | BPF_EXIT);
	     insns_cnt++) {
	}
	insns_cnt++;

	fd = bpf_prog_load(BPF_PROG_TYPE_CGROUP_SOCK, NULL, "GPL", insns,
			   insns_cnt, &opts);
	if (verbose && fd < 0)
		fprintf(stderr, "%s\n", bpf_log_buf);

	return fd;
}

static int run_test(int cgroup_fd, struct sock_create_test *test)
{
	int sock_fd, err, prog_fd, optval, ret = -1;
	socklen_t optlen = sizeof(optval);

	prog_fd = load_prog(test->insns, test->expected_attach_type);
	if (prog_fd < 0) {
		log_err("Failed to load BPF program");
		return -1;
	}

	err = bpf_prog_attach(prog_fd, cgroup_fd, test->attach_type, 0);
	if (err < 0) {
		log_err("Failed to attach BPF program");
		goto close_prog_fd;
	}

	sock_fd = socket(test->domain, test->type, test->protocol);
	if (sock_fd < 0) {
		if (test->error == DENY_CREATE)
			ret = 0;
		else
			log_err("Failed to create socket");

		goto detach_prog;
	}

	if (test->optname) {
		err = getsockopt(sock_fd, SOL_SOCKET, test->optname, &optval, &optlen);
		if (err) {
			log_err("Failed to call getsockopt");
			goto cleanup;
		}

		if (optval != test->optval) {
			errno = 0;
			log_err("getsockopt returned unexpected optval");
			goto cleanup;
		}
	}

	ret = test->error != OK;

cleanup:
	close(sock_fd);
detach_prog:
	bpf_prog_detach2(prog_fd, cgroup_fd, test->attach_type);
close_prog_fd:
	close(prog_fd);
	return ret;
}

void test_sock_create(void)
{
	int cgroup_fd, i;

	cgroup_fd = test__join_cgroup("/sock_create");
	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
		return;

	for (i = 0; i < ARRAY_SIZE(tests); i++) {
		if (!test__start_subtest(tests[i].descr))
			continue;

		ASSERT_OK(run_test(cgroup_fd, &tests[i]), tests[i].descr);
	}

	close(cgroup_fd);
}