#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/un.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/errno.h>
#include <pthread.h>
#include <assert.h>

#include <event2/listener.h>
#include <event2/util.h>

#include <MESA/MESA_prof_load.h>

#include <tfe_stream.h>
#include <kni_acceptor.h>
#include <proxy.h>
#include <platform.h>

#ifndef TFE_CONFIG_KNI_UXDOMAIN_PATH_DEFAULT
#define TFE_CONFIG_KNI_UXDOMAIN_PATH_DEFAULT            "/var/run/.tfe_kni_acceptor_handler"
#endif

/* The KNI and TFE communicate with each other by UNIX-based socket,
 * and the protocol between them is based on TLV format(Type-Length-Value).
 * The byte order for each entry in the protocol are Host-Ordered.
 *
 * I. Magic and Total counts of T-L-V tuples, at front of the SOCKET stream.
 * II. After Magic header, the stream consist of several T-L-V tuples.
 *
 * Note. Magic = 0x4d5a
 * Consider of the byte align problem, the minimum length of the value is 4bytes(32-bits).
 *
 *  0               1               2               3
 *  0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |             Magic             |      Total counts of TLV      |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |              Type             |             Length            |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |                             Value                             |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |              Type             |             Length            |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |                             Value                             |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 * |                            .......                            |
 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

enum
{
	KNI_TLV_MAGIC = 0x4d5a
};

enum KNI_TLV_TYPE
{
	KNI_TLV_TYPE_PROTOCOL = 0x0001,
	KNI_TLV_TYPE_KEYRING_ID = 0x0002
};

enum KNI_TLV_VALUE
{
	KNI_TLV_VALUE_HTTP = 0x1,
	KNI_TLV_VALUE_SSL = 0x2,
};

struct kni_tlv_header
{
	uint16_t magic;
	uint16_t counts;
};

struct kni_tlv_info
{
	uint16_t type;
	uint16_t len;

	union
	{
		uint8_t value_as_raw[0];
		uint16_t value_as_uint16;
		uint32_t value_as_uint32;
		uint16_t value_as_uint64;
	};
};

struct kni_acceptor
{
	/* INPUT */
	struct tfe_proxy * proxy;
	const char * profile;
	void * logger;

	/* CONFIG */
	char str_unixdomain_file[TFE_STRING_MAX];

	/* PERSIST RUNTIME RESOURCE */
	int fd_unixdomain;
	struct event_base * ev_base;
	struct evconnlistener * ev_listener;
	pthread_t thread;

	/* PEER CONNECTION RESOUCE
	 * should close by __kni_conn_close() */
	struct event * ev_kni_conn;
	int fd_kni_conn;
	pid_t pid_kni_conn;
};

void __kni_conn_close(struct kni_acceptor * ctx)
{
	if (ctx->fd_kni_conn != 0)
	{
		close(ctx->fd_kni_conn);
		ctx->fd_kni_conn = 0;
	}
	if (ctx->ev_kni_conn != NULL)
	{
		event_free(ctx->ev_kni_conn);
		ctx->ev_kni_conn = NULL;
	}
	if (ctx->pid_kni_conn != 0) { ctx->pid_kni_conn = 0; }
}

static int __kni_parse_tlv_data(struct kni_acceptor * ctx,
	struct tfe_proxy_accept_para * out_para, const char * data, size_t sz_data)
{
	const char * __cursor = data;
	size_t __left_data_len = sz_data;

	/* Parse the STREAM header */
	struct kni_tlv_header * tlv_header = (struct kni_tlv_header *) __cursor;
	if (__left_data_len < sizeof(struct kni_tlv_header))
	{
		TFE_LOG_ERROR(ctx->logger, "Invalid KNI TLV header format, length is not enough.");
		return -1;
	}

	if (tlv_header->magic != KNI_TLV_MAGIC)
	{
		TFE_LOG_ERROR(ctx->logger, "Invalid KNI TLV header magic number %x", tlv_header->magic);
		return -2;
	}

	unsigned int nr_tlv_tuples = tlv_header->counts;
	assert(nr_tlv_tuples != 0 && nr_tlv_tuples < 255);

	__left_data_len -= sizeof(struct kni_tlv_header);
	__cursor += sizeof(struct kni_tlv_header);

	for (unsigned int iter_tlv_tuples = 0; iter_tlv_tuples < nr_tlv_tuples; iter_tlv_tuples++)
	{
		struct kni_tlv_info * tlv_info = (struct kni_tlv_info *) __cursor;
		size_t sz_tlv_info = tlv_info->len + sizeof(struct kni_tlv_info);

		if (__left_data_len < sz_tlv_info)
		{
			TFE_LOG_ERROR(ctx->logger, "TLV info declares invalid value length, Too long.");
			return -3;
		}

		switch (tlv_info->type)
		{
			/* VALUE is uint32_t, length is 4 */
			case KNI_TLV_TYPE_PROTOCOL:
			{
				uint32_t __value = tlv_info->value_as_uint32;
				if (__value == KNI_TLV_VALUE_HTTP)
				{
					out_para->session_type = STREAM_PROTO_PLAIN;
				}
				if (__value == KNI_TLV_VALUE_SSL)
				{
					out_para->session_type = STREAM_PROTO_SSL;
				}
				assert(tlv_info->len == sizeof(uint32_t));
				break;
			}

				/* VALUE is uint32_t, length is 4 */
			case KNI_TLV_TYPE_KEYRING_ID:
			{
				uint32_t __value = tlv_info->value_as_uint32;
				out_para->keyring_id = __value;
				assert(tlv_info->len == sizeof(uint32_t));
				break;
			}

			default: assert(0);
		}

		__cursor += sz_tlv_info;
		__left_data_len -= sz_tlv_info;
	}

	return 0;
}

void __kni_event_cb(evutil_socket_t fd, short what, void * user)
{
	struct kni_acceptor * __ctx = (struct kni_acceptor *) user;
	struct cmsghdr * __cmsghdr;
	struct tfe_proxy_accept_para __accept_para{};

	int * __fds = NULL;
	assert(__ctx != NULL && __ctx->thread == pthread_self());
	assert(what & EV_READ);

	/* We use IOVEC to recieve the fds make by KNI.
	 * This is a kind of magic skill to share socket fds between two(or more) process.
     * http://man7.org/tlpi/code/online/dist/sockets/scm_rights_send.c.html
     */

	constexpr static int __TRANS_FDS_MAX = 2;
	constexpr static int __CONTROLLEN = CMSG_SPACE(__TRANS_FDS_MAX * sizeof(int));

	char __buffer[512] = {0};
	struct iovec __iovec[1];
	struct msghdr __msghdr;

	char __cmptr[__CONTROLLEN];

	__iovec[0].iov_base = __buffer;
	__iovec[0].iov_len = sizeof(__buffer);

	__msghdr.msg_iov = __iovec;
	__msghdr.msg_iovlen = 1;
	__msghdr.msg_name = NULL;
	__msghdr.msg_namelen = 0;
	__msghdr.msg_control = (void *) (__cmptr);
	__msghdr.msg_controllen = (size_t) __CONTROLLEN;

	ssize_t rd = recvmsg(fd, &__msghdr, 0);
	if (rd == -1 && (errno == EINTR || errno == EAGAIN))
	{
		return;
	}
	else if (rd < 0)
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at recving fds from KNI connection: %s. ", strerror(errno));
		goto __close_kni_connection;
	}
	else if (rd == 0)
	{
		TFE_LOG_INFO(__ctx->logger, "KNI disconnect (PID = %u). ", __ctx->pid_kni_conn);
		goto __close_kni_connection;
	}

	__cmsghdr = CMSG_FIRSTHDR(&__msghdr);
	__fds = (int *) (CMSG_DATA(__cmsghdr));

	if (unlikely(__fds == NULL))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at fetch CMSG_DATA() from incoming fds, close KNI connection.");
		goto __close_kni_connection;
	}

	if (unlikely(__kni_parse_tlv_data(__ctx, &__accept_para, __buffer, (size_t) rd) < 0))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at parsing TLV format, close KNI connection.");
		goto __close_kni_connection;
	}

	__accept_para.downstream_fd = __fds[0];
	__accept_para.upstream_fd = __fds[1];

	TFE_PROXY_STAT_INCREASE(STAT_FD_OPEN_BY_KNI_ACCEPT, 2);
	if (tfe_proxy_fds_accept(__ctx->proxy, &__accept_para) < 0)
	{
		goto __drop_recieved_fds;
	}

	return;

__close_kni_connection:
	__kni_conn_close(__ctx);

__drop_recieved_fds:
	TFE_PROXY_STAT_INCREASE(STAT_FD_CLOSE_BY_KNI_ACCEPT_FAIL, 2);
	if (__fds != NULL) evutil_closesocket(__fds[0]);
	if (__fds != NULL) evutil_closesocket(__fds[1]);
}

void __kni_listener_accept_cb(struct evconnlistener * listener, evutil_socket_t fd,
	struct sockaddr * sk_addr, int sk_len, void * user)
{
	struct kni_acceptor * __ctx = (struct kni_acceptor *) user;
	struct event * __event = NULL;

	struct ucred __cr{};
	socklen_t __cr_len = sizeof(struct ucred);

	int ret = 0;

	/* There is only one KNI process can connect to TFE.
	 * If ev_kni_conn is not NULL, there's already a KNI connected to TFE.
	 * We need to refuse this connection
	 */

	if (unlikely(__ctx->ev_kni_conn != NULL))
	{
		TFE_LOG_ERROR(__ctx->logger, "One KNI(PID = %d) has been connected to our program, "
									 "close this connection", __ctx->pid_kni_conn);
		evutil_closesocket(fd);
		return;
	}

	/* Get Peer's PID */
	if (getsockopt(fd, SOL_SOCKET, SO_PEERCRED, (void *) &__cr, &__cr_len) < 0)
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at getsockopt(SO_PEERCRED) for fd %d, close this connection", fd);
		goto __close_this_connection;
	}

	__event = event_new(__ctx->ev_base, fd, EV_READ | EV_PERSIST, __kni_event_cb, __ctx);
	if (unlikely(__event == NULL))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at creating event for fd %d, close this connection.", fd);
		goto __close_this_connection;
	}

	ret = event_add(__event, NULL);
	if (unlikely(ret < 0))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at adding event to evbase, close this connection. ");
		goto __close_this_connection;
	}

	__ctx->fd_kni_conn = fd;
	__ctx->ev_kni_conn = __event;
	__ctx->pid_kni_conn = __cr.pid;

	TFE_LOG_INFO(__ctx->logger, "KNI connected (PID = %u, fd = %d). ", __ctx->pid_kni_conn, __ctx->fd_kni_conn);
	return;

__close_this_connection:
	__kni_conn_close(__ctx);
}

void * __kni_listener_thread_entry(void * args)
{
	struct kni_acceptor * __ctx = (struct kni_acceptor *) args;
	assert(__ctx != NULL && __ctx->thread == pthread_self());

	TFE_LOG_DEBUG(__ctx->logger, "Starting KNI listener thread...");
	event_base_dispatch(__ctx->ev_base);
	TFE_LOG_DEBUG(__ctx->logger, "Stoping KNI listener thread...");
	return (void *) NULL;
}

void kni_acceptor_deinit(struct kni_acceptor * ctx)
{
	if (ctx != NULL && ctx->ev_listener != NULL)
	{
		evconnlistener_free(ctx->ev_listener);
	}

	if (ctx != NULL && ctx->ev_base != NULL)
	{
		event_base_free(ctx->ev_base);
	}

	if (ctx != NULL && ctx->fd_unixdomain != 0)
	{
		close(ctx->fd_unixdomain);
	}

	if (ctx != NULL)
	{
		free(ctx);
	}

	return;
}

struct kni_acceptor * kni_acceptor_init(struct tfe_proxy * proxy, const char * profile, void * logger)
{
	struct kni_acceptor * __ctx = ALLOC(struct kni_acceptor, 1);
	struct sockaddr_un __sockaddr_un;
	int ret = 0;

	__ctx->proxy = proxy;
	__ctx->profile = profile;
	__ctx->logger = logger;

	/* Read the unix domain socket file, this file is used to recieve fds from KNI */
	MESA_load_profile_string_def(profile, "kni", "uxdomain", __ctx->str_unixdomain_file,
		sizeof(__ctx->str_unixdomain_file), TFE_CONFIG_KNI_UXDOMAIN_PATH_DEFAULT);

	if (unlikely(unlink(__ctx->str_unixdomain_file) < 0))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at unlink undomain file %s: %s",
			__ctx->str_unixdomain_file, strerror(errno));
		goto __errout;
	}

	__sockaddr_un.sun_family = AF_UNIX;
	strncpy(__sockaddr_un.sun_path, __ctx->str_unixdomain_file, sizeof(__sockaddr_un.sun_path));

	/* Create new event base, this event base will be dispatched at separated thread */
	__ctx->ev_base = event_base_new();
	if (unlikely(__ctx->ev_base == NULL))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at creating event_base. ");
		goto __errout;
	}

	/* Create a listener */
	__ctx->ev_listener = evconnlistener_new_bind(__ctx->ev_base, __kni_listener_accept_cb, __ctx, 0,
		TFE_CONFIG_BACKLOG_DEFAULT, (struct sockaddr *) (&__sockaddr_un), sizeof(__sockaddr_un));

	if (unlikely(__ctx->ev_listener == NULL))
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at creating evconnlistener.");
		goto __errout;
	}

	/* Create a thread to dispatch ctx->evbase */
	ret = pthread_create(&__ctx->thread, NULL, __kni_listener_thread_entry, (void *) __ctx);
	if (ret < 0)
	{
		TFE_LOG_ERROR(__ctx->logger, "Failed at creating listener thread: %s", strerror(errno));
		goto __errout;
	}

	TFE_LOG_INFO(__ctx->logger, "KNI acceptor unixdomain file: %s", __ctx->str_unixdomain_file);
	return __ctx;

__errout:
	kni_acceptor_deinit(__ctx);
	return NULL;
}