// SPDX-License-Identifier: GPL-2.0
/* Copyright Amazon.com Inc. or its affiliates. */

#define _GNU_SOURCE
#include <sched.h>
#include <stdlib.h>
#include <net/if.h>

#include "test_progs.h"
#include "cgroup_helpers.h"
#include "network_helpers.h"
#include "test_tcp_custom_syncookie.skel.h"

static struct test_tcp_custom_syncookie_case {
	int family, type;
	char addr[16];
	char name[10];
} test_cases[] = {
	{
		.name = "IPv4 TCP",
		.family = AF_INET,
		.type = SOCK_STREAM,
		.addr = "127.0.0.1",
	},
	{
		.name = "IPv6 TCP",
		.family = AF_INET6,
		.type = SOCK_STREAM,
		.addr = "::1",
	},
};

static int setup_netns(void)
{
	if (!ASSERT_OK(unshare(CLONE_NEWNET), "create netns"))
		return -1;

	if (!ASSERT_OK(system("ip link set dev lo up"), "ip"))
		goto err;

	if (!ASSERT_OK(write_sysctl("/proc/sys/net/ipv4/tcp_ecn", "1"),
		       "write_sysctl"))
		goto err;

	return 0;
err:
	return -1;
}

static int setup_tc(struct test_tcp_custom_syncookie *skel)
{
	LIBBPF_OPTS(bpf_tc_hook, qdisc_lo, .attach_point = BPF_TC_INGRESS);
	LIBBPF_OPTS(bpf_tc_opts, tc_attach,
		    .prog_fd = bpf_program__fd(skel->progs.tcp_custom_syncookie));

	qdisc_lo.ifindex = if_nametoindex("lo");
	if (!ASSERT_OK(bpf_tc_hook_create(&qdisc_lo), "qdisc add dev lo clsact"))
		goto err;

	if (!ASSERT_OK(bpf_tc_attach(&qdisc_lo, &tc_attach),
		       "filter add dev lo ingress"))
		goto err;

	return 0;
err:
	return -1;
}

#define msg "Hello World"
#define msglen 11

static void transfer_message(int sender, int receiver)
{
	char buf[msglen];
	int ret;

	ret = send(sender, msg, msglen, 0);
	if (!ASSERT_EQ(ret, msglen, "send"))
		return;

	memset(buf, 0, sizeof(buf));

	ret = recv(receiver, buf, msglen, 0);
	if (!ASSERT_EQ(ret, msglen, "recv"))
		return;

	ret = strncmp(buf, msg, msglen);
	if (!ASSERT_EQ(ret, 0, "strncmp"))
		return;
}

static void create_connection(struct test_tcp_custom_syncookie_case *test_case)
{
	int server, client, child;

	server = start_server(test_case->family, test_case->type, test_case->addr, 0, 0);
	if (!ASSERT_NEQ(server, -1, "start_server"))
		return;

	client = connect_to_fd(server, 0);
	if (!ASSERT_NEQ(client, -1, "connect_to_fd"))
		goto close_server;

	child = accept(server, NULL, 0);
	if (!ASSERT_NEQ(child, -1, "accept"))
		goto close_client;

	transfer_message(client, child);
	transfer_message(child, client);

	close(child);
close_client:
	close(client);
close_server:
	close(server);
}

void test_tcp_custom_syncookie(void)
{
	struct test_tcp_custom_syncookie *skel;
	int i;

	if (setup_netns())
		return;

	skel = test_tcp_custom_syncookie__open_and_load();
	if (!ASSERT_OK_PTR(skel, "open_and_load"))
		return;

	if (setup_tc(skel))
		goto destroy_skel;

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

		skel->bss->handled_syn = false;
		skel->bss->handled_ack = false;

		create_connection(&test_cases[i]);

		ASSERT_EQ(skel->bss->handled_syn, true, "SYN is not handled at tc.");
		ASSERT_EQ(skel->bss->handled_ack, true, "ACK is not handled at tc");
	}

destroy_skel:
	system("tc qdisc del dev lo clsact");

	test_tcp_custom_syncookie__destroy(skel);
}