/** Various socket related functions
 *
 * Parse and print addresses, connect, close, etc...
 *
 * S2SS uses these functions to setup the network emulation feature.
 *
 * @author Steffen Vogel <stvogel@eonerc.rwth-aachen.de>
 * @copyright 2014-2016, Institute for Automation of Complex Power Systems, EONERC
 *   This file is part of S2SS. All Rights Reserved. Proprietary and confidential.
 *   Unauthorized copying of this file, via any medium is strictly prohibited.
 *********************************************************************************/

#include <string.h>
#include <unistd.h>
#include <poll.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/un.h>

#include <netinet/ether.h>
#include <netinet/ip.h>
#include <linux/if_packet.h>
#include <arpa/inet.h>

#include "if.h"
#include "nl.h"
#include "tc.h"
#include "config.h"
#include "utils.h"
#include "socket.h"
#include "checks.h"
#include "pool.h"

/* Forward declartions */
static struct node_type vt;

/* Private static storage */
struct list interfaces;

int socket_init(int argc, char * argv[], config_setting_t *cfg)
{
	if (check_root())
		error("The 'socket' node-type requires superuser privileges!");
	
	nl_init(); /* Fill link cache */
	list_init(&interfaces, (dtor_cb_t) if_destroy);

	/* Gather list of used network interfaces */
	list_foreach(struct node *n, &vt.instances) {
		struct socket *s = n->_vd;
		struct rtnl_link *link;

		/* Determine outgoing interface */
		if (if_get_egress((struct sockaddr *) &s->remote, &link)) {
			char *buf = socket_print_addr((struct sockaddr *) &s->remote);
			error("Failed to get interface for socket address '%s'", buf);
			free(buf);
		}

		/* Search of existing interface with correct ifindex */
		struct interface *i;
		list_foreach(i, &interfaces) {
			if (rtnl_link_get_ifindex(i->nl_link) == rtnl_link_get_ifindex(link))
				goto found;
		}
		
		/* If not found, create a new interface */
		i = if_create(link);
		list_push(&interfaces, i);

found:		list_push(&i->sockets, s);
	}
	
	/** @todo Improve mapping of NIC IRQs per path */
	int affinity;
	if (!config_setting_lookup_int(cfg, "affinity", &affinity))
		affinity = -1;

	list_foreach(struct interface *i, &interfaces)
		if_start(i, affinity);

	return 0;
}

int socket_deinit()
{
	list_foreach(struct interface *i, &interfaces)
		if_stop(i);

	list_destroy(&interfaces);

	return 0;
}

char * socket_print(struct node *n)
{
	struct socket *s = n->_vd;
	char *layer = NULL, *buf = NULL;
	
	switch (s->layer) {
		case LAYER_UDP: layer = "udp";	break;
		case LAYER_IP:	layer = "ip";	break;
		case LAYER_ETH:	layer = "eth";	break;
	}

	char *local = socket_print_addr((struct sockaddr *) &s->local);
	char *remote = socket_print_addr((struct sockaddr *) &s->remote);

	strcatf(&buf, "layer=%s, local=%s, remote=%s", layer, local, remote);
	
	free(local);
	free(remote);
	
	return buf;
}

int socket_open(struct node *n)
{
	struct socket *s = n->_vd;
	struct sockaddr_in *sin = (struct sockaddr_in *) &s->local;
	struct sockaddr_ll *sll = (struct sockaddr_ll *) &s->local;
	int ret;

	/* Create socket */
	switch (s->layer) {
		case LAYER_UDP:	s->sd = socket(sin->sin_family, SOCK_DGRAM,	IPPROTO_UDP); break;
		case LAYER_IP:	s->sd = socket(sin->sin_family, SOCK_RAW,	ntohs(sin->sin_port)); break;
		case LAYER_ETH:	s->sd = socket(sll->sll_family, SOCK_DGRAM,	sll->sll_protocol); break;
		default:
			error("Invalid socket type!");
	}

	if (s->sd < 0)
		serror("Failed to create socket");

	/* Bind socket for receiving */
	ret = bind(s->sd, (struct sockaddr *) &s->local, sizeof(s->local));
	if (ret < 0)
		serror("Failed to bind socket");

	/* Set fwmark for outgoing packets */
	if (setsockopt(s->sd, SOL_SOCKET, SO_MARK, &s->mark, sizeof(s->mark)))
		serror("Failed to set FW mark for outgoing packets");
	else
		debug(4, "Set FW mark for socket (sd=%u) to %u", s->sd, s->mark);

	/* Set socket priority, QoS or TOS IP options */
	int prio;
	switch (s->layer) {
		case LAYER_UDP:
		case LAYER_IP:
			prio = IPTOS_LOWDELAY;
			if (setsockopt(s->sd, IPPROTO_IP, IP_TOS, &prio, sizeof(prio)))
				serror("Failed to set type of service (QoS)");
			else
				debug(4, "Set QoS/TOS IP option for node %s to %#x", node_name(n), prio);
			break;

		default:
			prio = SOCKET_PRIO;
			if (setsockopt(s->sd, SOL_SOCKET, SO_PRIORITY, &prio, sizeof(prio)))
				serror("Failed to set socket priority");
			else
				debug(4, "Set socket priority for node %s to %d", node_name(n), prio);
			break;
	}

	return 0;
}

int socket_reverse(struct node *n)
{
	struct socket *s = n->_vd;
	
	SWAP(s->remote, s->local);
	
	return 0;
}

int socket_close(struct node *n)
{
	struct socket *s = n->_vd;

	if (s->sd >= 0)
		close(s->sd);

	return 0;
}

int socket_destroy(struct node *n)
{
	struct socket *s = n->_vd;
	
	rtnl_qdisc_put(s->tc_qdisc);
	rtnl_cls_put(s->tc_classifier);
	
	return 0;
}

int socket_read(struct node *n, struct pool *pool, int cnt)
{
	struct socket *s = n->_vd;

	int bytes;
	struct iovec iov[cnt];
	struct msghdr mhdr = {
		.msg_iov = iov,
		.msg_iovlen = ARRAY_LEN(iov)
	};

	/* Wait until next packet received */
	poll(&(struct pollfd) { .fd = s->sd, .events = POLLIN }, 1, -1);

	/* Get size of received packet in bytes */
	ioctl(s->sd, FIONREAD, &bytes);

	/* Check if packet length is correct */
	if (bytes % (cnt * 4) != 0)
		error("Packet length not dividable by 4: received=%u, cnt=%u", bytes, cnt);
	if (bytes / cnt > pool->stride)
		error("Packet length is too large: received=%u, cnt=%u, max=%zu", bytes, cnt, pool->stride);

	for (int i = 0; i < cnt; i++) {
		/* All messages of a packet must have equal length! */
		iov[i].iov_base = pool_getrel(pool, i);
		iov[i].iov_len  = bytes / cnt;
	}

	/* Receive message from socket */
	bytes = recvmsg(s->sd, &mhdr, 0);
	if (bytes == 0)
		error("Remote node %s closed the connection", node_name(n));
	else if (bytes < 0)
		serror("Failed recv");

	debug(17, "Received packet of %u bytes: %u samples a %u values per sample", bytes, cnt, (bytes / cnt) / 4 - 4);

	for (int i = 0; i < cnt; i++) {
		struct msg *m = pool_getrel(pool, i);

		/* Convert message to host endianess */
		if (m->endian != MSG_ENDIAN_HOST)
			msg_swap(m);

		/* Check integrity of packet */
		if (bytes / cnt != MSG_LEN(m->values))
			error("Invalid message len: %u for node %s", MSG_LEN(m->values), node_name(n));

		bytes -= MSG_LEN(m->values);
	}

	/* Check packet integrity */
	if (bytes != 0)
		error("Packet length does not match message header length! %u bytes left over.", bytes);

	return cnt;
}

int socket_write(struct node *n, struct pool *pool, int cnt)
{
	struct socket *s = n->_vd;
	int bytes, sent = 0;

	/** @todo we should check the MTU */

	struct iovec iov[cnt];
	for (int i = 0; i < cnt; i++) {
		struct msg *m = pool_getrel(pool, i);

		iov[sent].iov_base = m;
		iov[sent].iov_len  = MSG_LEN(m->values);

		sent++;
	}

	struct msghdr mhdr = {
		.msg_iov = iov,
		.msg_iovlen = sent
	};

	/* Specify destination address for connection-less procotols */
	switch (s->layer) {
		case LAYER_UDP:
		case LAYER_IP:
		case LAYER_ETH:
			mhdr.msg_name = (struct sockaddr *) &s->remote;
			mhdr.msg_namelen = sizeof(s->remote);
			break;
	}

	bytes = sendmsg(s->sd, &mhdr, 0);
	if (bytes < 0)
		serror("Failed send");

	debug(17, "Sent packet of %u bytes: %u samples a %u values per sample", bytes, cnt, (bytes / cnt) / 4 - 4);

	return sent;
}

int socket_parse(struct node *n, config_setting_t *cfg)
{
	const char *local, *remote, *layer;
	int ret;

	struct socket *s = n->_vd;

	if (!config_setting_lookup_string(cfg, "layer", &layer))
		cerror(cfg, "Missing layer for node %s", node_name(n));

	if (!strcmp(layer, "eth"))
		s->layer = LAYER_ETH;
	else if (!strcmp(layer, "ip"))
		s->layer = LAYER_IP;
	else if (!strcmp(layer, "udp"))
		s->layer = LAYER_UDP;
	else
		cerror(cfg, "Invalid layer '%s' for node %s", layer, node_name(n));

	if (!config_setting_lookup_string(cfg, "remote", &remote))
		cerror(cfg, "Missing remote address for node %s", node_name(n));

	if (!config_setting_lookup_string(cfg, "local", &local))
		cerror(cfg, "Missing local address for node %s", node_name(n));
	
	ret = socket_parse_addr(local, (struct sockaddr *) &s->local, s->layer, AI_PASSIVE);
	if (ret) {
		cerror(cfg, "Failed to resolve local address '%s' of node %s: %s",
			local, node_name(n), gai_strerror(ret));
	}

	ret = socket_parse_addr(remote, (struct sockaddr *) &s->remote, s->layer, 0);
	if (ret) {
		cerror(cfg, "Failed to resolve remote address '%s' of node %s: %s",
			remote, node_name(n), gai_strerror(ret));
	}

	config_setting_t *cfg_netem = config_setting_get_member(cfg, "netem");
	if (cfg_netem) {
		int enabled = 1;
		if (!config_setting_lookup_bool(cfg_netem, "enabled", &enabled) || enabled)
			tc_parse(cfg_netem, &s->tc_qdisc);
	}

	return 0;
}

char * socket_print_addr(struct sockaddr *saddr)
{
	union sockaddr_union *sa = (union sockaddr_union *) saddr;
	char *buf = alloc(64);
	
	/* Address */
	switch (sa->sa.sa_family) {
		case AF_INET6:
			inet_ntop(AF_INET6, &sa->sin6.sin6_addr, buf, 64);
			break;

		case AF_INET:
			inet_ntop(AF_INET, &sa->sin.sin_addr, buf, 64);
			break;
			
		case AF_PACKET:
			strcatf(&buf, "%02x", sa->sll.sll_addr[0]);
			for (int i = 1; i < sa->sll.sll_halen; i++)
				strcatf(&buf, ":%02x", sa->sll.sll_addr[i]);
			break;

		default:
			error("Unknown address family: '%u'", sa->sa.sa_family);
	}
	
	/* Port  / Interface */
	switch (sa->sa.sa_family) {
		case AF_INET6:
		case AF_INET:
			strcatf(&buf, ":%hu", ntohs(sa->sin.sin_port));
			break;

		case AF_PACKET: {
			struct nl_cache *cache = nl_cache_mngt_require("route/link");
			struct rtnl_link *link = rtnl_link_get(cache, sa->sll.sll_ifindex);
			if (!link)
				error("Failed to get interface for index: %u", sa->sll.sll_ifindex);
			
			strcatf(&buf, "%%%s", rtnl_link_get_name(link));
			strcatf(&buf, ":%hu", ntohs(sa->sll.sll_protocol));
			break;
		}
	}

	return buf;
}

int socket_parse_addr(const char *addr, struct sockaddr *saddr, enum socket_layer layer, int flags)
{
	/** @todo: Add support for IPv6 */
	union sockaddr_union *sa = (union sockaddr_union *) saddr;

	char *copy = strdup(addr);
	int ret;

	if (layer == LAYER_ETH) { /* Format: "ab:cd:ef:12:34:56%ifname:protocol" */
		/* Split string */
		char *node = strtok(copy, "%");
		char *ifname = strtok(NULL, ":");
		char *proto = strtok(NULL, "\0");

		/* Parse link layer (MAC) address */
		struct ether_addr *mac = ether_aton(node);
		if (!mac)
			error("Failed to parse MAC address: %s", node);

		memcpy(&sa->sll.sll_addr, &mac->ether_addr_octet, 6);
		
		/* Get interface index from name */
		struct nl_cache *cache = nl_cache_mngt_require("route/link");
		struct rtnl_link *link = rtnl_link_get_by_name(cache, ifname);
		if (!link)
			error("Failed to get network interface: '%s'", ifname);

		sa->sll.sll_protocol = htons((proto) ? strtol(proto, NULL, 0) : ETH_P_S2SS);
		sa->sll.sll_halen = 6;
		sa->sll.sll_family = AF_PACKET;
		sa->sll.sll_ifindex = rtnl_link_get_ifindex(link);

		ret = 0;
	}
	else {	/* Format: "192.168.0.10:12001" */
		struct addrinfo hint = {
			.ai_flags = flags,
			.ai_family = AF_UNSPEC
		};

		/* Split string */
		char *node = strtok(copy, ":");
		char *service = strtok(NULL, "\0");

		if (node && !strcmp(node, "*"))
			node = NULL;

		if (service && !strcmp(service, "*"))
			service = NULL;

		switch (layer) {
			case LAYER_IP:
				hint.ai_socktype = SOCK_RAW;
				hint.ai_protocol = (service) ? strtol(service, NULL, 0) : IPPROTO_S2SS;
				hint.ai_flags |= AI_NUMERICSERV;
				break;

			case LAYER_UDP:
				hint.ai_socktype = SOCK_DGRAM;
				hint.ai_protocol = IPPROTO_UDP;
				break;

			default:
				error("Invalid address type");
		}

		/* Lookup address */
		struct addrinfo *result;
		ret = getaddrinfo(node, (layer == LAYER_IP) ? NULL : service, &hint, &result);
		if (!ret) {

			if (layer == LAYER_IP) {
				/* We mis-use the sin_port field to store the IP protocol number on RAW sockets */
				struct sockaddr_in *sin = (struct sockaddr_in *) result->ai_addr;
				sin->sin_port = htons(result->ai_protocol);
			}

			memcpy(sa, result->ai_addr, result->ai_addrlen);

			freeaddrinfo(result);
		}
	}

	free(copy);

	return ret;
}

static struct node_type vt = {
	.name		= "socket",
	.description	= "Network socket (libnl3)",
	.vectorize	= 0, /* unlimited */
	.size		= sizeof(struct socket),
	.destroy	= socket_destroy,
	.reverse	= socket_reverse,
	.parse		= socket_parse,
	.print		= socket_print,
	.open		= socket_open,
	.close		= socket_close,
	.read		= socket_read,
	.write		= socket_write,
	.init		= socket_init,
	.deinit		= socket_deinit
};

REGISTER_NODE_TYPE(&vt)