/* OsmoUPF: Verify that skipping used ids works for: UP-SEID, GTP local TEID, nft ruleset chain_id. */

/* (C) 2023 by sysmocom - s.f.m.c. GmbH <info@sysmocom.de>
 *
 * All Rights Reserved
 *
 * Author: Neels Hofmeyr <nhofmeyr@sysmocom.de>
 *
 * SPDX-License-Identifier: GPL-2.0+
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <getopt.h>

#include <nftables/libnftables.h>

#include <osmocom/core/sockaddr_str.h>
#include <osmocom/core/application.h>

#include <osmocom/pfcp/pfcp_endpoint.h>

#include <osmocom/upf/upf.h>
#include <osmocom/upf/netinst.h>
#include <osmocom/upf/up_endpoint.h>
#include <osmocom/upf/up_peer.h>
#include <osmocom/upf/up_session.h>
#include <osmocom/upf/up_gtp_action.h>

#define log(FMT, ARGS...) fprintf(stderr, FMT, ##ARGS)
#define log_assert(COND) do { \
		log("assert(" #COND ")\n"); \
		OSMO_ASSERT(COND); \
	} while (0)

#define log_assert_expect_failure(COND) do { \
		log("assert(" #COND ") <-- EXPECTED TO FAIL (known error)\n"); \
		OSMO_ASSERT(!(COND)); \
	} while (0)


void *main_ctx;
void *ctx;

/* The override of osmo_pfcp_endpoint_tx() stores any Session Establishment Response's UP-SEID here, so that this test
 * can reference specific sessions later.
 */
uint64_t last_up_seid = 0;

void select_poll(void)
{
	while (osmo_select_main_ctx(1));
}

static void setup(const char *name)
{
	log("\n===== START of %s\n", name);
	ctx = talloc_named_const(main_ctx, 0, name);
	g_upf_alloc(ctx);
	osmo_talloc_replace_string(g_upf, &g_upf->pfcp.vty_cfg.local_addr, "1.1.1.1");

	OSMO_ASSERT(netinst_add(g_upf, &g_upf->netinst, "default", "1.1.1.1", NULL));

	/* PFCP endpoint recovery timestamp overridden by time() below */
	upf_pfcp_init();
	/* but do not upf_pfcp_listen() */

	upf_nft_init();

	select_poll();
	log("\n");
}

static void cleanup(void)
{
	up_endpoint_free(&g_upf->pfcp.ep);
	upf_gtp_devs_close();

	upf_gtp_genl_close();

	upf_nft_free();

	log("\n===== END of %s\n", talloc_get_name(ctx));
	talloc_free(ctx);
}

static struct osmo_sockaddr *str2addr(const char *addr, uint16_t port)
{
	static struct osmo_sockaddr osa;
	struct osmo_sockaddr_str str;
	osmo_sockaddr_str_from_str(&str, addr, port);
	osmo_sockaddr_str_to_sockaddr(&str, &osa.u.sas);
	return &osa;
}

static struct up_peer *have_peer(const char *remote_addr, uint16_t port)
{
	return up_peer_find_or_add(g_upf->pfcp.ep, str2addr(remote_addr, port));
}

static struct osmo_pfcp_msg *new_pfcp_msg_for_osmo_upf_rx(struct up_peer *from_peer, enum osmo_pfcp_message_type msg_type)
{
	/* pfcp_endpoint discards received messages immediately after dispatching; in this test, allocate them in
	 * OTC_SELECT so they get discarded on the next select_poll().
	 * osmo_pfcp_msg_alloc_rx() is not useful here, it creates a blank struct to be decoded from raw data; instead,
	 * use osmo_pfcp_msg_alloc_tx_req() which properly sets up the internal structures to match the given msg_type,
	 * and when that is done set m->rx = true to indicate it is a message received by osmo-upf. */
	struct osmo_pfcp_msg *m = osmo_pfcp_msg_alloc_tx_req(OTC_SELECT, &from_peer->remote_addr, msg_type);
	m->rx = true;
	return m;
}

static void peer_assoc(struct up_peer *peer)
{
	struct osmo_pfcp_msg *m = new_pfcp_msg_for_osmo_upf_rx(peer, OSMO_PFCP_MSGT_ASSOC_SETUP_REQ);
	m->ies.assoc_setup_req.recovery_time_stamp = 1234;
	osmo_fsm_inst_dispatch(peer->fi, UP_PEER_EV_RX_ASSOC_SETUP_REQ, m);
	select_poll();
}

static int next_teid = 0x100;
static int next_cp_seid = 0x100;

/* Send a PFCP Session Establishment Request, and return the created session */
static struct up_session *session_est_tunmap(struct up_peer *peer)
{
	struct osmo_pfcp_msg *m;

	struct osmo_pfcp_ie_f_seid cp_f_seid;

	struct osmo_pfcp_ie_f_teid f_teid_access_local;
	struct osmo_pfcp_ie_outer_header_creation ohc_access;

	struct osmo_pfcp_ie_f_teid f_teid_core_local;
	struct osmo_pfcp_ie_outer_header_creation ohc_core;

	struct osmo_pfcp_ie_apply_action aa = {};

	osmo_pfcp_bits_set(aa.bits, OSMO_PFCP_APPLY_ACTION_FORW, true);

	f_teid_access_local = (struct osmo_pfcp_ie_f_teid){
		.choose_flag = true,
		.choose = {
			.ipv4_addr = true,
		},
	};

	ohc_access = (struct osmo_pfcp_ie_outer_header_creation){
		.teid_present = true,
		.teid = next_teid++,
		.ip_addr = {
			.v4_present = true,
			.v4 = *str2addr("5.6.7.8", 0),
		},
	};
	osmo_pfcp_bits_set(ohc_access.desc_bits, OSMO_PFCP_OUTER_HEADER_CREATION_GTP_U_UDP_IPV4, true);

	f_teid_core_local = (struct osmo_pfcp_ie_f_teid){
		.choose_flag = true,
		.choose = {
			.ipv4_addr = true,
		},
	};

	ohc_core = (struct osmo_pfcp_ie_outer_header_creation){
		.teid_present = true,
		.teid = next_teid++,
		.ip_addr = {
			.v4_present = true,
			.v4 = *str2addr("13.14.15.16", 0),
		},
	};
	osmo_pfcp_bits_set(ohc_core.desc_bits, OSMO_PFCP_OUTER_HEADER_CREATION_GTP_U_UDP_IPV4, true);

	cp_f_seid = (struct osmo_pfcp_ie_f_seid){
		.seid = next_cp_seid++,
	};
	osmo_pfcp_ip_addrs_set(&cp_f_seid.ip_addr, osmo_pfcp_endpoint_get_local_addr(g_upf->pfcp.ep->pfcp_ep));

	m = new_pfcp_msg_for_osmo_upf_rx(peer, OSMO_PFCP_MSGT_SESSION_EST_REQ);
	m->h.seid_present = true;
	m->h.seid = 0;
	/* GTP tunmap: remove header from both directions, and add header in both directions */
	m->ies.session_est_req = (struct osmo_pfcp_msg_session_est_req){
		.node_id = m->ies.session_est_req.node_id,
		.cp_f_seid_present = true,
		.cp_f_seid = cp_f_seid,
		.create_pdr_count = 2,
		.create_pdr = {
			{
				.pdr_id = 1,
				.precedence = 255,
				.pdi = {
					.source_iface = OSMO_PFCP_SOURCE_IFACE_CORE,
					.local_f_teid_present = true,
					.local_f_teid = f_teid_core_local,
				},
				.outer_header_removal_present = true,
				.outer_header_removal = {
					.desc = OSMO_PFCP_OUTER_HEADER_REMOVAL_GTP_U_UDP_IPV4,
				},
				.far_id_present = true,
				.far_id = 1,
			},
			{
				.pdr_id = 2,
				.precedence = 255,
				.pdi = {
					.source_iface = OSMO_PFCP_SOURCE_IFACE_ACCESS,
					.local_f_teid_present = true,
					.local_f_teid = f_teid_access_local,
				},
				.outer_header_removal_present = true,
				.outer_header_removal = {
					.desc = OSMO_PFCP_OUTER_HEADER_REMOVAL_GTP_U_UDP_IPV4,
				},
				.far_id_present = true,
				.far_id = 2,
			},
		},
		.create_far_count = 2,
		.create_far = {
			{
				.far_id = 1,
				.forw_params_present = true,
				.forw_params = {
					.destination_iface = OSMO_PFCP_DEST_IFACE_ACCESS,
					.outer_header_creation_present = true,
					.outer_header_creation = ohc_access,
				},
				.apply_action = aa,
			},
			{
				.far_id = 2,
				.forw_params_present = true,
				.forw_params = {
					.destination_iface = OSMO_PFCP_DEST_IFACE_CORE,
					.outer_header_creation_present = true,
					.outer_header_creation = ohc_core,
				},
				.apply_action = aa,
			},
		},
	};

	osmo_fsm_inst_dispatch(peer->fi, UP_PEER_EV_RX_SESSION_EST_REQ, m);
	select_poll();

	return up_session_find_by_up_seid(peer, last_up_seid);
}

static void session_del(struct up_session *session)
{
	struct osmo_pfcp_msg *m;

	log_assert(session);

	m = new_pfcp_msg_for_osmo_upf_rx(session->up_peer, OSMO_PFCP_MSGT_SESSION_DEL_REQ);
	m->h.seid_present = true;
	m->h.seid = session->up_seid;

	osmo_fsm_inst_dispatch(session->fi, UP_SESSION_EV_RX_SESSION_DEL_REQ, m);
	select_poll();
}

static void dump_state(void)
{
	struct up_peer *peer;
	log("\n state:\n");
	llist_for_each_entry(peer, &g_upf->pfcp.ep->peers, entry) {
		struct up_session *session;
		int bkt;
		log(" | peer %s %s\n", peer->fi->name, osmo_fsm_inst_state_name(peer->fi));
		hash_for_each(peer->sessions_by_up_seid, bkt, session, node_by_up_seid) {
			struct up_gtp_action *a;
			llist_for_each_entry(a, &session->active_gtp_actions, entry) {
				if (a->kind != UP_GTP_U_TUNMAP)
					continue;
				log(" |  session[%s]: UP-SEID 0x%"PRIx64"; chain_id access=%u core=%u;"
				    " local TEID access=0x%x core=0x%x\n",
				       osmo_fsm_inst_state_name(session->fi),
				       session->up_seid,
				       a->tunmap.access.chain_id, a->tunmap.core.chain_id,
				       a->tunmap.access.tun.local.teid, a->tunmap.core.tun.local.teid);
			}
		}
	}
	log("\n");
}

static void test_skip_used_id(void)
{
	struct up_peer *peer;
	struct up_session *s1;
	uint64_t s1_up_seid;
	struct up_session *s2;
	struct up_session *s3;
	struct up_session *s4;
	struct up_gtp_action *a;

	setup(__func__);

	log("PFCP Associate peer\n");
	peer = have_peer("1.2.3.4", 1234);
	peer_assoc(peer);
	dump_state();

	/* Make sure to start out all IDs with 1 */
	g_upf->pfcp.ep->next_up_seid_state = 0;
	g_upf->gtp.next_local_teid_state = 0;
	g_upf->tunmap.next_chain_id_state = 0;

	log("set up tunmap, which assigns first UP-SEID 0x1, local-TEID 0x1 and 0x2, chain_ids 1 and 2\n");
	s1 = session_est_tunmap(peer);
	dump_state();

	log_assert(s1->up_seid == 1);
	a = llist_first_entry_or_null(&s1->active_gtp_actions, struct up_gtp_action, entry);
	log_assert(a);
	log_assert(a->kind == UP_GTP_U_TUNMAP);
	log_assert(a->tunmap.core.tun.local.teid == 1);
	log_assert(a->tunmap.access.tun.local.teid == 2);
	log_assert(a->tunmap.access.chain_id == 1);
	log_assert(a->tunmap.core.chain_id == 2);
	log("\n");

	log("simulate wrapping of IDs back to 1\n");
	g_upf->pfcp.ep->next_up_seid_state = 0;
	g_upf->gtp.next_local_teid_state = 0;
	g_upf->tunmap.next_chain_id_state = 0;

	log("set up second tunmap, should use distinct IDs\n");
	s2 = session_est_tunmap(peer);
	dump_state();

	log_assert(s2->up_seid == 2);
	a = llist_first_entry_or_null(&s2->active_gtp_actions, struct up_gtp_action, entry);
	log_assert(a);
	log_assert(a->kind == UP_GTP_U_TUNMAP);
	log_assert(a->tunmap.core.tun.local.teid == 3);
	log_assert(a->tunmap.access.tun.local.teid == 4);
	log_assert(a->tunmap.access.chain_id == 3);
	log_assert(a->tunmap.core.chain_id == 4);
	log("\n");

	log("drop first tunmap (%s)\n", s1->fi->name);
	s1_up_seid = s1->up_seid;
	session_del(s1);
	dump_state();
	log_assert(up_session_find_by_up_seid(peer, s1_up_seid) == NULL);
	log("\n");

	log("again wrap all ID state back to 1\n");
	g_upf->pfcp.ep->next_up_seid_state = 0;
	g_upf->gtp.next_local_teid_state = 0;
	g_upf->tunmap.next_chain_id_state = 0;

	log("set up third tunmap, should now re-use same IDs as the first session\n");
	s3 = session_est_tunmap(peer);
	dump_state();

	log_assert(s3->up_seid == 1);
	a = llist_first_entry_or_null(&s3->active_gtp_actions, struct up_gtp_action, entry);
	log_assert(a);
	log_assert(a->kind == UP_GTP_U_TUNMAP);
	log_assert(a->tunmap.core.tun.local.teid == 1);
	log_assert(a->tunmap.access.tun.local.teid == 2);
	log_assert(a->tunmap.access.chain_id == 1);
	log_assert(a->tunmap.core.chain_id == 2);
	log("\n");

	log("set up 4th tunmap; chain_id state would use 3 and 4, but they are in use, so should assign 5 and 6\n");
	s4 = session_est_tunmap(peer);
	dump_state();

	log_assert(s4->up_seid == 3);
	a = llist_first_entry_or_null(&s4->active_gtp_actions, struct up_gtp_action, entry);
	log_assert(a);
	log_assert(a->kind == UP_GTP_U_TUNMAP);
	log_assert(a->tunmap.core.tun.local.teid == 5);
	log_assert(a->tunmap.access.tun.local.teid == 6);
	log_assert(a->tunmap.access.chain_id == 5);
	log_assert(a->tunmap.core.chain_id == 6);
	log("\n");

	cleanup();
}

static const struct log_info_cat test_default_categories[] = {
	[DREF] = {
		.name = "DREF",
		.description = "Reference Counting",
		.enabled = 1, .loglevel = LOGL_DEBUG,
		.color = OSMO_LOGCOLOR_DARKGREY,
	},
	[DPEER] = {
		.name = "DPEER",
		.description = "PFCP peer association",
		.enabled = 1, .loglevel = LOGL_DEBUG,
		.color = OSMO_LOGCOLOR_YELLOW,
	},
	[DSESSION] = {
		.name = "DSESSION",
		.description = "PFCP sessions",
		.enabled = 1, .loglevel = LOGL_DEBUG,
		.color = OSMO_LOGCOLOR_BLUE,
	},
	[DGTP] = {
		.name = "DGTP",
		.description = "GTP tunneling",
		.enabled = 1, .loglevel = LOGL_DEBUG,
		.color = OSMO_LOGCOLOR_PURPLE,
	},
	[DNFT] = {
		.name = "DNFT",
		.description = "GTP forwarding rules via linux netfilter",
		.enabled = 1, .loglevel = LOGL_DEBUG,
		.color = OSMO_LOGCOLOR_PURPLE,
	},
};

const struct log_info log_info = {
	.cat = test_default_categories,
	.num_cat = ARRAY_SIZE(test_default_categories),
};

static struct {
	bool verbose;
} cmdline_opts = {
	.verbose = false,
};

static void print_help(const char *program)
{
	printf("Usage:\n"
	       "  %s [-v]\n"
	       "Options:\n"
	       "  -h --help      show this text.\n"
	       "  -v --verbose   print source file and line numbers\n",
	       program
	       );
}

static void handle_options(int argc, char **argv)
{
	while (1) {
		int option_index = 0, c;
		static struct option long_options[] = {
			{"help", 0, 0, 'h'},
			{"verbose", 1, 0, 'v'},
			{0, 0, 0, 0}
		};

		c = getopt_long(argc, argv, "hv",
				long_options, &option_index);
		if (c == -1)
			break;

		switch (c) {
		case 'h':
			print_help(argv[0]);
			exit(0);
		case 'v':
			cmdline_opts.verbose = true;
			break;
		default:
			/* catch unknown options *as well as* missing arguments. */
			fprintf(stderr, "Error in command line options. Exiting.\n");
			exit(-1);
			break;
		}
	}
}

int main(int argc, char **argv)
{
	handle_options(argc, argv);

	main_ctx = talloc_named_const(NULL, 0, "main");

	msgb_talloc_ctx_init(main_ctx, 0);

	osmo_fsm_set_dealloc_ctx(OTC_SELECT);

	osmo_init_logging2(main_ctx, &log_info);
	log_set_print_category_hex(osmo_stderr_target, 0);
	log_set_print_category(osmo_stderr_target, 1);
	log_set_print_level(osmo_stderr_target, 1);
	log_set_print_timestamp(osmo_stderr_target, 0);
	log_set_print_extended_timestamp(osmo_stderr_target, 0);
	log_set_all_filter(osmo_stderr_target, 1);

	if (cmdline_opts.verbose) {
		log_set_print_filename2(osmo_stderr_target, LOG_FILENAME_BASENAME);
		log_set_print_filename_pos(osmo_stderr_target, LOG_FILENAME_POS_LINE_END);
		log_set_use_color(osmo_stderr_target, 1);
	} else {
		log_set_print_filename2(osmo_stderr_target, LOG_FILENAME_NONE);
		log_set_use_color(osmo_stderr_target, 0);
	}

	osmo_fsm_log_timeouts(true);
	osmo_fsm_log_addr(false);

	/* actual tests */
	test_skip_used_id();

	log_fini();
	talloc_free(main_ctx);
	return 0;
}

/* overrides */

int osmo_pfcp_endpoint_tx(struct osmo_pfcp_endpoint *ep, struct osmo_pfcp_msg *m)
{
	enum osmo_pfcp_cause *cause;

	log("\n[test override] PFCP tx:\n%s\n\n", osmo_pfcp_msg_to_str_c(OTC_SELECT, m));

	last_up_seid = 0;

	cause = osmo_pfcp_msg_cause(m);
	switch (m->h.message_type) {
	case OSMO_PFCP_MSGT_SESSION_EST_RESP:
		if (*cause == OSMO_PFCP_CAUSE_REQUEST_ACCEPTED) {
			last_up_seid = m->ies.session_est_resp.up_f_seid.seid;
			log("osmo-upf created session 0x%"PRIx64"\n\n", last_up_seid);
		}
		break;
	default:
		break;
	};
	osmo_pfcp_msg_free(m);
	return 0;
}

static void *fake_nft_ctx = (void *)0x1;

struct nft_ctx *nft_ctx_new(uint32_t flags)
{
	log("[test override] %s()\n", __func__);
	return fake_nft_ctx;
}

void nft_ctx_free(struct nft_ctx *ctx)
{
	log("[test override] %s()\n", __func__);
	log_assert(ctx == fake_nft_ctx);
}

int nft_run_cmd_from_buffer(struct nft_ctx *nft, const char *buf)
{
	log("\n[test override] %s():\n%s\n", __func__, buf);
	return 0;
}

/* for deterministic recovery_time_stamp */
time_t time(time_t *tloc)
{
	log("[test override] %s()\n", __func__);
	return 0;
}