// getsrvbyname.c -- A trivial implementation of a DNS SRV [RFC2782]
// resolver.

// Copyright (c) 2003, John S. Bucy 
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
//     * Redistributions of source code must retain the above
//     copyright notice, this list of conditions and the following
//     disclaimer.
//
//     * 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.
//
//     * The names of the author(s) may not 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 OWNER 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.



// NAME
// getsrvbyname
//
// SYNOPSIS
// int getsrvbyname(char *srvname, 
// 	            struct srvent *targets, 
// 	            int *targets_len) 
//
// Get DNS SRV records associated with srvname.  Return up to
// targets_len of them in targets.  targets_len is a value-result
// parameter which returns the number of records stored in targets.
//
// The srvent structure is defined as follows:
//
// struct srvent {
//   char host[NS_MAXDNAME];
//   uint16_t port; // local byte order
//   int weight;
//   int priority;
// };
//
// The members of the srvent structure are:
//
// host: the host providing the service.
// port: the port number on <host> for the service.
//
// 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.  [RFC2782]
//
// 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. [RFC2782]
//
// RETURN VALUE
// getsrvbyname returns 0 on success, -1 on error.  Additional error
// information may be available via hstrerror()/h_errno depending on
// the underlying resolver library.
//
// BUGS
// The resolver library is not be thread-safe or reentrant on some
// platforms.  getsrvbyname should be thread safe and reentrant
// provided the underlying resolver library is.


#include <inttypes.h>

#include <netdb.h>
#include <arpa/nameser.h>

#include <assert.h>
#include <stdio.h>

// This doesn't expect much of the underlying resolver; we just need
// res_search() and the various constants.

struct srvent {
  char host[NS_MAXDNAME];
  uint16_t port; // local byte order
  int weight;
  int priority;
};

#define MAXNAMES 32 // maximum number of names appearing in a message

// parsing state
struct pstate {
  char *buff;
  int bufflen;

  char *ptr; // current offset
  int resid; // remaining bytes to consume

  int names_count;

  // for header compression
  struct { char n[NS_MAXDNAME]; uint16_t offset; } names[MAXNAMES];
};

#define PSTATE_INIT {0,0,0,0,0}

struct dns_header {
  uint16_t id;
  unsigned qr : 1;
  unsigned opcode : 4;
  unsigned aa : 1;       // authoritative answer
  unsigned tc : 1;       // truncation
  unsigned rd : 1;       // recursion desired
  unsigned ra : 1;       // recursion available
  unsigned resv : 3;
  unsigned rcode : 4;
  uint16_t qdcount;  // query length
  uint16_t ancount;  // answer
  uint16_t nscount;  // nameserver
  uint16_t arcount;  // additional
};


static inline uint16_t
unmarshal_uint16(struct pstate *s)
{
  uint16_t result = (uint8_t)s->ptr[0] << 8 | (uint8_t)s->ptr[1];
  s->ptr += 2;
  s->resid -= 2;
  return result;
}

// returns length of name in ebuff on success, -1 on error.
// result must be MAX_DNAME size
int
unmarshal_name(struct pstate *s, char *result)
{
  int len = 0;
  uint16_t off0;
  int res_len = 0; // check against NS_MAXDNAME
  char *rptr = result;

  // "name compression"
  if(((uint8_t)*s->ptr & 0xc0) == 0xc0) {
    int i;
    off0 = ntohs(*(uint16_t *)s->ptr);
    off0 &= ~0xc000;

    if(s->resid < 2) return -1;

    for(i = 0; i < s->names_count; i++) {
      if(s->names[i].offset == off0) {
	strcpy(rptr, s->names[i].n);
	s->ptr += 2;
	s->resid -= 2;
	return 2;
      }
    }
    return -1;
  }


  if((uint8_t)*s->buff >= s->resid) {
    return -1;
  } 
  off0 = (uint16_t)(s->ptr - s->buff);

  do {
    int len2 = (uint8_t)*s->ptr;
    memcpy(rptr, s->ptr + 1, len2);
    res_len += len2;

    len += (len2 + 1);
    s->resid -= len2;
    rptr += len2;
    s->ptr += (len2 + 1);

    *rptr++ = '.';
    res_len++;

  } while(s->resid > 0 && *s->ptr);
  
  s->ptr++;
  s->resid--;
  len++; // the trailing 0
  *rptr = 0;

  assert(s->names_count < MAXNAMES);
  s->names[s->names_count].offset = off0;
  strcpy(s->names[s->names_count].n, result);
  s->names_count++;

  return 0;
}


// leaves s->ptr at the beininning of rdata (if any)
int 
unmarshal_rr(struct pstate *s,
	     char *name,
	     uint16_t *type,
	     uint16_t *class,
	     uint32_t *ttl)
{
  int rv = unmarshal_name(s, name);
  if(rv == -1) return -1;
  uint16_t rdata_len;

  if(s->resid < 8) return -1;

  rdata_len = unmarshal_uint16(s);
  
  if(s->resid < (rdata_len + 8)) return -1;

  // type  
  *type = unmarshal_uint16(s);

  // class 
  *class = unmarshal_uint16(s); 

  // ttl 
  *ttl = ntohl(*(uint32_t *)s->ptr);  s->ptr += 4; s->resid -= 4;

  return 0;
}

int
unmarshal_srv_rdata(struct pstate *s, 
		    struct srvent *result)
{
  //        _Service._Proto.Name SRV IN TTL 

  // Priority Weight Port Target

  if(s->resid < 6) return -1;
  result->priority = unmarshal_uint16(s);
  result->weight = unmarshal_uint16(s);
  result->port = unmarshal_uint16(s);
  return unmarshal_name(s, result->host);
}
		  


#define BUFFLEN 2048


int
getsrvbyname(char *dom, 
	     struct srvent *targets, 
	     int *targets_len) 
{
  int i;
  char buff[BUFFLEN];
  struct pstate st = PSTATE_INIT, *s = &st;
  int rv;
  struct dns_header *h;
  int targets_count = 0; // how many have we gotten?

  s->buff = s->ptr = buff;

  rv = res_search(dom, ns_c_in, ns_t_srv, buff, BUFFLEN);
  if(rv == -1 || rv < sizeof(struct dns_header)) 
  {
    return -1;
  }
  s->bufflen = s->resid = rv;

  // unmarshall the header 
  h = (struct dns_header *)buff;
  h->qdcount = ntohs(h->qdcount);
  h->ancount = ntohs(h->ancount);
  h->nscount = ntohs(h->nscount);
  h->arcount = ntohs(h->arcount);

  s->ptr += sizeof(struct dns_header);
  s->resid -= sizeof(struct dns_header);
    
  for(i = 0; i < h->qdcount; i++) {
    char name[NS_MAXDNAME];
    uint16_t type, class, rdata_len;
    uint32_t ttl;

    rv = unmarshal_name(s, name);
    if(rv == -1) return -1;

    if(s->resid < 4) return -1;
    
    // type  
    type = unmarshal_uint16(s);
    
    // class 
    class = unmarshal_uint16(s);
  }

  for(i = 0; i < h->ancount; i++) {
    char name[NS_MAXDNAME];
    uint16_t type, class, rdata_len;
    uint32_t ttl;

    char rdata[1024];

    rv = unmarshal_rr(s, name, &type, &class, &ttl);

    if(rv == -1) return -1;

    rv = unmarshal_srv_rdata(s, targets);
    if(rv == -1) return -1;
    
    targets++;
    targets_count++;
    if(targets_count > *targets_len) {
      goto out;
    }
  }

 out:
  *targets_len = targets_count;
  return 0;
}



int 
main(int argc, char **argv)
{
  struct srvent sent;
  int rv;
  int i, j;
  int res_len;

  argc--; argv++;

  for(i = 0; i < argc; i++) {
    res_len = 1;
    rv = getsrvbyname(argv[i], &sent, &res_len);
    if(rv == -1) {
      fprintf(stderr, "*** error: %s\n", hstrerror(h_errno));
    }
    else {
      for(j = 0; j < res_len; j++) {
	printf("%s %d %d %d\n", sent.host, sent.port, sent.priority, sent.weight);
      }
    }
  }
  

  exit(0);
  return 0;
}