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 #include <chrono>
  17
  18 #include <algorithm>
  19
  20 #include <apt-pkg/configuration.h>
  21 #include <apt-pkg/error.h>
  22 #include <apt-pkg/strutl.h>
  23
  24
  25 #include "srvrec.h"
  26
  27
  28 bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result)
  29 {
  30    std::string target;
  31    struct servent *s_ent = getservbyport(htons(port), "tcp");
  32    if (s_ent == NULL)
  33       return false;
  34
  35    strprintf(target, "_%s._tcp.%s", s_ent->s_name, host.c_str());
  36    return GetSrvRecords(target, Result);
  37 }
  38
  39 bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
  40 {
  41    unsigned char answer[PACKETSZ];
  42    int answer_len, compressed_name_len;
  43    int answer_count;
  44
  45    if (res_init() != 0)
  46       return _error->Errno("res_init", "Failed to init resolver");
  47
  48    answer_len = res_query(name.c_str(), C_IN, T_SRV, answer, sizeof(answer));
  49    if (answer_len == -1)
  50       return false;
  51    if (answer_len < (int)sizeof(HEADER))
  52       return _error->Warning("Not enough data from res_query (%i)", answer_len);
  53
  54    // check the header
  55    HEADER *header = (HEADER*)answer;
  56    if (header->rcode != NOERROR)
  57       return _error->Warning("res_query returned rcode %i", header->rcode);
  58    answer_count = ntohs(header->ancount);
  59    if (answer_count <= 0)
  60       return _error->Warning("res_query returned no answers (%i) ", answer_count);
  61
  62    // skip the header
  63    compressed_name_len = dn_skipname(answer+sizeof(HEADER), answer+answer_len);
  64    if(compressed_name_len < 0)
  65       return _error->Warning("dn_skipname failed %i", compressed_name_len);
  66
  67    // pt points to the first answer record, go over all of them now
  68    unsigned char *pt = answer+sizeof(HEADER)+compressed_name_len+QFIXEDSZ;
  69    while ((int)Result.size() < answer_count && pt < answer+answer_len)
  70    {
  71       SrvRec rec;
  72       u_int16_t type, klass, priority, weight, port, dlen;
  73       char buf[MAXDNAME];
  74
  75       compressed_name_len = dn_skipname(pt, answer+answer_len);
  76       if (compressed_name_len < 0)
  77          return _error->Warning("dn_skipname failed (2): %i",
  78                                 compressed_name_len);
  79       pt += compressed_name_len;
  80       if (((answer+answer_len) - pt) < 16)
  81          return _error->Warning("packet too short");
  82
  83       // extract the data out of the result buffer
  84       #define extract_u16(target, p) target = *p++ << 8; target |= *p++;
  85
  86       extract_u16(type, pt);
  87       if(type != T_SRV)
  88          return _error->Warning("Unexpected type excepted %x != %x",
  89                                 T_SRV, type);
  90       extract_u16(klass, pt);
  91       if(klass != C_IN)
  92          return _error->Warning("Unexpected class excepted %x != %x",
  93                                 C_IN, klass);
  94       pt += 4;  // ttl
  95       extract_u16(dlen, pt);
  96       extract_u16(priority, pt);
  97       extract_u16(weight, pt);
  98       extract_u16(port, pt);
  99
 100       #undef extract_u16
 101
 102       compressed_name_len = dn_expand(answer, answer+answer_len, pt, buf, sizeof(buf));
 103       if(compressed_name_len < 0)
 104          return _error->Warning("dn_expand failed %i", compressed_name_len);
 105       pt += compressed_name_len;
 106
 107       // add it to our class
 108       rec.priority = priority;
 109       rec.weight = weight;
 110       rec.port = port;
 111       rec.target = buf;
 112       Result.push_back(rec);
 113    }
 114
 115    // implement load balancing as specified in RFC-2782
 116
 117    // sort them by priority
 118    std::stable_sort(Result.begin(), Result.end());
 119
 120    for(std::vector<SrvRec>::iterator I = Result.begin();
 121       I != Result.end(); ++I)
 122    {
 123       if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
 124       {
 125          std::cerr << "SrvRecs: got " << I->target
 126                    << " prio: " << I->priority
 127                    << " weight: " << I->weight
 128                    << std::endl;
 129       }
 130    }
 131
 132    return true;
 133 }
 134
 135 SrvRec PopFromSrvRecs(std::vector<SrvRec> &Recs)
 136 {
 137    // FIXME: instead of the simplistic shuffle below use the algorithm
 138    //        described in rfc2782 (with weights)
 139    //        and figure out how the weights need to be adjusted if
 140    //        a host refuses connections
 141
 142 #if 0  // all code below is only needed for the weight adjusted selection
 143    // assign random number ranges
 144    int prev_weight = 0;
 145    int prev_priority = 0;
 146    for(std::vector<SrvRec>::iterator I = Result.begin();
 147       I != Result.end(); ++I)
 148    {
 149       if(prev_priority != I->priority)
 150          prev_weight = 0;
 151       I->random_number_range_start = prev_weight;
 152       I->random_number_range_end = prev_weight + I->weight;
 153       prev_weight = I->random_number_range_end;
 154       prev_priority = I->priority;
 155
 156       if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
 157          std::cerr << "SrvRecs: got " << I->target
 158                    << " prio: " << I->priority
 159                    << " weight: " << I->weight
 160                    << std::endl;
 161    }
 162
 163    // go over the code in reverse order and note the max random range
 164    int max = 0;
 165    prev_priority = 0;
 166    for(std::vector<SrvRec>::iterator I = Result.end();
 167       I != Result.begin(); --I)
 168    {
 169       if(prev_priority != I->priority)
 170          max = I->random_number_range_end;
 171       I->random_number_range_max = max;
 172    }
 173 #endif
 174
 175    // shuffle in a very simplistic way for now (equal weights)
 176    std::vector<SrvRec>::iterator I, J;
 177    I = J = Recs.begin();
 178    for(;I != Recs.end(); ++I)
 179    {
 180       if(I->priority != J->priority)
 181          break;
 182    }
 183
 184    // FIXME: meeeeh, where to init this properly
 185    unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
 186    std::shuffle(J, I, std::default_random_engine(seed));
 187
 188    // meh, no pop_front() in std::vector?
 189    SrvRec selected = *Recs.begin();
 190    Recs.erase(Recs.begin());
 191
 192    if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
 193       std::cerr << "PopFromSrvRecs: selecting " << selected.target << std::endl;
 194
 195    return selected;
 196 }