/* msgtap support for Unbound */

/*
 * Copyright (c) 2021 The University of Queensland
 * Copyright (c) 2013-2014, Farsight Security, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "config.h"
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/un.h>

#include <msgtap.h>
#include <msgtap/dns.h>

#include "sldns/sbuffer.h"
#include "util/config_file.h"
#include "util/ub_event.h"
#include "util/net_help.h"
#include "util/netevent.h"
#include "util/log.h"

#include "msgtap/msgtap.h"
#include "msgtap/libmsgtap.h"

#ifndef nitems
#define nitems(_a)	(sizeof((_a)) / sizeof((_a)[0]))
#endif

struct mt_msg;

static int		 mt_connect(struct mt_env *);
static void		 mt_reconnect(struct mt_env *);

static struct mt_msg	*mt_msg_new(struct mt_env *, uint8_t, uint8_t);
static void		 mt_msg_send(struct mt_env *, struct mt_msg *,
			     struct sldns_buffer *);
static void		 mt_msg_free(struct mt_msg *);

static int		 mt_msg_add_bytes(struct mt_msg *, uint8_t, uint8_t,
			     const void *, size_t);
static int		 mt_msg_add_net(struct mt_msg *, uint8_t, uint8_t,
			    const struct sockaddr_storage *ss,
			    enum comm_point_type);

static void		 mt_log_addr(const struct sockaddr_storage *, socklen_t,
			     const char *);

static const char mt_name[] = "unbound";

static void
mt_apply_cfg(struct mt_env *env, struct config_file *cfg)
{
	if ((env->log_resolver_query_messages = (unsigned int)
	     cfg->msgtap_log_resolver_query_messages)) {
		verbose(VERB_OPS, "msgtap RESOLVER_QUERY enabled");
	}
	if ((env->log_resolver_response_messages = (unsigned int)
	     cfg->msgtap_log_resolver_response_messages))
	{
		verbose(VERB_OPS, "msgtap RESOLVER_RESPONSE enabled");
	}
	if ((env->log_client_query_messages = (unsigned int)
	     cfg->msgtap_log_client_query_messages))
	{
		verbose(VERB_OPS, "msgtap CLIENT_QUERY enabled");
	}
	if ((env->log_client_response_messages = (unsigned int)
	     cfg->msgtap_log_client_response_messages))
	{
		verbose(VERB_OPS, "msgtap CLIENT_RESPONSE enabled");
	}
	if ((env->log_forwarder_query_messages = (unsigned int)
	     cfg->msgtap_log_forwarder_query_messages))
	{
		verbose(VERB_OPS, "msgtap FORWARDER_QUERY enabled");
	}
	if ((env->log_forwarder_response_messages = (unsigned int)
	     cfg->msgtap_log_forwarder_response_messages))
	{
		verbose(VERB_OPS, "msgtap FORWARDER_RESPONSE enabled");
	}
}

struct mt_env *
mt_create(struct config_file *cfg)
{
	char hostname[MAXHOSTNAMELEN];
	struct mt_env *env;

	if (gethostname(hostname, sizeof(hostname)) == -1) {
		log_err("msgtap gethostname: %s", strerror(errno));
		return (NULL);
	}

	if (cfg->msgtap_socket_path == NULL ||
	    strlen(cfg->msgtap_socket_path) == 0) {
		log_err("msgtap: msgtap-socket-path not defined");
		return (NULL);
	}

	env = calloc(1, sizeof(*env));
	if (env == NULL)
		return (NULL);

	env->hostname = strdup(hostname);
	if (env->hostname == NULL) {
		log_err("msgtap hostname copy: %s", strerror(errno));
		goto free;
	}
	env->hostname_len = strlen(env->hostname);

	env->fd = -1;
	env->seq = 0;
	env->socket_path = strdup(cfg->msgtap_socket_path);
	if (env->socket_path == NULL) {
		log_err("msgtap: msgtap-socket-path copy: %s",
		    strerror(ENOMEM));
		goto free_hostname;
	}

	mt_apply_cfg(env, cfg);

	return (env);

free_hostname:
	free(env->hostname);
free:
	free(env);
	return (NULL);
}

/**
 * Initialize per-worker state.
 */
int
mt_init(struct mt_env *env, struct comm_base *cb)
{
	int fd;

	env->cb = cb;

	env->pid = getpid();
	env->thrid = getthrid();

	env->fd = mt_connect(env);
	if (env->fd == -1) {
		log_warn("msgtap-socket-path %s: %s",
		    env->socket_path, strerror(errno));
		mt_reconnect(env);
	}

	return (0);
}

/**
 * Deletes the per-worker state created by mt_init.
 */
void
mt_deinit(struct mt_env *env)
{
	if (env->reconnect_ev != NULL) {
		log_assert(env->fd == -1);
		ub_timer_del(env->reconnect_ev);
		ub_event_free(env->reconnect_ev);
		env->reconnect_ev = NULL;
	}
	if (env->fd != -1) {
		close(env->fd);
		env->fd = -1;
	}
}

/**
 * Delete msgtap environment object.
 */
void
mt_delete(struct mt_env *env)
{
	free(env->socket_path);
	free(env);
}

void
mt_msg_send_client_query(struct mt_env *env,
    const struct sockaddr_storage *ss, socklen_t sslen,
    enum comm_point_type cptype,
    struct sldns_buffer *qmsg)
{
	struct mt_msg *mt_msg;

	mt_msg = mt_msg_new(env, MSGTAP_T_NET_DIR_IN, DNSMSGTAP_MSGTYPE_CQ);
	if (mt_msg == NULL)
		return;

	if (mt_msg_add_net(mt_msg, MSGTAP_T_IPSRCADDR, MSGTAP_T_IPSRCPORT,
	    ss, cptype) == -1) {
		/* log */
		goto free;
	}

	mt_msg_send(env, mt_msg, qmsg);
free:
	mt_msg_free(mt_msg);
}

void
mt_msg_send_client_response(struct mt_env *env,
    const struct sockaddr_storage *ss, socklen_t sslen,
    enum comm_point_type cptype,
    struct sldns_buffer *rmsg)
{
	struct mt_msg *mt_msg;

	mt_msg = mt_msg_new(env, MSGTAP_T_NET_DIR_OUT, DNSMSGTAP_MSGTYPE_CR);
	if (mt_msg == NULL)
		return;

	if (mt_msg_add_net(mt_msg, MSGTAP_T_IPDSTADDR, MSGTAP_T_IPDSTPORT,
	    ss, cptype) == -1) {
		/* log */
		goto free;
	}

	mt_msg_send(env, mt_msg, rmsg);
free:
	mt_msg_free(mt_msg);
}

void
mt_msg_send_outside_query(struct mt_env *env,
    const struct sockaddr_storage *ss, socklen_t sslen,
    enum comm_point_type cptype,
    const uint8_t *zone, size_t zone_len,
    struct sldns_buffer *qmsg)
{
	struct mt_msg *mt_msg;
	uint16_t qflags;
	uint8_t type;

	qflags = sldns_buffer_read_u16_at(qmsg, 2);
	if (qflags & BIT_RD) {
		if (!env->log_forwarder_query_messages)
			return;
		type = DNSMSGTAP_MSGTYPE_FQ;
	} else {
		if (!env->log_resolver_query_messages)
			return;
		type = DNSMSGTAP_MSGTYPE_RQ;
	}

	mt_msg = mt_msg_new(env, MSGTAP_T_NET_DIR_OUT, type);
	if (mt_msg == NULL)
		return;

	if (mt_msg_add_net(mt_msg, MSGTAP_T_IPDSTADDR, MSGTAP_T_IPDSTPORT,
	    ss, cptype) == -1) {
		/* log */
		goto free;
	}

	if (mt_msg_add_bytes(mt_msg,
	    MSGTAP_CLASS_TYPED, DNSMSGTAP_QUERY_ZONE,
	    zone, zone_len) == -1) {
		/* log */
		goto free;
	}

	mt_msg_send(env, mt_msg, qmsg);
free:
	mt_msg_free(mt_msg);
}

void
mt_msg_send_outside_response(struct mt_env *env,
    const struct sockaddr_storage *ss, socklen_t sslen,
    enum comm_point_type cptype,
    const uint8_t *zone, size_t zone_len,
    const uint8_t *qbuf, size_t qbuf_len,
    const struct timeval *qtime, const struct timeval *rtime,
    struct sldns_buffer *rmsg)
{
	struct mt_msg *mt_msg;
	uint16_t qflags;
	uint8_t type;

	qflags = sldns_buffer_read_u16_at(rmsg, 2);
	if (qflags & BIT_RD) {
		if (!env->log_forwarder_response_messages)
			return;
		type = DNSMSGTAP_MSGTYPE_FR;
	} else {
		if (!env->log_resolver_response_messages)
			return;
		type = DNSMSGTAP_MSGTYPE_RR;
	}

	mt_msg = mt_msg_new(env, MSGTAP_T_NET_DIR_IN, type);
	if (mt_msg == NULL)
		return;

	if (mt_msg_add_net(mt_msg, MSGTAP_T_IPSRCADDR, MSGTAP_T_IPSRCPORT,
	    ss, cptype) == -1) {
		/* log */
		goto free;
	}

	if (mt_msg_add_bytes(mt_msg,
	    MSGTAP_CLASS_TYPED, DNSMSGTAP_QUERY_ZONE,
	    zone, zone_len) == -1) {
		/* log */
		goto free;
	}

	mt_msg_send(env, mt_msg, rmsg);
free:
	mt_msg_free(mt_msg);
}

struct mt_msg {
	void		*msg;
	size_t		 msglen;
};

static void *
mt_msg_md_realloc(struct mt_msg *mt_msg, size_t len)
{
	uint8_t *msg;
	size_t olen, nlen;

	if (len > 0xffff) {
		errno = EMSGSIZE;
		return (NULL);
	}

	olen = mt_msg->msglen;
	nlen = olen + sizeof(struct msgtap_metadata) + len;
	msg = realloc(mt_msg->msg, nlen);
	if (msg == NULL)
		return (NULL);

	mt_msg->msg = msg;
	mt_msg->msglen = nlen;

	return (msg + olen);
}

static int
mt_msg_add_ts(struct mt_msg *mt_msg, const struct timespec *ts)
{
	struct msgtap_metadata *md;
	uint64_t nsec;

	nsec = (ts->tv_sec * 1000000000ULL) + ts->tv_nsec;

	md = mt_msg_md_realloc(mt_msg, sizeof(nsec));
	if (md == NULL)
		return (-1);

	msgtap_md_u64(md, MSGTAP_CLASS_BASE, MSGTAP_T_TS, nsec);

	return (0);
}

static int
mt_msg_add_flag(struct mt_msg *mt_msg, uint8_t class, uint8_t type)
{
	struct msgtap_metadata *md;

	md = mt_msg_md_realloc(mt_msg, 0);
	if (md == NULL)
		return (-1);

	msgtap_md(md, class, type, 0);

	return (0);
}

static int
mt_msg_add_bytes(struct mt_msg *mt_msg, uint8_t class, uint8_t type,
    const void *buf, size_t buflen)
{
	struct msgtap_metadata *md;

	md = mt_msg_md_realloc(mt_msg, buflen);
	if (md == NULL)
		return (-1);

	msgtap_md_mem(md, class, type, buf, buflen);

	return (0);
}

static int
mt_msg_add_u32(struct mt_msg *mt_msg, uint8_t class, uint8_t type, uint32_t u32)
{
	struct msgtap_metadata *md;

	md = mt_msg_md_realloc(mt_msg, sizeof(u32));
	if (md == NULL)
		return (-1);

	msgtap_md_u32(md, class, type, u32);

	return (0);
}

static int
mt_msg_add_u16(struct mt_msg *mt_msg, uint8_t class, uint8_t type, uint16_t u16)
{
	struct msgtap_metadata *md;

	md = mt_msg_md_realloc(mt_msg, sizeof(u16));
	if (md == NULL)
		return (-1);

	msgtap_md_u16(md, class, type, u16);

	return (0);
}

static int
mt_msg_add_u8(struct mt_msg *mt_msg, uint8_t class, uint8_t type, uint8_t u8)
{
	struct msgtap_metadata *md;

	md = mt_msg_md_realloc(mt_msg, sizeof(u8));
	if (md == NULL)
		return (-1);

	msgtap_md_u8(md, class, type, u8);

	return (0);
}

static int
mt_msg_add_net(struct mt_msg *mt_msg, uint8_t cl_addr, uint8_t cl_port,
    const struct sockaddr_storage *ss, enum comm_point_type cptype)
{
	uint8_t iptype, ipproto;
	const void *addr;
	in_port_t port;
	size_t addrlen;

	switch (cptype) {
	case comm_udp:
		ipproto = MSGTAP_T_IPPROTO_UDP;
		break;
	case comm_tcp:
		ipproto = MSGTAP_T_IPPROTO_TCP;
		break;
	default:
		fatal_exit("%s: unsupported cptype %u", __func__, cptype);
		/* NOTREACHED */
	}

	switch (ss->ss_family) {
	case AF_INET: {
		const struct sockaddr_in *sin = (struct sockaddr_in *)ss;

		iptype = MSGTAP_T_IPV4;

		addr = &sin->sin_addr;
		addrlen = sizeof(sin->sin_addr);

		port = sin->sin_port;

		break;
	}
	case AF_INET6: {
		const struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *)ss;

		iptype = MSGTAP_T_IPV6;

		addr = &sin6->sin6_addr;
		addrlen = sizeof(sin6->sin6_addr);

		port = sin6->sin6_port;

		break;
	}
	default:
		fatal_exit("%s: unsupported address family %u", __func__,
		    ss->ss_family);
		/* NOTREACHED */
	}

	if (mt_msg_add_flag(mt_msg, MSGTAP_CLASS_BASE, iptype) == -1)
		return (-1);

	if (mt_msg_add_bytes(mt_msg, MSGTAP_CLASS_BASE, cl_addr,
	    addr, addrlen) == -1)
		return (-1);

	if (mt_msg_add_u8(mt_msg, MSGTAP_CLASS_BASE, MSGTAP_T_IPPROTO,
	    ipproto) == -1)
		return (-1);

	if (mt_msg_add_u16(mt_msg, MSGTAP_CLASS_BASE, cl_port, port) == -1)
		return (-1);

	return (0);
}

static struct mt_msg *
mt_msg_new(struct mt_env *env, uint8_t dir, uint8_t msgtype)
{
	struct mt_msg *mt_msg;
	struct msgtap_header *mh;
	struct msgtap_metadata *md;

	struct timespec ts;
	uint32_t seq;

	seq = env->seq++;

	if (env->fd == -1)
		return (NULL);

	if (clock_gettime(CLOCK_REALTIME, &ts) == -1) {
		log_warn("could not get time for msg %u seq %u: %s",
		    msgtype, seq, strerror(errno));
		return (NULL);
	}

	mt_msg = malloc(sizeof(*mt_msg));
	if (mt_msg == NULL)
		goto warn;

	mt_msg->msg = NULL;
	mt_msg->msglen = 0;

	if (mt_msg_add_ts(mt_msg, &ts) == -1)
		goto free;

	if (mt_msg_add_bytes(mt_msg, MSGTAP_CLASS_BASE, MSGTAP_T_HOSTNAME,
	    env->hostname, env->hostname_len) == -1)
		goto free;

	if (mt_msg_add_bytes(mt_msg, MSGTAP_CLASS_BASE, MSGTAP_T_NAME,
	    mt_name, sizeof(mt_name) - 1) == -1)
		goto free;

	if (mt_msg_add_u32(mt_msg,
	    MSGTAP_CLASS_BASE, MSGTAP_T_PID, env->pid) == -1)
		goto free;

	if (mt_msg_add_u32(mt_msg,
	    MSGTAP_CLASS_BASE, MSGTAP_T_TID, env->thrid) == -1)
		goto free;

	if (mt_msg_add_u32(mt_msg,
	    MSGTAP_CLASS_BASE, MSGTAP_T_SEQ32, seq) == -1)
		goto free;

	if (mt_msg_add_u8(mt_msg,
	    MSGTAP_CLASS_TYPED, DNSMSGTAP_MSGTYPE, msgtype) == -1)
		goto free;

	if (mt_msg_add_u8(mt_msg,
	    MSGTAP_CLASS_BASE, MSGTAP_T_NET_DIR, dir) == -1)
		goto free;

	return (mt_msg);

free:
	free(mt_msg->msg);
free_mt_msg:
	free(mt_msg);
warn:
	log_warn("unable to allocate buffer for msg %u seq %u: %s",
	    msgtype, seq, strerror(errno));

	return (NULL);
}

static void
mt_msg_send(struct mt_env *env, struct mt_msg *mt_msg,
    struct sldns_buffer *buf)
{
	struct msgtap_header mh;
	struct iovec iov[3];
	struct msghdr msg = {
		.msg_iov = iov,
		.msg_iovlen = nitems(iov),
	};
	uint32_t buflen;
	ssize_t rv;

	if (mt_msg->msglen > 0xffffffff) {
		/* metadata is too long */
		return;
	}
	if (sldns_buffer_limit(buf) > 0xffffffff) {
		/* buffer is too long */
		return;
	}
	buflen = sldns_buffer_limit(buf);

	mh.mh_flags = htons(MSGTAP_F_VERSION_0);
	mh.mh_type = htons(MSGTAP_TYPE_DNS);
	mh.mh_metalen = htonl(mt_msg->msglen);
	mh.mh_msglen = htonl(buflen);
	mh.mh_caplen = htonl(buflen);

	iov[0].iov_base = &mh;
	iov[0].iov_len = sizeof(mh);
	iov[1].iov_base = mt_msg->msg;
	iov[1].iov_len = mt_msg->msglen;
	iov[2].iov_base = sldns_buffer_begin(buf);
	iov[2].iov_len = buflen;

	rv = sendmsg(env->fd, &msg, 0);
	if (rv == -1) {
		switch (errno) {
		case ENOBUFS:
			/* oh well */
			return;
		case EPIPE:
			/* we'll try and fix this */
			break;
		default:
			log_warn("msgtap send: %s", strerror(errno));
		}

		close(env->fd);
		env->fd = mt_connect(env);
		if (env->fd == -1) {
			log_warn("msgtap retry: %s", strerror(errno));
			mt_reconnect(env);
			return;
		}

		rv = sendmsg(env->fd, &msg, 0);
		if (rv == -1) {
			log_warn("msgtap send: %s", strerror(errno));
			close(env->fd);
			env->fd = -1;
			mt_reconnect(env);
			return;
		}
	}
}

static void
mt_msg_free(struct mt_msg *mt_msg)
{
	free(mt_msg->msg);
	free(mt_msg);
}

/*
 *
 */

static void	mt_reconnect_cb(int, short, void *);

static int
mt_connect(struct mt_env *env)
{
	struct sockaddr_un sun = { .sun_family = AF_UNIX };
	int fd = -1;

	if (strlcpy(sun.sun_path, env->socket_path,
	    sizeof(sun.sun_path)) >= sizeof(sun.sun_path)) {
		errno = ENAMETOOLONG;
		return (-1);
	}

	fd = socket(AF_UNIX, SOCK_SEQPACKET | SOCK_NONBLOCK, 0);
	if (fd == -1)
		return (-1);

	if (connect(fd, (struct sockaddr *)&sun, sizeof(sun)) == -1) {
		close(fd);
		return (-1);
	}

	/* all good */
	return (fd);
}

static void
mt_reconnect_add(struct mt_env *env)
{
	struct timeval tv;

	tv.tv_sec = 1;
	tv.tv_usec = arc4random_uniform(1000000);

	if (ub_timer_add(env->reconnect_ev, comm_base_internal(env->cb),
	    mt_reconnect_cb, env, &tv) != 0)
		fatal_exit("msgtap reconnect: timer add failed");
}

static void
mt_reconnect_cb(int zero, short events, void *arg)
{
	struct mt_env *env = arg;
	int fd;

	fd = mt_connect(env);
	if (fd == -1) {
		mt_reconnect_add(env);
		return;
	}

	verbose(VERB_OPS, "msgtap %s reconnected (seq %u)",
	    env->socket_path, env->seq);

	ub_event_free(env->reconnect_ev);
	env->reconnect_ev = NULL;

	env->fd = fd;
}

static void
mt_reconnect(struct mt_env *env)
{
	env->reconnect_ev = ub_event_new(comm_base_internal(env->cb),
	    0, UB_EV_TIMEOUT, mt_reconnect_cb, env);
	if (env->reconnect_ev == NULL)
		fatal_exit("msgtap reconnect: %s", strerror(ENOMEM));

	mt_reconnect_add(env);
}