/* util.c - miscellaneous utility functions
 *
 * Copyright (C) 2016  Luke Shumaker
 *
 * 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 3 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 <ctype.h> /* for isdigit */
#include <errno.h>
#include <error.h>
#include <netdb.h> /* for {get,free}addrinfo() */
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/stat.h> /* for open */
#include <sys/un.h>
#include <unistd.h> /* for unlink */

#include "util.h"

void *xrealloc(void *ptr, size_t size)
{
	void *ret = realloc(ptr, size);
	if (ret == NULL) {
		if (ptr==NULL)
			error(1, errno, "Could not allocate memory");
		else
			error(1, errno, "Could not re-allocate memory");
	}
	return ret;
}

bool is_numeric(const char *str) {
	for (size_t i = 0; str[i] != '\0'; i++)
		if (!isdigit(str[i]))
			return false;
	if (str[0] == '\0')
		return false;
	return true;
}	

int get_fd(const char *addr) {
	int sock;
	if (strcmp(addr, "stdin") == 0) {
		sock = 0;
	} else if (strcmp(addr, "stdout") == 0) {
		sock = 1;
	} else if (strcmp(addr, "stderr") == 0) {
		sock = 2;
	} else if (strncmp(addr, "systemd", strlen("systemd")) == 0) {
		sock = 3; /* <systemd/sd-daemon.h>:SD_LISTEN_FDS_START */
		addr = &addr[strlen("systemd")];
		switch (addr[0]) {
		case '\0':
			/* do nothing */
			break;
		case ':':
			addr = &addr[1];
			if (is_numeric(addr)) {
				sock += atoi(addr);
			} else {
				const char *e = getenv("LISTEN_FDNAMES");
				if (e == NULL)
					return -ENOTCONN;
				char *names = strdupa(e);
				char *name = NULL;
				int i = -1;
				do {
					name = strsep(&names, ":");
					i++;
				} while (name != NULL && strcmp(name, addr) != 0);
				if (name == NULL)
					return -ENOENT;
				sock += i;
			}
		}
	} else {
		if (!is_numeric(addr))
			return -EINVAL;
		sock = atoi(addr);
	}
	return sock;
}

FILE *xfopen(const char *path, const char *mode) {
	FILE *stream = fopen(path, mode);
	if (stream == NULL)
		error(EXIT_FAILURE, errno, "fopen: %s", path);
	return stream;
}

FILE *xfdopen(const char *path, const char *mode) {
	int fd = get_fd(path);
	if (fd < 0)
		error(EXIT_FAILURE, -fd, "get_fd: %s", path);
	FILE *stream = fdopen(fd, mode);
	if (stream == NULL)
		error(EXIT_FAILURE, errno, "fdopen: %s (%d)", path, fd);
	return stream;
}

/* same error codes values as -getaddrinfo(); */
int sockstream_listen(const char *type, const char *addr) {
	int sock;
	if (strcmp(type, "fd") == 0) {
		sock = get_fd(addr);
		if (sock < 0) {
			errno = -sock;
			return -EAI_SYSTEM;
		}
		/* make sure it's a socket */
		struct stat st;
		if (fstat(sock, &st) < 0)
			return -EAI_SYSTEM;
		if (!S_ISSOCK(st.st_mode)) {
			errno = ENOTSOCK;
			return -EAI_SYSTEM;
		}
		/* make sure the socket is a stream */
		int fd_socktype = 0;
		socklen_t l = sizeof(fd_socktype);
		if (getsockopt(sock, SOL_SOCKET, SO_TYPE, &fd_socktype, &l) < 0)
			return -EAI_SYSTEM;
		if (l != sizeof(fd_socktype)) {
			errno = EINVAL;
			return -EAI_SYSTEM;
		}
		if (fd_socktype != SOCK_STREAM) {
			errno = ENOSTR;
			return -EAI_SYSTEM;
		}
		/* make sure the socket is listening */
		int fd_listening = 0;
		l = sizeof(fd_listening);
		if (getsockopt(sock, SOL_SOCKET, SO_ACCEPTCONN, &fd_listening, &l) < 0)
			return -EAI_SYSTEM;
		if (l != sizeof(fd_listening)) {
			errno = EINVAL;
			return -EAI_SYSTEM;
		}
		if (!fd_listening)
			listen(sock, SOMAXCONN);
		/* return the socket */
		return sock;
	} else if (strcmp(type, "unix") == 0) {
		struct sockaddr_un un_addr = { 0 };
		if (strlen(addr)+1 > sizeof(un_addr.sun_path)) {
			errno = ENAMETOOLONG;
			return -EAI_SYSTEM;
		}
		un_addr.sun_family = AF_UNIX;
		strcpy(un_addr.sun_path, addr);
		int sock = socket(AF_UNIX, SOCK_STREAM, 0);
		if (sock < 0)
			return -EAI_SYSTEM;
		unlink(addr);
		if (bind(sock, (struct sockaddr *)&un_addr, sizeof(un_addr)) < 0)
			return -EAI_SYSTEM;
		if (listen(sock, SOMAXCONN) < 0)
			return -EAI_SYSTEM;
		return sock;
	} else if (strcmp(type, "tcp4") == 0 || strcmp(type, "tcp6") == 0) {
		struct addrinfo *ai_addr = NULL;
		struct addrinfo ai_hints = { 0 };

		char *host = strdupa(addr);
		char *col = strrchr(host, ':');
		if (col == NULL) {
			errno = EINVAL;
			return -EAI_SYSTEM;
		}
		char *port = &col[1];
		*col = '\0';
		if (host[0] == '\0')
			host = NULL;

		ai_hints.ai_family = (strcmp(type, "tcp6") == 0) ? AF_INET6 : AF_INET;
		ai_hints.ai_socktype = SOCK_STREAM;
		ai_hints.ai_flags = AI_PASSIVE;

		int r = getaddrinfo(host, port, &ai_hints, &ai_addr);
		if (r < 0)
			return r;
		if (r > 0)
			return -r;

		int sock = socket(ai_addr->ai_family, ai_addr->ai_socktype, ai_addr->ai_protocol);
		if (sock < 0)
			return -EAI_SYSTEM;

		int yes = 1;
		setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

		if (bind(sock, ai_addr->ai_addr, ai_addr->ai_addrlen) < 0)
			return -EAI_SYSTEM;

		if (listen(sock, SOMAXCONN) < 0)
			return -EAI_SYSTEM;

		freeaddrinfo(ai_addr);
		return sock;
	} else {
		errno = EINVAL;
		return -EAI_SYSTEM;
	}
}