/*
 * Copyright 2009 Red Hat, Inc.
 *
 * 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.
 *
 * 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; see the file COPYING.  If not, write to
 * the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/nameser.h>
#include <netdb.h>
#include <resolv.h>
#include "getsrvinfo.h"

#define __must_be_array(a) \
	(__builtin_types_compatible_p(typeof(a), typeof(&a[0])))
#define ARRAY_SIZE(arr) (sizeof(arr) / sizeof((arr)[0]) + __must_be_array(arr))

enum {
	gsi_max_dom_name_sz		= 64,
	gsi_dns_buf_sz			= 1024,
};

static int fill_srvinfo(struct srvinfo *si,
			const struct addrinfo *hints,
			unsigned int gsi_flags,
			unsigned int priority,
			unsigned int weight, unsigned int port,
			unsigned int nlen, const char *name)
{
	char portstr[11];
	char *hostname;
	struct addrinfo *res0 = NULL;
	int rc = 0;

	sprintf(portstr, "%u", port);

	hostname = malloc(nlen + 1);
	if (!hostname)
		return -ENOMEM;

	memcpy(hostname, name, nlen);
	hostname[nlen] = 0;

	if (!(gsi_flags & FSI_NO_ADDR)) {
		rc = getaddrinfo(hostname, portstr, hints, &res0);

		/* because we do not wish failure of one lookup
		 * to cause failure of the entire operation,
		 * we simply record the lack of an address (NULL)
		 * and move on.
		 */
		if (rc)
			res0 = NULL;
	}

	si->si_prio = priority;
	si->si_weight = weight;
	si->si_target = hostname;
	si->si_port = port;

	si->si_addr = res0;
	si->si_next = NULL;

	return 0;
}

void freesrvinfo(struct srvinfo *res)
{
	while (res) {
		struct srvinfo *tmp;
		
		tmp = res;
		res = res->si_next;

		if (tmp->si_addr)
			freeaddrinfo(tmp->si_addr);
		free(tmp->si_target);
		free(tmp);
	}
}

int getsrvinfo(const char *service_name, const char *domain_name,
	       const struct addrinfo *hints, unsigned int gsi_flags,
	       struct srvinfo **res_out)
{
	unsigned char resp[gsi_dns_buf_sz];
	int rlen;
	ns_msg nsb;
	ns_rr rrb;
	int rrlen;
	char hostb[gsi_max_dom_name_sz];
	struct srvinfo *si, *res = NULL, *res_last = NULL;
	const unsigned char *p;
	int rc, gsi_rc, i;

	if (!domain_name || !res_out)
		return ESI_INVAL;

	*res_out = NULL;

	/* we concatencate service_name and domain_name as a helpful
	 * service for the caller, because it is very common
	 * that service_name is either completely static, or at least
	 * stored in a separate variable from domain_name.
	 */
	if (service_name) {
		char *name;
		size_t name_len;
		int has_dot;

		has_dot = (service_name[strlen(service_name) - 1] == '.');
		name_len = strlen(service_name) + strlen(domain_name) +
			   (has_dot ? 0 : 1) + 1;

		name = malloc(name_len);
		if (!name)
			return ESI_OOM;

		snprintf(name, name_len, "%s%s%s", service_name,
			     has_dot ? "" : ".",
			     domain_name);
			
		rc = res_search(name, ns_c_in, ns_t_srv, resp, sizeof(resp));

		free(name);
	} else {
		rc = res_search(domain_name, ns_c_in, ns_t_srv,
				resp, sizeof(resp));
	}

	/* parse resolver return value */
	if (rc < 0) {
		switch (h_errno) {
		case TRY_AGAIN:
			return ESI_AGAIN;
		case HOST_NOT_FOUND:
		case NO_DATA:
		case NO_RECOVERY:
		default:
			return ESI_FAIL;
		}
	}
	rlen = rc;

	if (rlen == 0)
		return ESI_FAIL;

	/* set up DNS result parse */
	if (ns_initparse(resp, rlen, &nsb) < 0)
		return ESI_CORRUPT;

	/* iterate through each answer.  Because DNS packets may
	 * be truncated, we do not signal an error on
	 * short-length faults found during packet parsing
	 */
	for (i = 0; i < ns_msg_count(nsb, ns_s_an); i++) {
		rc = ns_parserr(&nsb, ns_s_an, i, &rrb);
		if (rc < 0)
			continue;

		if (ns_rr_class(rrb) != ns_c_in)
			continue;

		switch (ns_rr_type(rrb)) {
		case ns_t_srv:
			rrlen = ns_rr_rdlen(rrb);
			if (rrlen < 8) {	/* 2+2+2 and 2 for host */
				break;
			}
			p = ns_rr_rdata(rrb);
			rc = dn_expand(resp, resp+rlen, p+6,
				       hostb, gsi_max_dom_name_sz);
			if (rc < 0) {
				break;
			}
			if (rc < 2) {
				break;
			}

			si = malloc(sizeof(*si));
			if (!si) {
				gsi_rc = ESI_OOM;
				goto err_out;
			}

			if (fill_srvinfo(si, hints, gsi_flags,
					 ns_get16(p+0),
					 ns_get16(p+2),
					 ns_get16(p+4),
					 rc, hostb)) {
				free(si);
				gsi_rc = ESI_OOM;
				goto err_out;
			}

			/* if first item, set BOL-ptr */
			if (!res)
				res = si;

			/* append to EOL */
			if (res_last)
				res_last->si_next = si;
			res_last = si;

			break;

		case ns_t_cname:	/* impossible, but ... ? */
		default:
			break;
		}
	}

	*res_out = res;
	return ESI_NONE;

err_out:
	freesrvinfo(res);
	return gsi_rc;
}

static const char *gsi_error_str[] = {
	[ESI_NONE]		= "no error",
	[ESI_CORRUPT]		= "server returned bad data",
	[ESI_FAIL]		= "server returned permanent failure",
	[ESI_AGAIN]		= "server returned temporary failure",
	[ESI_OOM]		= "internal memory alloc failed",
	[ESI_INVAL]		= "invalid argument(s)",
};

const char *gsi_strerror(int errcode)
{
	if (errcode < 0 || errcode >= ARRAY_SIZE(gsi_error_str))
		return NULL;
	
	return gsi_error_str[errcode];
}