/* This file is part of varnish-mib -*- c -*-
   Copyright (C) 2014-2019 Sergey Poznyakoff

   Varnish-mib 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, or (at your option)
   any later version.

   Varnish-mib 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 varnish-mib.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "varnish_mib.h"
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/ioctl.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <errno.h>

struct vcli_sockaddr {
	socklen_t len;
	struct sockaddr addr[1];
};

#define ISSPACE(c) ((c)==' '||(c)=='\t'||(c)=='\n')

vcli_sockaddr_t 
vcli_parse_sockaddr(char const *arg)
{
	char *node;
	size_t n;
	int rc;
	struct addrinfo *res;
	char const *p;
	vcli_sockaddr_t sa;
	
	n = strcspn(arg, ": \t");
	node = malloc(n + 1);
	if (!node) {
		snmp_log(LOG_ERR, "out of memory\n");
		return NULL;
	}

	memcpy(node, arg, n);
	node[n] = 0;
	p = arg + n;
	if (*p == ':')
		p++;
	else
		while (*p && ISSPACE(*p))
			p++;

	if (!p)
		p = "6082";
	
	rc = getaddrinfo(node, p, NULL, &res);
	free(node);
	if (rc) {
		snmp_log(LOG_ERR, "can't resolve %s\n", arg);
		return NULL;
	}

	sa = malloc(sizeof(*sa) - sizeof(sa->addr) + res->ai_addrlen);
	if (!sa) {
		snmp_log(LOG_ERR, "out of memory\n");
		return NULL;
	}
	sa->len = res->ai_addrlen;
	memcpy(sa->addr, res->ai_addr, res->ai_addrlen);
	freeaddrinfo(res);
	return sa;
}

vcli_sockaddr_t
vcli_sockaddr_dup(vcli_sockaddr_t src)
{
	vcli_sockaddr_t dst;
	dst = malloc(sizeof(*dst) - sizeof(dst->addr) + src->len);
	memcpy(dst, src, src->len);
	return dst;
}

#define VCLI_INIT_ALLOC 16

static int
vcli_alloc(struct vcli_conn *conn, size_t size)
{
	char *p;

	if (size < conn->bufmax)
		return 0;
	p = realloc(conn->base, size);
	if (!p) {
		snmp_log(LOG_ERR, "out of memory\n");
		return -1;
	}
	conn->base = p;
	conn->bufmax = size;
	return 0;
}
	
static int
vcli_extra_size(struct vcli_conn *conn, size_t incr)
{
	if (incr + conn->bufsize > conn->bufmax) {
		size_t size;
		if (conn->bufmax == 0)
			size = incr > VCLI_INIT_ALLOC
				? incr
				: VCLI_INIT_ALLOC;
		else {
			for (size = conn->bufmax; size < conn->bufsize + incr;
			     size *= 2)
				if (size <= conn->bufmax) {
					snmp_log(LOG_ERR, "out of memory\n");
					return -1;
				}
		}
		return vcli_alloc(conn, size);
	}
	return 0;
}

static int
vcli_read(struct vcli_conn *conn, size_t size)
{
	fd_set rd;
	time_t start;
	struct timeval tv;
	long ttl;
	int ret = 0;
	int rc;
	
	++size;
	rc = vcli_alloc(conn, size + 1);
	if (rc)
		return SNMP_ERR_GENERR;

	conn->bufsize = 0;
	time(&start);
	while (size) {
		FD_ZERO(&rd);
		FD_SET(conn->fd, &rd);
		
		ttl = vcli_timeout - (time(NULL) - start);
		if (ttl <= 0) {
			snmp_log(LOG_ERR, "timed out reading from varnish\n");
			ret = -1;
			break;
		}

		tv.tv_sec = ttl;
		tv.tv_usec = 0;
		rc = select(conn->fd + 1, &rd, NULL, NULL, &tv);
		if (rc < 0) {
			if (errno == EINTR || errno == EAGAIN)
				continue;
			snmp_log(LOG_ERR, "select: %s\n", strerror(errno));
			ret = -1;
			break;
		}

		if (FD_ISSET(conn->fd, &rd)) {
			int n;
			
			if (ioctl(conn->fd, FIONREAD, &n) < 0) {
				snmp_log(LOG_ERR, "ioctl: %s\n",
					 strerror(errno));
				ret = -1;
				break;
			}
			if (n > size)
				n = size;
			rc = read(conn->fd, conn->base + conn->bufsize, n);
			if (rc > 0) {
				conn->bufsize += rc;
				size -= rc;
			} else if (rc == 0
				   || errno == EINTR || errno == EAGAIN) {
				continue;
			} else {
				snmp_log(LOG_ERR, "read: %s\n",
					 strerror(errno));
				ret = -1;
				break;
			}
		}
	}

	if (ret == 0) {
		if (conn->bufsize == 0)
			ret = -1;
		conn->base[conn->bufsize] = 0;
		DEBUGMSGTL(("vcli:transcript", "<<varnish: %s\n", conn->base));
	}
	
	return ret;
}

static int
vcli_getline(struct vcli_conn *conn)
{
	fd_set rd;
	time_t start;
	struct timeval tv;
	long ttl;
	int ret = 0;
	int rc;

	conn->bufsize = 0;
	time(&start);
	while (1) {
		FD_ZERO(&rd);
		FD_SET(conn->fd, &rd);
		
		ttl = vcli_timeout - (time(NULL) - start);
		if (ttl <= 0) {
			snmp_log(LOG_ERR, "timed out reading from varnish\n");
			ret = -1;
			break;
		}

		tv.tv_sec = ttl;
		tv.tv_usec = 0;
		rc = select(conn->fd + 1, &rd, NULL, NULL, &tv);
		if (rc < 0) {
			if (errno == EINTR || errno == EAGAIN)
				continue;
			snmp_log(LOG_ERR, "select: %s\n", strerror(errno));
			ret = -1;
			break;
		}

		if (FD_ISSET(conn->fd, &rd)) {
			char c;

			rc = read(conn->fd, &c, 1);
			if (rc == 1) {
				if (vcli_extra_size(conn, 1)) {
					ret = -1;
					break;
				}
				conn->base[conn->bufsize++] = c;
				if (c == '\n')
					break;
			} else if (rc == 0
				   || errno == EINTR || errno == EAGAIN) {
				continue;
			} else {
				snmp_log(LOG_ERR, "read: %s\n",
					 strerror(errno));
				ret = -1;
				break;
			}
		}
	}

	if (ret == 0) {
		if (conn->bufsize == 0)
			ret = -1;
		else if (conn->base[conn->bufsize-1] == '\n') {
			conn->bufsize--;
			if (conn->base[conn->bufsize-1] == '\r')
				conn->bufsize--;
		}
		conn->base[conn->bufsize] = 0;
	}
	
	return ret;
}

int
vcli_write(struct vcli_conn *conn)
{
	size_t size;

	DEBUGMSGTL(("vcli:transcript", ">>varnish: %s\n", conn->base));
	for (size = 0; size < conn->bufsize; ) {
		int n = write(conn->fd, conn->base + size,
			      conn->bufsize - size);
		if (n < 0) {
			snmp_log(LOG_ERR, "write error: %s\n",
				 strerror(errno));
			return -1;
		}
		size += n;
	}
	return 0;
}

int
vcli_read_response(struct vcli_conn *conn)
{
	char *p;
	unsigned long n;
	
	if (vcli_getline(conn)) {
		vcli_disconnect(conn);
		return SNMP_ERR_GENERR;
	}

	n = strtoul(conn->base, &p, 10);
	conn->resp = n;
	if (!ISSPACE(*p)) {
		snmp_log(LOG_ERR, "unrecognized response from Varnish: %s\n",
			 conn->base);
		return SNMP_ERR_GENERR;
	}
	while (*p && ISSPACE(*p))
		++p;
	n = strtoul(p, &p, 10);
	if (n > 0)
		return vcli_read(conn, n);
	return 0;
}

int
vcli_vasprintf(struct vcli_conn *conn, const char *fmt, va_list ap)
{
	int rc = 0;
  
	rc = vcli_alloc(conn, VCLI_INIT_ALLOC);
	if (rc)
		return -1;
  
	for (;;) {
		va_list aq;
		ssize_t n;

		va_copy(aq, ap);
		n = vsnprintf(conn->base, conn->bufmax, fmt, aq);
		va_end(aq);
		if (n < 0 || n >= conn->bufmax ||
		    !memchr(conn->base, '\0', n + 1)) {
			size_t newlen = conn->bufmax * 2;
			if (newlen < conn->bufmax) {
				snmp_log(LOG_ERR, "out of memory\n");
				rc = -1;
				break;
			}
			rc = vcli_alloc(conn, newlen);
			if (rc)
				break;
		} else {
			conn->bufsize = n;
			break;
		}
	}

	return rc;
}  

int
vcli_asprintf(struct vcli_conn *conn, const char *fmt, ...)
{
	int rc;
	va_list ap;
	
	va_start(ap, fmt);
	rc = vcli_vasprintf(conn, fmt, ap);
	va_end(ap);
	return rc;
}

static int
open_socket(vcli_sockaddr_t sa)
{
	int fd;
	struct timeval start, connect_tv;
	
	fd = socket(sa->addr->sa_family, SOCK_STREAM, 0);
	if (fd == -1) {
		snmp_log(LOG_ERR, "socket: %s\n", strerror(errno));
		return -1;
	}

	connect_tv.tv_sec = vcli_timeout;
	connect_tv.tv_usec = 0;
	
	gettimeofday(&start, NULL);
	for (;;) {
		int rc = connect(fd, sa->addr, sa->len);
		if (rc == 0)
			break;
		else if (errno == ECONNREFUSED) {
			struct timeval tv, now;
			fd_set rset, wset, xset;
			
			gettimeofday(&now, NULL);
			timersub(&now, &start, &tv);
			
			if (timercmp(&tv, &connect_tv, < )) {
				FD_ZERO(&rset);
				FD_SET(fd, &rset);
				FD_ZERO(&wset);
				FD_SET(fd, &wset);
				FD_ZERO(&xset);
				FD_SET(fd, &xset);
				select(fd + 1, &rset, &wset, &xset, &tv);
				continue;
			}
		} else {
			char hbuf[NI_MAXHOST], sbuf[NI_MAXSERV];
			int ec = errno;
			if (getnameinfo(sa->addr, sa->len,
					hbuf, sizeof(hbuf), sbuf, sizeof(sbuf),
					NI_NUMERICHOST | NI_NUMERICSERV) == 0)
				snmp_log(LOG_ERR,
					 "cannot connect to %s:%s: %s\n",
					 hbuf, sbuf, strerror(ec));
			else
				snmp_log(LOG_ERR,
					 "cannot connect to varnish: %s\n",
					 strerror(ec));
			close(fd);
			return -1;
		}
	}

	return fd;
}

void
vcli_disconnect(struct vcli_conn *conn)
{
	close(conn->fd);
	free(conn->base);
	free(conn->secret);
	memset(conn, 0, sizeof(*conn));
}

static int
vcli_handshake(struct vcli_conn *conn)
{
	if (vcli_read_response(conn))
		return 1;

	if (conn->resp == CLIS_AUTH) {
		char buf[CLI_AUTH_RESPONSE_LEN + 1];
		char *p = strchr(conn->base, '\n');

		if (!p) {
			snmp_log(LOG_ERR,
				 "unrecognized response from Varnish: %s\n",
				 conn->base);
			return SNMP_ERR_GENERR;
		}
		*p = 0;
		if (varnish_auth_response(conn->secret, conn->base, buf))
			return 1;
		
		if (vcli_asprintf(conn, "auth %s\n", buf) ||
		    vcli_write(conn))
			return 1;

		if (vcli_read_response(conn))
			return 1;
	}

	if (conn->resp != CLIS_OK) {
		snmp_log(LOG_ERR, "Varnish connection rejected: %u %s\n",
			 conn->resp, conn->base);
		return 1;
	}

	if (vcli_asprintf(conn, "ping\n") || vcli_write(conn))
		return 1;

	if (vcli_read_response(conn))
		return 1;

	if (conn->resp != CLIS_OK || strstr(conn->base, "PONG") == NULL) {
		snmp_log(LOG_ERR, "expected PONG, but got %s\n",
			 conn->base);
		return 1;
	}
	
	return 0;
}

static int
parse_connection(struct vsm *vsm, vcli_sockaddr_t *addr, char **secret)
{
	if (vcli_sockaddr) {
		DEBUGMSGTL(("vcli", "using configured sockaddr\n"));
		*addr = vcli_sockaddr_dup(vcli_sockaddr);
	} else {
		char *T_arg, *p;
		T_arg = VSM_Dup(vsm, "Arg", "-T");
		if (!T_arg) {
			snmp_log(LOG_ERR, "no -T arg in shared memory\n");
			return SNMP_ERR_GENERR;
		}
		p = T_arg + strlen(T_arg) - 1;
		if (*p == '\n')
			*p = 0;
		DEBUGMSGTL(("vcli", "-T '%s'\n", T_arg));
		*addr = vcli_parse_sockaddr(T_arg);
		free(T_arg);
	}
	if (!*addr)
		return SNMP_ERR_GENERR;

	if (vcli_secret) {
		DEBUGMSGTL(("vcli", "using configured secret\n"));
		*secret = strdup(vcli_secret);
		if (!*secret) {
			snmp_log(LOG_ERR, "out of memory");
			return SNMP_ERR_GENERR;
		}
	} else {
		char *S_arg = VSM_Dup(vsm, "Arg", "-S");
		if (!S_arg) {
			snmp_log(LOG_ERR, "no -S arg in shared memory\n");
			return SNMP_ERR_GENERR;
		}
		DEBUGMSGTL(("vcli", "-S '%s'\n", S_arg));
		*secret = S_arg;
        }
	return 0;
}

int
vcli_connect(struct vsm *vsm, struct vcli_conn *conn)
{
	vcli_sockaddr_t addr = NULL;
	char *secret = NULL;
	int rc;
	
	rc = parse_connection(vsm, &addr, &secret);
	if (rc == 0) {
		memset(conn, 0, sizeof(*conn));
		conn->fd = open_socket(addr);
		if (conn->fd == -1)
			rc = SNMP_ERR_GENERR;
		else {
			conn->secret = secret;
			secret = NULL;
			if (vcli_handshake(conn)) {
				vcli_disconnect(conn);
				rc = SNMP_ERR_GENERR;
			}
		}
	}
	free(addr);
	free(secret);
	return rc;
}