/*
 * (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/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\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(const char *ruleset)
{
	int rc;

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

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

	rc = nft_run_cmd_from_buffer(g_upf->nft.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;
	}

	LOGP(DNFT, LOGL_DEBUG, "set up nft ruleset: %s\n", osmo_quote_str_c(OTC_SELECT, ruleset, -1));
	return 0;
}

int upf_nft_init()
{
	int rc;

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

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

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

	rc = upf_nft_run(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->nft.table_name, -1));
		return rc;
	}
	LOGP(DNFT, LOGL_NOTICE, "Created nft table %s\n", osmo_quote_str_c(OTC_SELECT, g_upf->nft.table_name, -1));

	rc = upf_nft_run(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->nft.table_name);
		return rc;
	}
	return 0;
}

int upf_nft_free()
{
	if (!g_upf->nft.nft_ctx)
		return 0;
	nft_ctx_free(g_upf->nft.nft_ctx);
	g_upf->nft.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 @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, " @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 char *upf_nft_ruleset_tunmap_create_c(void *ctx, const struct upf_nft_args *args)
{
	OSMO_NAME_C_IMPL(ctx, 1024, "ERROR", upf_nft_ruleset_tunmap_create_buf, args)
}

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;
}

static char *upf_nft_ruleset_tunmap_delete_c(void *ctx, const struct upf_nft_args *args)
{
	OSMO_NAME_C_IMPL(ctx, 512, "ERROR", upf_nft_ruleset_tunmap_delete_buf, args)
}

int upf_nft_tunmap_to_str_buf(char *buf, size_t buflen, const struct upf_nft_tunmap_desc *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.gtp_remote_addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x <---> ", tunmap->access.remote_teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->access.gtp_local_addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x UPF ", tunmap->access.local_teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->core.gtp_local_addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x <---> ", tunmap->core.local_teid);
	OSMO_STRBUF_APPEND(sb, osmo_sockaddr_to_str_buf2, &tunmap->core.gtp_remote_addr);
	OSMO_STRBUF_PRINTF(sb, ":0x%x CORE", tunmap->core.remote_teid);
	return sb.chars_needed;
}

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

static void upf_nft_args_from_tunmap_desc(struct upf_nft_args *args, const struct upf_nft_tunmap_desc *tunmap)
{
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->access.gtp_remote_addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->access.gtp_local_addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->core.gtp_remote_addr.u.sa) == 0);
	OSMO_ASSERT(osmo_sockaddr_port(&tunmap->core.gtp_local_addr.u.sa) == 0);

	*args = (struct upf_nft_args){
		.table_name = g_upf->nft.table_name,
		.peer_a = {
			.addr_remote = &tunmap->access.gtp_remote_addr,
			.teid_remote = tunmap->access.remote_teid,
			.addr_local = &tunmap->access.gtp_local_addr,
			.teid_local = tunmap->access.local_teid,
			.chain_id = tunmap->access.chain_id,
		},
		.peer_b = {
			.addr_remote = &tunmap->core.gtp_remote_addr,
			.teid_remote = tunmap->core.remote_teid,
			.addr_local = &tunmap->core.gtp_local_addr,
			.teid_local = tunmap->core.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->nft.table_name);
}

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

static uint32_t chain_id_next(void)
{
	g_upf->nft.next_chain_id_state++;
	if (!g_upf->nft.next_chain_id_state)
		g_upf->nft.next_chain_id_state++;
	return g_upf->nft.next_chain_id_state;
}

char *upf_nft_tunmap_get_ruleset_str(void *ctx, struct upf_nft_tunmap_desc *tunmap)
{
	struct upf_nft_args args;

	if (!tunmap->access.chain_id)
		tunmap->access.chain_id = chain_id_next();
	if (!tunmap->core.chain_id)
		tunmap->core.chain_id = chain_id_next();

	upf_nft_args_from_tunmap_desc(&args, tunmap);
	return upf_nft_ruleset_tunmap_create_c(ctx, &args);
}

char *upf_nft_tunmap_get_ruleset_del_str(void *ctx, struct upf_nft_tunmap_desc *tunmap)
{
	struct upf_nft_args args;
	upf_nft_args_from_tunmap_desc(&args, tunmap);
	return upf_nft_ruleset_tunmap_delete_c(ctx, &args);
}

int upf_nft_tunmap_create(struct upf_nft_tunmap_desc *tunmap)
{
	return upf_nft_run(upf_nft_tunmap_get_ruleset_str(OTC_SELECT, tunmap));
}

int upf_nft_tunmap_delete(struct upf_nft_tunmap_desc *tunmap)
{
	return upf_nft_run(upf_nft_tunmap_get_ruleset_del_str(OTC_SELECT, tunmap));
}