/*
 * (C) 2021-2022 by sysmocom - s.f.m.c. GmbH <info@sysmocom.de>
 * All Rights Reserved.
 *
 * Author: Neels Janosch 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 <errno.h>
#include <nftables/libnftables.h>

#include <osmocom/core/talloc.h>
#include <osmocom/core/logging.h>
#include <osmocom/core/timer.h>

#include <osmocom/upf/upf.h>
#include <osmocom/upf/upf_nft.h>

static char *upf_nft_ruleset_table_create(void *ctx, const char *table_name)
{
	return talloc_asprintf(ctx, "add table inet %s { flags owner; };\n", table_name);
}

static char *upf_nft_ruleset_vmap_init(void *ctx, const char *table_name, int priority_pre, int priority_post)
{
	/* add chain inet osmo-upf pre { type filter hook prerouting priority -300; policy accept; }
	 * add chain inet osmo-upf post { type filter hook postrouting priority 400; policy accept; }
	 * add map inet osmo-upf tunmap-pre { typeof ip daddr . @ih,32,32 : verdict; }
	 * add map inet osmo-upf tunmap-post { typeof meta mark : verdict; }
	 * add rule inet osmo-upf pre udp dport 2152 ip daddr . @ih,32,32 vmap @tunmap-pre
	 * add rule inet osmo-upf post meta mark vmap @tunmap-post
	 */
	return talloc_asprintf(ctx,
		"add chain inet %s pre { type filter hook prerouting priority %d; policy accept; };\n"
		"add chain inet %s post { type filter hook postrouting priority %d; policy accept; };\n"
		"add map inet %s tunmap-pre { typeof ip daddr . @ih,32,32 : verdict; };\n"
		"add map inet %s tunmap-post { typeof meta mark : verdict; };\n"
		"add rule inet %s pre udp dport %u ip daddr . @ih,32,32 vmap @tunmap-pre;\n"
		"add rule inet %s post meta mark vmap @tunmap-post;\n",
		table_name, priority_pre,
		table_name, priority_post,
		table_name,
		table_name,
		table_name, PORT_GTP1_U,
		table_name);
}

static int upf_nft_run_now(const char *ruleset)
{
	int rc;
	const int logmax = 256;

	if (g_upf->tunmap.mockup) {
		LOGP(DNFT, LOGL_NOTICE, "tunmap/mockup active: not running nft ruleset: '%s'\n", ruleset);
		return 0;
	}

	if (!g_upf->tunmap.nft_ctx) {
		rc = upf_nft_init();
		if (rc)
			return rc;
	}

	rc = nft_run_cmd_from_buffer(g_upf->tunmap.nft_ctx, ruleset);
	if (rc < 0) {
		LOGP(DNFT, LOGL_ERROR, "error running nft ruleset: rc=%d ruleset=%s\n",
		     rc, osmo_quote_str_c(OTC_SELECT, ruleset, -1));
		return -EIO;
	}

	if (log_check_level(DNFT, LOGL_DEBUG)) {
		size_t l = strlen(ruleset);
		LOGP(DNFT, LOGL_DEBUG, "ran nft ruleset, %zu chars: \"%s%s\"\n",
		     l,
		     osmo_escape_cstr_c(OTC_SELECT, ruleset, OSMO_MIN(logmax, l)),
		     l > logmax ? "..." : "");
	}

	return 0;
}

struct nft_queue {
	struct osmo_tdef *flush_time_tdef;
	struct osmo_tdef *ruleset_max_tdef;
	struct osmo_strbuf sb;
	/* 128 NFT rulesets amount to about 110 kb of char */
	char buf[1<<17];
	unsigned int ruleset_count;
	struct osmo_timer_list timer;
};

static void nft_queue_clear_buf(struct nft_queue *q)
{
	q->sb = (struct osmo_strbuf){ .buf = q->buf, .len = sizeof(q->buf) };
	q->buf[0] = '\0';
}

static void nft_queue_init(void *ctx, struct nft_queue *q,
			   struct osmo_tdef *flush_time_tdef,
			   struct osmo_tdef *ruleset_max_tdef)
{
	*q = (struct nft_queue){
		.flush_time_tdef = flush_time_tdef,
		.ruleset_max_tdef = ruleset_max_tdef,
	};
	nft_queue_clear_buf(q);
}

static void nft_queue_flush(struct nft_queue *q, const char *reason)
{
	static unsigned int flush_count = 0;
	static unsigned int ruleset_count = 0;

	/* We will now flush the queue empty. A timer needs to run only when the next pending entry is added. */
	osmo_timer_del(&q->timer);

	/* Nothing to send? */
	if (!q->sb.chars_needed)
		return;

	flush_count++;
	ruleset_count += q->ruleset_count;
	LOGP(DNFT, LOGL_INFO, "Flushing NFT ruleset queue: %s: n:%u strlen:%zu (flush count: %u avg rules per flush: %s)\n",
	     reason,
	     q->ruleset_count, q->sb.chars_needed,
	     flush_count, osmo_int_to_float_str_c(OTC_SELECT, 10 * ruleset_count / flush_count, 1));

	q->ruleset_count = 0;

	upf_nft_run_now(q->sb.buf);

	nft_queue_clear_buf(q);
}

static void nft_queue_flush_cb(void *q)
{
	nft_queue_flush(q, "timeout");
}

static int nft_enqueue(struct nft_queue *q,
		       int (*tunmap_to_str_buf)(char *buf, size_t len, struct upf_tunmap *tunmap),
		       struct upf_tunmap *tunmap)
{
	int ruleset_max;
	struct osmo_strbuf q_sb_was = q->sb;

	OSMO_STRBUF_APPEND(q->sb, tunmap_to_str_buf, tunmap);

	/* is that being cut off? then revert the addition. This should never happen in practice. */
	if (q->sb.chars_needed >= q->sb.len) {
		q->sb = q_sb_was;
		if (q->sb.pos)
			*q->sb.pos = '\0';
		nft_queue_flush(q, "reached max nr of chars");
		OSMO_STRBUF_APPEND(q->sb, tunmap_to_str_buf, tunmap);
	}

	/* Append separator -- no problem if that gets cut off. */
	OSMO_STRBUF_PRINTF(q->sb, "\n");

	q->ruleset_count++;

	LOGP(DNFT, LOGL_INFO, "Added NFT ruleset to queue: n:%u strlen:%zu\n",
	     q->ruleset_count, q->sb.chars_needed);

	/* Added a rule, see if it has reached ruleset_max. */
	ruleset_max = osmo_tdef_get(q->ruleset_max_tdef, q->ruleset_max_tdef->T, OSMO_TDEF_CUSTOM, 128);
	if (q->ruleset_count >= ruleset_max) {
		nft_queue_flush(q, "reached max nr of rules");
		return 0;
	}

	/* Item added. If the timer is not running yet, schedule a flush in given timeout */
	if (!osmo_timer_pending(&q->timer)) {
		struct osmo_tdef *t;
		unsigned long us;
		osmo_timer_setup(&q->timer, nft_queue_flush_cb, q);
		t = q->flush_time_tdef;
		us = osmo_tdef_get(t, t->T, OSMO_TDEF_US, 100000);
		osmo_timer_schedule(&q->timer, us / 1000000, us % 1000000);
	}
	return 0;
}

static void nft_queue_free(struct nft_queue *q)
{
	osmo_timer_del(&q->timer);
}

static struct nft_queue g_nft_queue = {};

int upf_nft_init()
{
	int rc;

	nft_queue_init(g_upf, &g_nft_queue,
		       osmo_tdef_get_entry(g_upf_nft_tdefs, -32),
		       osmo_tdef_get_entry(g_upf_nft_tdefs, -33));

	/* Always set up the default settings, also in mockup mode, so that the VTY reflects sane values */
	if (!g_upf->tunmap.table_name)
		g_upf->tunmap.table_name = talloc_strdup(g_upf, "osmo-upf");

	/* When in mockup mode, do not set up nft_ctx and netfilter table */
	if (g_upf->tunmap.mockup) {
		LOGP(DNFT, LOGL_NOTICE,
		     "tunmap/mockup active: not allocating libnftables nft_ctx. FOR TESTING PURPOSES ONLY.\n");
		return 0;
	}

	g_upf->tunmap.nft_ctx = nft_ctx_new(NFT_CTX_DEFAULT);
	if (!g_upf->tunmap.nft_ctx) {
		LOGP(DNFT, LOGL_ERROR, "cannot allocate libnftables nft_ctx\n");
		return -EIO;
	}

	rc = upf_nft_run_now(upf_nft_tunmap_get_table_init_str(OTC_SELECT));
	if (rc) {
		LOGP(DNFT, LOGL_ERROR, "Failed to create nft table %s\n",
		     osmo_quote_str_c(OTC_SELECT, g_upf->tunmap.table_name, -1));
		return rc;
	}
	LOGP(DNFT, LOGL_NOTICE, "Created nft table %s\n", osmo_quote_str_c(OTC_SELECT, g_upf->tunmap.table_name, -1));

	rc = upf_nft_run_now(upf_nft_tunmap_get_vmap_init_str(OTC_SELECT));
	if (rc) {
		LOGP(DNFT, LOGL_ERROR, "Failed to initialize nft verdict map in table %s\n", g_upf->tunmap.table_name);
		return rc;
	}
	return 0;
}

int upf_nft_free()
{
	nft_queue_free(&g_nft_queue);
	if (!g_upf->tunmap.nft_ctx)
		return 0;
	nft_ctx_free(g_upf->tunmap.nft_ctx);
	g_upf->tunmap.nft_ctx = NULL;
	return 0;
}

struct upf_nft_args_peer {
	/* The source IP address in packets received from this peer */
	const struct osmo_sockaddr *addr_remote;
	/* The TEID that we send to the peer in GTP packets. */
	uint32_t teid_remote;
	/* The local destination IP address in packets received from this peer */
	const struct osmo_sockaddr *addr_local;
	/* The TEID that the peer sends to us in GTP packets. */
	uint32_t teid_local;
	/* The nft chain id that forwards packets received on addr_local,teid_local. Also used for the 'mark' id in
	 * the verdict map ruleset. */
	uint32_t chain_id;
};

struct upf_nft_args {
	/* global table name */
	const char *table_name;

	struct upf_nft_args_peer peer_a;
	struct upf_nft_args_peer peer_b;
};

static int tunmap_add_single_direction(char *buf, size_t buflen,
				       const struct upf_nft_args *args,
				       bool dir_a2b)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };
	const struct upf_nft_args_peer *from_peer;
	const struct upf_nft_args_peer *to_peer;

	if (dir_a2b) {
		from_peer = &args->peer_a;
		to_peer = &args->peer_b;
	} else {
		from_peer = &args->peer_b;
		to_peer = &args->peer_a;
	}

	/* # add chain for verdict map in prerouting
	 * add chain inet osmo-upf tunmap-pre-123
	 * # mangle destination address at prerouting
	 * add rule inet osmo-upf tunmap-pre-123 ip daddr set 1.1.1.1 meta mark set 123 counter accept
	 *
	 * # add chain for verdict map in postrouting
	 * add chain inet osmo-upf tunmap-post-123
	 * # mangle source address and GTP TID at postrouting
	 * add rule inet osmo-upf tunmap-post-123 ip saddr set 2.2.2.1 udp sport set 2152 @ih,32,32 set 0x00000102 counter accept
	 *
	 * # add elements to verdict map, jump to chain
	 * add element inet osmo-upf tunmap-pre { 2.2.2.3 . 0x00000203 : jump tunmap-pre-123 }
	 * add element inet osmo-upf tunmap-post { 123 : jump tunmap-post-123 }
	 */

	OSMO_STRBUF_PRINTF(sb, "add chain inet %s tunmap-pre-%u;\n",
			   args->table_name, from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "add rule inet %s tunmap-pre-%u",
			   args->table_name, from_peer->chain_id);
	OSMO_STRBUF_PRINTF(sb, " ip daddr set ");
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, to_peer->addr_remote);
	OSMO_STRBUF_PRINTF(sb, " meta mark set %u counter accept;\n", from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "add chain inet %s tunmap-post-%u;\n",
			   args->table_name, from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "add rule inet %s tunmap-post-%u",
			   args->table_name, from_peer->chain_id);
	OSMO_STRBUF_PRINTF(sb, " ip saddr set ");
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, to_peer->addr_local);
	OSMO_STRBUF_PRINTF(sb, " udp sport set 2152");
	OSMO_STRBUF_PRINTF(sb, " @ih,32,32 set 0x%x", to_peer->teid_remote);
	OSMO_STRBUF_PRINTF(sb, " counter accept;\n");

	OSMO_STRBUF_PRINTF(sb, "add element inet %s tunmap-pre { ",
			   args->table_name);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, from_peer->addr_local);
	OSMO_STRBUF_PRINTF(sb, " . 0x%x : jump tunmap-pre-%u };\n",
			   from_peer->teid_local, from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "add element inet %s tunmap-post { %u : jump tunmap-post-%u };\n",
			   args->table_name, from_peer->chain_id, from_peer->chain_id);

	return sb.chars_needed;
}

static int tunmap_del_single_direction(char *buf, size_t buflen,
				       const struct upf_nft_args *args,
				       bool dir_a2b)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };
	const struct upf_nft_args_peer *from_peer;

	if (dir_a2b)
		from_peer = &args->peer_a;
	else
		from_peer = &args->peer_b;

	/* delete element inet osmo-upf tunmap-pre { 2.2.2.3 . 0x203 }
	 * delete element inet osmo-upf tunmap-post { 123 }
	 * delete chain inet osmo-upf tunmap-pre-123
	 * delete chain inet osmo-upf tunmap-post-123
	 */

	OSMO_STRBUF_PRINTF(sb, "delete element inet %s tunmap-pre { ",
			   args->table_name);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, from_peer->addr_local);
	OSMO_STRBUF_PRINTF(sb, " . 0x%x };\n", from_peer->teid_local);

	OSMO_STRBUF_PRINTF(sb, "delete element inet %s tunmap-post { %u };\n",
			   args->table_name, from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "delete chain inet %s tunmap-pre-%u;\n",
			   args->table_name, from_peer->chain_id);

	OSMO_STRBUF_PRINTF(sb, "delete chain inet %s tunmap-post-%u;\n",
			   args->table_name, from_peer->chain_id);

	return sb.chars_needed;
}

static int upf_nft_ruleset_tunmap_create_buf(char *buf, size_t buflen, const struct upf_nft_args *args)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };

	/* Forwarding from peer_a to peer_b */
	OSMO_STRBUF_APPEND(sb, tunmap_add_single_direction, args, true);
	/* And from peer_b to peer_a */
	OSMO_STRBUF_APPEND(sb, tunmap_add_single_direction, args, false);

	return sb.chars_needed;
}

static int upf_nft_ruleset_tunmap_delete_buf(char *buf, size_t buflen, const struct upf_nft_args *args)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };

	/* Forwarding from peer_a to peer_b */
	OSMO_STRBUF_APPEND(sb, tunmap_del_single_direction, args, true);
	/* And from peer_b to peer_a */
	OSMO_STRBUF_APPEND(sb, tunmap_del_single_direction, args, false);

	return sb.chars_needed;
}

int upf_nft_tunmap_to_str_buf(char *buf, size_t buflen, const struct upf_tunmap *tunmap)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };

	/* ACCESS 1.1.1.2:0x102 <---> 2.2.2.1:0x201 UPF 2.2.2.3:0x203 <---> 3.3.3.2:0x302 CORE */
	OSMO_STRBUF_PRINTF(sb, "ACCESS ");
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->access.tun.remote.addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x <---> ", tunmap->access.tun.remote.teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->access.tun.local.addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x UPF ", tunmap->access.tun.local.teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->core.tun.local.addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x <---> ", tunmap->core.tun.local.teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->core.tun.remote.addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x CORE", tunmap->core.tun.remote.teid);
	return sb.chars_needed;
}

char *upf_nft_tunmap_to_str_c(void *ctx, const struct upf_tunmap *tunmap)
{
	OSMO_NAME_C_IMPL(ctx, 128, "ERROR", upf_nft_tunmap_to_str_buf, tunmap)
}

static void upf_nft_args_from_tunmap(struct upf_nft_args *args, const struct upf_tunmap *tunmap)
{
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->access.tun.remote.addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->access.tun.local.addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->core.tun.remote.addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->core.tun.local.addr.u.sa) == 0);

	*args = (struct upf_nft_args){
		.table_name = g_upf->tunmap.table_name,
		.peer_a = {
			.addr_remote = &tunmap->access.tun.remote.addr,
			.teid_remote = tunmap->access.tun.remote.teid,
			.addr_local = &tunmap->access.tun.local.addr,
			.teid_local = tunmap->access.tun.local.teid,
			.chain_id = tunmap->access.chain_id,
		},
		.peer_b = {
			.addr_remote = &tunmap->core.tun.remote.addr,
			.teid_remote = tunmap->core.tun.remote.teid,
			.addr_local = &tunmap->core.tun.local.addr,
			.teid_local = tunmap->core.tun.local.teid,
			.chain_id = tunmap->core.chain_id,
		},
	};
}

char *upf_nft_tunmap_get_table_init_str(void *ctx)
{
	return upf_nft_ruleset_table_create(ctx, g_upf->tunmap.table_name);
}

char *upf_nft_tunmap_get_vmap_init_str(void *ctx)
{
	return upf_nft_ruleset_vmap_init(ctx, g_upf->tunmap.table_name, g_upf->tunmap.priority_pre,
					 g_upf->tunmap.priority_post);
}

int upf_nft_tunmap_get_ruleset_str_buf(char *buf, size_t len, struct upf_tunmap *tunmap)
{
	struct upf_nft_args args;
	upf_nft_args_from_tunmap(&args, tunmap);
	return upf_nft_ruleset_tunmap_create_buf(buf, len, &args);
}

char *upf_nft_tunmap_get_ruleset_str(void *ctx, struct upf_tunmap *tunmap)
{
	OSMO_NAME_C_IMPL(ctx, 1024, "ERROR", upf_nft_tunmap_get_ruleset_str_buf, tunmap)
}

int upf_nft_tunmap_get_ruleset_del_str_buf(char *buf, size_t len, struct upf_tunmap *tunmap)
{
	struct upf_nft_args args;
	upf_nft_args_from_tunmap(&args, tunmap);
	return upf_nft_ruleset_tunmap_delete_buf(buf, len, &args);
}

char *upf_nft_tunmap_get_ruleset_del_str(void *ctx, struct upf_tunmap *tunmap)
{
	OSMO_NAME_C_IMPL(ctx, 1024, "ERROR", upf_nft_tunmap_get_ruleset_del_str_buf, tunmap)
}

static int upf_nft_tunmap_ensure_chain_id(struct upf_nft_tun *tun)
{
	if (tun->chain_id)
		return 0;
	tun->chain_id = upf_next_chain_id();
	if (!tun->chain_id)
		return -ENOSPC;
	return 0;
}

int upf_nft_tunmap_create(struct upf_tunmap *tunmap)
{
	if (upf_nft_tunmap_ensure_chain_id(&tunmap->access)
	    || upf_nft_tunmap_ensure_chain_id(&tunmap->core))
		return -ENOSPC;
	return nft_enqueue(&g_nft_queue, upf_nft_tunmap_get_ruleset_str_buf, tunmap);
}

int upf_nft_tunmap_delete(struct upf_tunmap *tunmap)
{
	return nft_enqueue(&g_nft_queue, upf_nft_tunmap_get_ruleset_del_str_buf, tunmap);
}