apt-pkg/contrib/srvrec.cc

   1 // -*- mode: cpp; mode: fold -*-
   2 // Description                                                          /*{{{*/
   3 /* ######################################################################
   4
   5    SRV record support
   6
   7    ##################################################################### */
   8                                                                         /*}}}*/
   9 #include <config.h>
  10
  11 #include <netdb.h>
  12
  13 #include <netinet/in.h>
  14 #include <arpa/nameser.h>
  15 #include <resolv.h>
  16
  17 #include <algorithm>
  18
  19 #include <apt-pkg/strutl.h>
  20 #include <apt-pkg/error.h>
  21 #include "srvrec.h"
  22
  23 bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result)
  24 {
  25    std::string target;
  26    struct servent *s_ent = getservbyport(htons(port), "tcp");
  27    if (s_ent == NULL)
  28       return false;
  29
  30    strprintf(target, "_%s._tcp.%s", s_ent->s_name, host.c_str());
  31    return GetSrvRecords(target, Result);
  32 }
  33
  34 bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
  35 {
  36    unsigned char answer[PACKETSZ];
  37    int answer_len, compressed_name_len;
  38    int answer_count;
  39
  40    if (res_init() != 0)
  41       return _error->Errno("res_init", "Failed to init resolver");
  42
  43    answer_len = res_query(name.c_str(), C_IN, T_SRV, answer, sizeof(answer));
  44    if (answer_len == -1)
  45       return false;
  46    if (answer_len < (int)sizeof(HEADER))
  47       return _error->Warning("Not enough data from res_query (%i)", answer_len);
  48
  49    // check the header
  50    HEADER *header = (HEADER*)answer;
  51    if (header->rcode != NOERROR)
  52       return _error->Warning("res_query returned rcode %i", header->rcode);
  53    answer_count = ntohs(header->ancount);
  54    if (answer_count <= 0)
  55       return _error->Warning("res_query returned no answers (%i) ", answer_count);
  56
  57    // skip the header
  58    compressed_name_len = dn_skipname(answer+sizeof(HEADER), answer+answer_len);
  59    if(compressed_name_len < 0)
  60       return _error->Warning("dn_skipname failed %i", compressed_name_len);
  61
  62    // pt points to the first answer record, go over all of them now
  63    unsigned char *pt = answer+sizeof(HEADER)+compressed_name_len+QFIXEDSZ;
  64    while ((int)Result.size() < answer_count && pt < answer+answer_len)
  65    {
  66       SrvRec rec;
  67       u_int16_t type, klass, priority, weight, port, dlen;
  68       char buf[MAXDNAME];
  69
  70       compressed_name_len = dn_skipname(pt, answer+answer_len);
  71       if (compressed_name_len < 0)
  72          return _error->Warning("dn_skipname failed (2): %i",
  73                                 compressed_name_len);
  74       pt += compressed_name_len;
  75       if (((answer+answer_len) - pt) < 16)
  76          return _error->Warning("packet too short");
  77
  78       // extract the data out of the result buffer
  79       #define extract_u16(target, p) target = *p++ << 8; target |= *p++;
  80
  81       extract_u16(type, pt);
  82       if(type != T_SRV)
  83          return _error->Warning("Unexpected type excepted %x != %x",
  84                                 T_SRV, type);
  85       extract_u16(klass, pt);
  86       if(klass != C_IN)
  87          return _error->Warning("Unexpected class excepted %x != %x",
  88                                 C_IN, klass);
  89       pt += 4;  // ttl
  90       extract_u16(dlen, pt);
  91       extract_u16(priority, pt);
  92       extract_u16(weight, pt);
  93       extract_u16(port, pt);
  94
  95       #undef extract_u16
  96
  97       compressed_name_len = dn_expand(answer, answer+answer_len, pt, buf, sizeof(buf));
  98       if(compressed_name_len < 0)
  99          return _error->Warning("dn_expand failed %i", compressed_name_len);
 100       pt += compressed_name_len;
 101
 102       // add it to our class
 103       rec.priority = priority;
 104       rec.weight = weight;
 105       rec.port = port;
 106       rec.target = buf;
 107       Result.push_back(rec);
 108    }
 109
 110    // implement load balancing as specified in RFC-2782
 111
 112    // sort them by priority
 113    std::stable_sort(Result.begin(), Result.end());
 114
 115    // assign random number ranges
 116    int prev_weight = 0;
 117    int prev_priority = 0;
 118    for(std::vector<SrvRec>::iterator I = Result.begin();
 119       I != Result.end(); ++I)
 120    {
 121       if(prev_priority != I->priority)
 122          prev_weight = 0;
 123       I->random_number_range_start = prev_weight;
 124       I->random_number_range_end = prev_weight + I->weight;
 125       prev_weight = I->random_number_range_end;
 126       prev_priority = I->priority;
 127    }
 128
 129    // go over the code in reverse order and note the max random range
 130    int max = 0;
 131    prev_priority = 0;
 132    for(std::vector<SrvRec>::iterator I = Result.end();
 133       I != Result.begin(); --I)
 134    {
 135       if(prev_priority != I->priority)
 136          max = I->random_number_range_end;
 137       I->random_number_range_max = max;
 138    }
 139
 140    // FIXME: now shuffle
 141
 142    return true;
 143 }