// DNS support -- return information about LDAP servers,
// by looking up a SRV record.
//
// SRV records, as discussed in RFC 2782, point to locations of
// services.  In the words of that RFC,
//
//     If a SRV-cognizant LDAP client wants to discover a LDAP server that
//     supports TCP protocol and provides LDAP service for the domain
//     example.com., it does a lookup of
//
//         _ldap._tcp.example.com
//
// The get_sorted_srv_records() function below will take a domain example.com,
// prepend the "_ldap._tcp",
// do the lookup to obtain the SRV record,
// and return a string containing the hosts and ports,
// sorted as described in the RFC.
//
// It's unexpectedly hard to find clear documentation about DNS lookups,
// but the following does seem to be portable,
// though without much in the way of intelligible manpages on any platform.
//
// See:
//   Scrappy documentation: https://docstore.mik.ua/orelly/networking_2ndEd/dns/ch15_02.htm
//   Wikipedia: https://en.wikipedia.org/wiki/List_of_DNS_record_types
//   RFC 1035, Domain names -- implementation and specification
//
// This module exposes a function
//
//    char* get_sorted_srv_records(const char* domain);
//
// which returns a space-separated list of ldap:// URIs obtained from
// the SRV record corresponding to the given domain, in a random order
// appropriate to the weights obtained.
//
// This can be compiled as a standalone (test?) program; see STANDALONE below.


// CentOS features.h requires _BSD_SOURCE in order to get types.h to define u_char,
// which resolv.h refers to.
#define _BSD_SOURCE 1
// Debian prefers _DEFAULT_SOURCE, and issues a deprecation warning
// if _BSD_SOURCE is defined and _DEFAULT_SOURCE isn't.
#define _DEFAULT_SOURCE 1
// ... and _XOPEN_SOURCE to get getopt
// (on eg CentOS, this also happens with _POSIX_C_SOURCE >= 2, but
// that _prevents_ macOS types.h from defining u_char).
#define _XOPEN_SOURCE 1

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <errno.h>

#include "config.h"

// See the docs for autoconf AC_HEADER_RESOLV
#ifdef HAVE_SYS_TYPES_H
#  include <sys/types.h>
#endif
#ifdef HAVE_NETINET_IN_H
#  include <netinet/in.h>   /* inet_ functions / structs */
#endif
#ifdef HAVE_ARPA_NAMESER_H
#  include <arpa/nameser.h> /* DNS HEADER struct */
#endif
#ifdef HAVE_NETDB_H
#  include <netdb.h>
#endif
#include <resolv.h>

// There is a test main program at the bottom of this file.
// To compile:
//      cc -Wall -std=c99 -o dns-support -DSTANDALONE=1 -lresolv dns-support.c
//
#ifndef STANDALONE
#define STANDALONE 0
#endif

#if STANDALONE
#include <stdarg.h>
#include <string.h>

// dummy out functions provided by the other modules in this kit
static int _chatter = 2;
int get_verbosity(void) { return _chatter; }
void change_verbosity(int inc) { _chatter += inc; }
void log_err_message(void* dummy, const char* fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);
    vfprintf(stderr, fmt, ap);
    va_end(ap);
}
#else
// we call get_verbosity and log_err_message
#include "ldaputil.h"
#include "ldap-config.h"
#endif

typedef unsigned char byte;

// Return a uniform random deviate in [0,1)
//
// random(3) returns a number in [0, 2^31-1].
// We convert this to a double by dividing it by 2^31 exactly.
// Thus the return value is always strictly less than one.
static double random_f(void)
{
    static byte initialised_p = 0;
    static double denom;

    if (! initialised_p) {
        srandom(time(NULL));
        denom = ldexp(1, 31);
        initialised_p = 1;

        // Explore a few properties of 'denom'
        // int exp;
        // double mantissa = frexp(denom, &exp);
        // printf("%g -> %f.2^%d;  ", denom, mantissa, exp);
        // double two31 = (double)2147483647; // 2^31-1 as an integer
        // mantissa = frexp(two31, &exp);
        // printf("%g -> %f.2^%d;  ", two31, mantissa, exp);
        // printf("ratio = %g, less than one? %s\n",
        //        two31/denom, two31/denom < 1.0 ? "yes" : "no");
    }

    double univariate = (double)random();
    return univariate / denom;
}

// Uncompress a domain-name.
// The returned pointer points to storage which is reused on each invocation.
static const char* uncompress(ns_msg msg, const byte* p, int* skiplen_p)
{
    static char ans[NS_MAXDNAME];
    int skiplen = ns_name_uncompress(ns_msg_base(msg), ns_msg_end(msg),
                                     p,
                                     ans, sizeof(ans));
    if (skiplen_p != NULL) *skiplen_p = skiplen;

    return (skiplen >= 0 ? ans : NULL);
}

typedef struct srv_result_s {
    char host[NS_MAXDNAME+1];
    unsigned int port;
    unsigned int priority;
    float order;
} srv_result;

static int srv_result_compar(const void* a_v, const void* b_v)
{
    srv_result* a = (srv_result*)a_v;
    srv_result* b = (srv_result*)b_v;

    if (a->priority == b->priority) {
        return a->order < b->order ? -1 : +1;
    } else {
        return a->priority - b->priority;
    }
}

/*
 * For a given domain, retrieve a sequence of SRV records,
 * sort them, and display the result.
 *
 * RFC 2782
 *
 * The SRV record is defined (textually) to be (priority, weight,
 * port, domain-name), where the first three are
 * two-byte integers.  The domain-name is not to be compressed,
 * according to the RFC, 'unless and until permitted by future
 * standards action', but uncompress doesn't seem to mind, and it
 * seems future-proof.
 *
 * See the RFC for other observations on how best to use SRV
 * records, as both a server and as a client.  The RFC says,
 * of the host selection algorithm:
 *
 *   Priority
 *        The priority of this target host.  A client MUST attempt to
 *        contact the target host with the lowest-numbered priority it can
 *        reach; target hosts with the same priority SHOULD be tried in an
 *        order defined by the weight field.  The range is 0-65535.  This
 *        is a 16 bit unsigned integer in network byte order.
 *
 *   Weight
 *        A server selection mechanism.  The weight field specifies a
 *        relative weight for entries with the same priority. Larger
 *        weights SHOULD be given a proportionately higher probability of
 *        being selected. The range of this number is 0-65535.  This is a
 *        16 bit unsigned integer in network byte order.  Domain
 *        administrators SHOULD use Weight 0 when there isn't any server
 *        selection to do, to make the RR easier to read for humans (less
 *        noisy).  In the presence of records containing weights greater
 *        than 0, records with weight 0 should have a very small chance of
 *        being selected.
 *
 *        In the absence of a protocol whose specification calls for the
 *        use of other weighting information, a client arranges the SRV
 *        RRs of the same Priority in the order in which target hosts,
 *        specified by the SRV RRs, will be contacted.
 *
 * The algorithm there is rather cumbersome.
 *
 * Instead, for each host, i, record the priority and an ordering
 * parameter which is drawn from an exponential distribution with
 * parameter weight_i (which we can obtain with
 * -log(1-uniform())/weight_i, where uniform() is a uniform random
 * deviate in [0,1)).  The minimum of an ensemble of such random
 * variables is item i, with probability weight_i/(sum_j(weight_j)).
 * See eg <https://en.wikipedia.org/wiki/Exponential_distribution#Distribution_of_the_minimum_of_exponential_random_variables>.
 *
 * With this done, sort the records.  The comparator function orders
 * host with lower priority first, and when two hosts have equal
 * priority, it orders first the one with the lower ordering parameter.
 */
char* get_sorted_srv_records(const char* domain)
{
    byte buf[NS_PACKETSZ];
    size_t retlen = 0; // length of return string
    int chatter = get_verbosity();

    char ldapdomain[NS_MAXDNAME+1];
    snprintf(ldapdomain, NS_MAXDNAME, "_ldap._tcp.%s", domain);
    ldapdomain[NS_MAXDNAME] = '\0'; // just in case

    errno = 0;
    int res_len = res_query(ldapdomain, ns_c_in, ns_t_srv, buf, sizeof(buf));
    if (res_len < 0) {
        if (errno != 0) {
            log_err_message(NULL, "res_query: error %s\n", strerror(errno));
        }
        return NULL;
    }

    ns_msg msg;
    ns_rr rr;

    ns_initparse(buf, res_len, &msg);
    int nmsg = ns_msg_count(msg, ns_s_an);

    srv_result* srv = (srv_result*)malloc(nmsg * sizeof(srv_result));
    if (srv == NULL) {
        log_err_message(NULL, "Unable to allocate %zu\n", nmsg * sizeof(srv_result));
	return NULL;
    }

    if (chatter > 2) fprintf(stderr, "SRV for %s:\n", domain);
    for (int recnum=0; recnum<nmsg; recnum++) {
        if (ns_parserr(&msg, ns_s_an, recnum, &rr)) {
            log_err_message(NULL, "ns_parserr: %s\n", strerror(errno));
	    return NULL;
        }

        const char* domainname = uncompress(msg, ns_rr_rdata(rr)+6, NULL);
        if (domainname) {
            const byte* prefs = ns_rr_rdata(rr);

            strncpy(srv[recnum].host, domainname, NS_MAXDNAME);
            srv[recnum].host[NS_MAXDNAME] = '\0'; // just in case

            // add 7 for "ldap://", 5 for portnum, 1 for space
            retlen += strlen(domainname) + 7 + 5 + 1;

            srv[recnum].priority = ns_get16(prefs);

            int16_t weight = ns_get16(prefs+2);
            if (weight == 0) {
                // This should only happen when there is no
                // weight-based selection to do, but insert a
                // placeholder and very high value, in case this has
                // been bungled.
                srv[recnum].order = 1e100;
            } else {
                // log(random_f()) would also work, but using
                // (1-random_f()) guarantees that we will never try to
                // evaluate log(0) (yes, I know that would be unlikely).
                //
                // Also, choose the log function appropriately, if the
                // argument is likely to be close to 1 (yes, this
                // is a degree of numerical care that is quite
                // unnecessary for this particular application).
                double r = random_f(); // in [0,1)
                double lograndom = (r < 0.5) ? log1p(-r) : log(1-r);
                srv[recnum].order = -lograndom/weight;
            }
            srv[recnum].port = ns_get16(prefs+4);

	    if (chatter > 2) {
            	fprintf(stderr, "%30s:%d\t%d %d -> %f\n",
	                srv[recnum].host, srv[recnum].port,
			srv[recnum].priority,
                        weight, srv[recnum].order);
	    }
        } else {
            log_err_message(NULL, "failed to uncompress domainname! %s\n", strerror(errno));
	    return NULL;
        }
    }

    qsort(srv, nmsg, sizeof(srv_result), srv_result_compar);

    char* rbuf = (char*)malloc(retlen);
    if (rbuf == NULL) {
	log_err_message(NULL, "Can't allocate %zd!\n", retlen);
        return NULL;
    }

    if (chatter > 2) {
	fprintf(stderr, "Order:\n");
	for (int recnum=0; recnum<nmsg; recnum++) {
	    fprintf(stderr, "  %d: %s:%d\n",
		    recnum, srv[recnum].host, srv[recnum].port);
	}
    }

    char* p = rbuf;
    for (int recnum=0; recnum<nmsg; recnum++) {
	size_t l = sprintf(p, "ldap://%s:%d ",
			   srv[recnum].host, srv[recnum].port);
	p += l;
    }
    p--;
    *p = '\0';

    return rbuf;
}

#if STANDALONE
// This is essentially a test program, but it might be useful in some other way.
static const char* progname;
static void Usage(void)
{
    fprintf(stderr, "Usage: %s [-v] domain\n", progname);
    exit(1);
}
int main(int argc, char** argv)
{
    progname = argv[0];

    res_init();

    const char* domain = NULL;

    for (argc--, argv++; argc>0; argc--, argv++) {
        if (**argv == '-') {
            while (*++*argv) {
                switch (**argv) {
                  case 'v':
                    change_verbosity(+1);
                    break;
                  default:
                    Usage();
                }
            }
        } else {
            if (domain) Usage();
            domain = *argv;
        }
    }
    if (domain == NULL) Usage();

    char* ldap_search_list = get_sorted_srv_records(domain);
    printf("%s\n", ldap_search_list);

    exit(0);
}
#endif