/*
 * nsock.c -- portable BSD sockets wrapper
 *
 * Copyright (c) 2013 David Demelier <markand@malikania.fr>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#  include <stdio.h>
#  include <stdlib.h>

#if defined(_WIN32)

  typedef SOCKET		nsock_socket_t;
  typedef const char *		nsock_carg_t;
  typedef char *		nsock_arg_t;

#else
#  include <errno.h>
#  include <string.h>

#  define ioctlsocket(s)	ioctl(s)
#  define closesocket(s)	close(s)

#  define gai_strerrorA		gai_strerror

#  define INVALID_SOCKET	-1
#  define SOCKET_ERROR		-1

  typedef int			nsock_socket_t;
  typedef const void *		nsock_carg_t;
  typedef void *		nsock_arg_t;

#endif

#include "nsock.h"

#define	TAILQ_HEAD(name, type)						\
struct name {								\
	struct type *tqh_first;	/* first element */			\
	struct type **tqh_last;	/* addr of last next element */		\
	size_t noclients;						\
}

#define	TAILQ_ENTRY(type)						\
struct {								\
	struct type *tqe_next;	/* next element */			\
	struct type **tqe_prev;	/* address of previous next element */	\
}

#define	TAILQ_EMPTY(head)	((head)->tqh_first == NULL)

#define	TAILQ_FIRST(head)	((head)->tqh_first)

#define	TAILQ_FOREACH(var, head, field)					\
	for ((var) = TAILQ_FIRST((head));				\
	    (var);							\
	    (var) = TAILQ_NEXT((var), field))

#define	TAILQ_FOREACH_SAFE(var, head, field, tvar)			\
	for ((var) = TAILQ_FIRST((head));				\
	    (var) && ((tvar) = TAILQ_NEXT((var), field), 1);		\
	    (var) = (tvar))

#define	TAILQ_INIT(head) do {						\
	TAILQ_FIRST((head)) = NULL;					\
	(head)->tqh_last = &TAILQ_FIRST((head));			\
} while (0)

#define	TAILQ_INSERT_HEAD(head, elm, field) do {			\
	if ((TAILQ_NEXT((elm), field) = TAILQ_FIRST((head))) != NULL)	\
		TAILQ_FIRST((head))->field.tqe_prev =			\
		    &TAILQ_NEXT((elm), field);				\
	else								\
		(head)->tqh_last = &TAILQ_NEXT((elm), field);		\
	TAILQ_FIRST((head)) = (elm);					\
	(elm)->field.tqe_prev = &TAILQ_FIRST((head));			\
	(head)->noclients ++;						\
} while (0)

#define	TAILQ_INSERT_TAIL(head, elm, field) do {			\
	TAILQ_NEXT((elm), field) = NULL;				\
	(elm)->field.tqe_prev = (head)->tqh_last;			\
	*(head)->tqh_last = (elm);					\
	(head)->tqh_last = &TAILQ_NEXT((elm), field);			\
	(head)->noclients ++;						\
} while (0)

#define	TAILQ_NEXT(elm, field) ((elm)->field.tqe_next)

#define	TAILQ_REMOVE(head, elm, field) do {				\
	if ((TAILQ_NEXT((elm), field)) != NULL)				\
		TAILQ_NEXT((elm), field)->field.tqe_prev = 		\
		    (elm)->field.tqe_prev;				\
	else {								\
		(head)->tqh_last = (elm)->field.tqe_prev;		\
	}								\
	*(elm)->field.tqe_prev = TAILQ_NEXT((elm), field);		\
	(head)->noclients --;						\
} while (0)

struct nsock {
	nsock_socket_t		fd;
	char			error[128];
};

struct nsock_address {
	struct sockaddr_storage	addr;
	socklen_t		addrlen;
	char			error[128];
};

struct nsock_listener {
	const struct nsock	*sock;
	char			error[128];
	TAILQ_HEAD(, nsock_clt)	head;
};

struct nsock_clt {
	const struct nsock	*sock;
	TAILQ_ENTRY(nsock_clt)	link;
};

/* --------------------------------------------------------
 * Private helpers
 * -------------------------------------------------------- */

static void
nsock_set_error(char *buffer, size_t bufsize)
{
	memset(buffer, 0, bufsize);

#if defined(_WIN32)
	LPSTR str;

	FormatMessageA(
		FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
		NULL,
		WSAGetLastError(),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPSTR)&str, 0, NULL);

	if (str) {
		strncpy(buffer, str, bufsize);
		LocalFree(str);
	}
#else
	strncpy(buffer, strerror(errno), bufsize);
#endif
}

static void
nsock_set_errno(char *buffer, size_t bufsize, int no)
{
	memset(buffer, 0, bufsize);
	strncpy(buffer, strerror(no), bufsize - 1);
}

static struct nsock_clt *
nsock_clt_new(const struct nsock *ns)
{
	struct nsock_clt *clt;

	if ((clt = malloc(sizeof (struct nsock_clt))) == NULL)
		return NULL;

	clt->sock = ns;

	return clt;
}

/* --------------------------------------------------------
 * Sockets functions
 * -------------------------------------------------------- */

void
nsock_init(void)
{
#if defined(_WIN32)
	WSADATA wsa;
	WSAStartup(MAKEWORD(2, 2), &wsa);
#endif
}

struct nsock *
nsock_new(void)
{
	struct nsock *ns;

	if ((ns = malloc(sizeof (struct nsock))) == NULL)
		return NULL;

	return ns;
}

int
nsock_create(struct nsock *ns, int domain, int type, int protocol)
{
	ns->fd = socket(domain, type, protocol);

	if (ns->fd == INVALID_SOCKET)
		nsock_set_error(ns->error, sizeof (ns->error));

	return ns->fd == INVALID_SOCKET ? -1 : 0;
}

const char *
nsock_error(struct nsock *ns)
{
	if (ns == NULL)
		return strerror(ENOMEM);

	return ns->error;
}

int
nsock_bind(struct nsock *ns, const struct nsock_address *ep)
{
	int ret;

	ret = bind(ns->fd, (const struct sockaddr *)&ep->addr, ep->addrlen);

	if (ret == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return ret == SOCKET_ERROR ? -1 : 0;
}

int
nsock_listen(struct nsock *ns, int max)
{
	int ret;

	ret = listen(ns->fd, max);

	if (ret == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return ret == SOCKET_ERROR ? -1 : 0;
}

int
nsock_accept(struct nsock *ns, struct nsock **client, struct nsock_address **clientinfo)
{
	struct sockaddr_storage *st = NULL;
	socklen_t *len = NULL;
	int ret;

	if ((*client = malloc(sizeof (struct nsock))) == NULL) {
		nsock_set_errno(ns->error, sizeof (ns->error), ENOMEM);
		return -1;
	}

	if (clientinfo != NULL) {
		if ((*clientinfo = malloc(sizeof (struct nsock_address))) == NULL) {
			free(client);
			nsock_set_errno(ns->error, sizeof (ns->error), ENOMEM);
			return -1;
		}

		st = &(*clientinfo)->addr;
		len = &(*clientinfo)->addrlen;

		/* Set the addrlen to sockaddr_storage first */
		*len = sizeof (struct sockaddr_storage);
	}

	/* Prepare client */
	memset((*client)->error, 0, sizeof ((*client)->error));
	(*client)->fd = accept(ns->fd, (struct sockaddr *)st, len);

	if ((*client)->fd == INVALID_SOCKET) {
		nsock_set_error(ns->error, sizeof (ns->error));

		/* free clients and set to NULL so user will not use it */
		free(*client);
		*client = NULL;

		if (clientinfo != NULL) {
			free(*clientinfo);
			*clientinfo = NULL;
		}

		ret = -1;
	} else
		ret = 0;

	return ret;
}

int
nsock_connect(struct nsock *ns, const struct nsock_address *ep)
{
	int ret;

	ret = connect(ns->fd, (const struct sockaddr *)&ep->addr, ep->addrlen);

	if (ret == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return ret == SOCKET_ERROR ? -1 : 0;
}

int
nsock_set(struct nsock *ns, int level, int name, const void *arg, unsigned arglen)
{
	int ret;

	ret = setsockopt(ns->fd, level, name, (nsock_carg_t)arg, arglen);

	if (ret == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return ret == SOCKET_ERROR ? -1 : 0;
}

long
nsock_recv(struct nsock *ns, void *data, size_t datasz, int flags)
{
	long nbread;

	nbread = recv(ns->fd, data, datasz, flags);

	if (nbread == -1)
		nsock_set_error(ns->error, sizeof (ns->error));

	return nbread;
}

long
nsock_recvfrom(struct nsock *ns, struct nsock_address *ep, void *data, size_t datasz, int flags)
{
	struct sockaddr_storage *st = NULL;
	socklen_t *len = NULL;
	long nbread;

	if (ep != NULL) {
		st = &ep->addr;
		len = &ep->addrlen;

		/* Set the addrlen to sockaddr_storage first */
		*len = sizeof (struct sockaddr_storage);
	}

	nbread = recvfrom(ns->fd, data, datasz, flags,
	    (struct sockaddr *)st, len);

	if (nbread == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return nbread;
}

long
nsock_send(struct nsock *ns, const void *data, size_t datasz, int flags)
{
	long nbsent;

	nbsent = send(ns->fd, data, datasz, flags);

	if (nbsent == -1)
		nsock_set_error(ns->error, sizeof (ns->error));

	return nbsent;
}

long
nsock_sendto(struct nsock *ns, const struct nsock_address *ep, const void *data, size_t datasz, int flags)
{
	long nbsent;

	nbsent = sendto(ns->fd, data, datasz, flags,
	    (const struct sockaddr *)&ep->addr, ep->addrlen);

	if (nbsent == SOCKET_ERROR)
		nsock_set_error(ns->error, sizeof (ns->error));

	return nbsent;
}

void
nsock_close(struct nsock *ns)
{
	closesocket(ns->fd);
}

void
nsock_free(struct nsock *ns)
{
	free(ns);
}

void
nsock_finish(void)
{
#if defined(_WIN32)
	WSACleanup();
#endif
}

/* --------------------------------------------------------
 * End point functions
 * -------------------------------------------------------- */

struct nsock_address *
nsock_addr_new(void)
{
	struct nsock_address *ep;

	if ((ep = calloc(1, sizeof (struct nsock_address))) == NULL)
		return NULL;

	return ep;
}

const char *
nsock_addr_error(struct nsock_address *ep)
{
	if (ep == NULL)
		return strerror(ENOMEM);

	return ep->error;
}

int
nsock_addr_bind_ip(struct nsock_address *ep, const char *iface, unsigned port, int family)
{
	if (family == AF_INET6) {
		struct sockaddr_in6 *ptr = (struct sockaddr_in6 *)&ep->addr;

		memset(ptr, 0, sizeof (struct sockaddr_in6));
		ptr->sin6_family = AF_INET6;
		ptr->sin6_port = htons(port);

		if (iface == NULL || strcmp(iface, "*") == 0)
			ptr->sin6_addr = in6addr_any;
		else if (inet_pton(AF_INET6, iface, &ptr->sin6_addr) <= 0) {
			nsock_set_error(ep->error, sizeof (ep->error));
			return -1;
		}

		ep->addrlen = sizeof (struct sockaddr_in6);
	} else {
		struct sockaddr_in *ptr = (struct sockaddr_in *)&ep->addr;

		memset(ptr, 0, sizeof (struct sockaddr_in));
		ptr->sin_family = AF_INET;
		ptr->sin_port = htons(port);

		if (iface == NULL || strcmp(iface, "*") == 0)
			ptr->sin_addr.s_addr = INADDR_ANY;
		else if (inet_pton(AF_INET, iface, &ptr->sin_addr) <= 0) {
			nsock_set_error(ep->error, sizeof (ep->error));
			return -1;
		}

		ep->addrlen = sizeof (struct sockaddr_in);
	}
	
	return 0;
}

int
nsock_addr_connect_ip(struct nsock_address *ep, const char *host, unsigned port, int family)
{
	struct addrinfo hints, *res;
	char portstr[32];
	int error;

	memset(&hints, 0, sizeof (hints));
	hints.ai_family = family;
	hints.ai_socktype = SOCK_STREAM;

	memset(portstr, 0, sizeof (portstr));
	sprintf(portstr, "%u", port);

	error = getaddrinfo(host, portstr, &hints, &res);
	if (error) {
		memset(ep->error, 0, sizeof (ep->error));
		strncpy(ep->error, gai_strerrorA(error), sizeof (ep->error) - 1);
		return -1;
	}

	memcpy(&ep->addr, res->ai_addr, res->ai_addrlen);
	ep->addrlen = res->ai_addrlen;

	freeaddrinfo(res);

	return 0;
}

#if !defined(_WIN32)

void
nsock_addr_unix(struct nsock_address *ep, const char *path)
{
	struct sockaddr_un *ptr= (struct sockaddr_un *)&ep->addr;

	/* Path */
	memset(ptr, 0, sizeof (struct sockaddr_un));
	strncpy(ptr->sun_path, path, sizeof (ptr->sun_path) - 1);
	ptr->sun_family = AF_UNIX;

	/* Len is computed with SUN_LEN */
	ep->addrlen = SUN_LEN(ptr);
}

#endif

struct sockaddr *
nsock_addr_getaddr(struct nsock_address *ep)
{
	return (struct sockaddr *)&ep->addr;
}

socklen_t
nsock_addr_getaddrlen(const struct nsock_address *ep)
{
	return ep->addrlen;
}

void
nsock_addr_free(struct nsock_address *ep)
{
	free(ep);
}

/* --------------------------------------------------------
 * listener functions
 * -------------------------------------------------------- */

struct nsock_listener *
nsock_lst_new(const struct nsock *ns)
{
	struct nsock_listener *ls;	

	if ((ls = malloc(sizeof (struct nsock_listener))) == NULL)
		return NULL;

	ls->sock = ns;	
	TAILQ_INIT(&ls->head);

	return ls;
}

const char *
nsock_lst_error(struct nsock_listener *ls)
{
	if (ls == NULL)
		return strerror(ENOMEM);

	return ls->error;
}

int
nsock_lst_push(struct nsock_listener *ls, struct nsock *ns)
{
	struct nsock_clt *clt;

	if ((clt = nsock_clt_new(ns)) == NULL)
		return -1;

	TAILQ_INSERT_HEAD(&ls->head, clt, link);

	return 0;
}

int
nsock_lst_append(struct nsock_listener *ls, struct nsock *ns)
{
	struct nsock_clt *clt;

	if ((clt = nsock_clt_new(ns)) == NULL)
		return -1;

	TAILQ_INSERT_TAIL(&ls->head, clt, link);

	return 0;
}

size_t
nsock_lst_count(const struct nsock_listener *ls)
{
	return ls->head.noclients;
}

void
nsock_lst_remove(struct nsock_listener *ls, const struct nsock *ns)
{
	struct nsock_clt *clt, *tmp;

	TAILQ_FOREACH_SAFE(clt, &ls->head, link, tmp) {
		if (clt->sock == ns) {
			TAILQ_REMOVE(&ls->head, clt, link);
			free(clt);
			break;
		}
	}
}

struct nsock *
nsock_lst_select(struct nsock_listener *ls, long sec, long usec)
{
	fd_set fds;
	struct timeval maxwait, *towait;
	int error;
	int fdmax;
	struct nsock_clt *clt;

	fdmax = TAILQ_FIRST(&ls->head)->sock->fd;

	FD_ZERO(&fds);
	TAILQ_FOREACH(clt, &ls->head, link) {
		FD_SET(clt->sock->fd, &fds);
		if ((int)clt->sock->fd > fdmax)
			fdmax = clt->sock->fd;
	}

	maxwait.tv_sec = sec;
	maxwait.tv_usec = usec;

        // Set to NULL for infinite timeout.
	towait = (sec == 0 && usec == 0) ? NULL : &maxwait;
	error = select(fdmax + 1, &fds, NULL, NULL, towait);

	TAILQ_FOREACH(clt, &ls->head, link)
		if (FD_ISSET(clt->sock->fd, &fds))
			return (struct nsock *)clt->sock;

	return NULL;
}

void
nsock_lst_map(struct nsock_listener *ls, nsock_lst_map_t map, void *data)
{
	const struct nsock_clt *clt;

	TAILQ_FOREACH(clt, &ls->head, link)
		map(clt->sock, data);
}

void
nsock_lst_free(struct nsock_listener *ls)
{
	struct nsock_clt *clt, *tmp;

	TAILQ_FOREACH_SAFE(clt, &ls->head, link, tmp)
		free(clt);

	free(ls);
}