// SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause

/*
 * BPF-based flow shaping
 *
 * The test brings up two veth in two isolated namespaces, attach some flow
 * shaping program onto it, and ensures that a manual speedtest maximum
 * value matches the rate set in the BPF shapers.
 */

#include <asm-generic/socket.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <math.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <bpf/libbpf.h>
#include <pthread.h>
#include "test_progs.h"
#include "network_helpers.h"
#include "test_tc_edt.skel.h"

#define SERVER_NS "tc-edt-server-ns"
#define CLIENT_NS "tc-edt-client-ns"
#define IP4_ADDR_VETH1 "192.168.1.1"
#define IP4_ADDR_VETH2 "192.168.1.2"
#define IP4_ADDR_VETH2_HEX 0xC0A80102

#define TIMEOUT_MS		2000
#define TEST_PORT		9000
#define TARGET_RATE_MBPS	5.0
#define TX_BYTES_COUNT		(1 * 1000 * 1000)
#define RATE_ERROR_PERCENT	2.0

struct connection {
	int server_listen_fd;
	int server_conn_fd;
	int client_conn_fd;
};

static int setup(struct test_tc_edt *skel)
{
	struct nstoken *nstoken_client, *nstoken_server;
	int ret;

	if (!ASSERT_OK(make_netns(CLIENT_NS), "create client ns"))
		goto fail;
	if (!ASSERT_OK(make_netns(SERVER_NS), "create server ns"))
		goto fail_delete_client_ns;

	nstoken_client = open_netns(CLIENT_NS);
	if (!ASSERT_OK_PTR(nstoken_client, "open client ns"))
		goto fail_delete_server_ns;
	SYS(fail_close_client_ns, "ip link add veth1 type veth peer name %s",
	    "veth2 netns " SERVER_NS);
	SYS(fail_close_client_ns, "ip -4 addr add " IP4_ADDR_VETH1 "/24 dev veth1");
	SYS(fail_close_client_ns, "ip link set veth1 up");

	nstoken_server = open_netns(SERVER_NS);
	if (!ASSERT_OK_PTR(nstoken_server, "enter server ns"))
		goto fail_close_client_ns;
	SYS(fail_close_server_ns, "ip -4 addr add " IP4_ADDR_VETH2 "/24 dev veth2");
	SYS(fail_close_server_ns, "ip link set veth2 up");
	SYS(fail_close_server_ns, "tc qdisc add dev veth2 root fq");
	ret = tc_prog_attach("veth2", -1, bpf_program__fd(skel->progs.tc_prog));
	if (!ASSERT_OK(ret, "attach bpf prog"))
		goto fail_close_server_ns;
	skel->bss->target_rate = TARGET_RATE_MBPS * 1000 * 1000;
	close_netns(nstoken_server);
	close_netns(nstoken_client);

	return 0;

fail_close_server_ns:
	close_netns(nstoken_server);
fail_close_client_ns:
	close_netns(nstoken_client);
fail_delete_server_ns:
	remove_netns(SERVER_NS);
fail_delete_client_ns:
	remove_netns(CLIENT_NS);
fail:
	return -1;
}

static void cleanup(void)
{
	remove_netns(CLIENT_NS);
	remove_netns(SERVER_NS);
}

static void run_test(void)
{
	int server_fd, client_fd, err;
	double rate_mbps, rate_error;
	struct nstoken *nstoken;
	__u64 ts_start, ts_end;

	nstoken = open_netns(SERVER_NS);
	if (!ASSERT_OK_PTR(nstoken, "open server ns"))
		return;
	server_fd = start_server(AF_INET, SOCK_STREAM, IP4_ADDR_VETH2,
			TEST_PORT, TIMEOUT_MS);
	if (!ASSERT_OK_FD(server_fd, "start server"))
		return;

	close_netns(nstoken);
	nstoken = open_netns(CLIENT_NS);
	if (!ASSERT_OK_PTR(nstoken, "open client ns"))
		return;
	client_fd = connect_to_fd(server_fd, 0);
	if (!ASSERT_OK_FD(client_fd, "connect client"))
		return;

	ts_start = get_time_ns();
	err = send_recv_data(server_fd, client_fd, TX_BYTES_COUNT);
	ts_end = get_time_ns();
	close_netns(nstoken);
	ASSERT_OK(err, "send_recv_data");

	rate_mbps = TX_BYTES_COUNT / ((ts_end - ts_start) / 1000.0);
	rate_error =
		fabs((rate_mbps - TARGET_RATE_MBPS) * 100.0 / TARGET_RATE_MBPS);

	ASSERT_LE(rate_error, RATE_ERROR_PERCENT,
		  "rate error is lower than threshold");
}

void test_tc_edt(void)
{
	struct test_tc_edt *skel;

	skel = test_tc_edt__open_and_load();
	if (!ASSERT_OK_PTR(skel, "skel open and load"))
		return;

	if (!ASSERT_OK(setup(skel), "global setup"))
		return;

	run_test();

	cleanup();
	test_tc_edt__destroy(skel);
}