Merge remote-tracking branch 'mvo/feature/srv-records' into debian/experimental

8 files changed, 353 insertions, 14 deletions

diff --git a/apt-pkg/contrib/srvrec.cc b/apt-pkg/contrib/srvrec.cc
new file mode 100644
index 000000000..b4a3d97d2
--- /dev/null
+++ b/apt-pkg/contrib/srvrec.cc
@@ -0,0 +1,196 @@
+// -*- mode: cpp; mode: fold -*-
+// Description								/*{{{*/
+/* ######################################################################
+
+   SRV record support
+
+   ##################################################################### */
+									/*}}}*/
+#include <config.h>
+
+#include <netdb.h>
+
+#include <netinet/in.h>
+#include <arpa/nameser.h>
+#include <resolv.h>
+#include <chrono>
+
+#include <algorithm>
+
+#include <apt-pkg/configuration.h>
+#include <apt-pkg/error.h>
+#include <apt-pkg/strutl.h>
+
+
+#include "srvrec.h"
+
+
+bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result)
+{
+   std::string target;
+   struct servent *s_ent = getservbyport(htons(port), "tcp");
+   if (s_ent == NULL)
+      return false;
+
+   strprintf(target, "_%s._tcp.%s", s_ent->s_name, host.c_str());
+   return GetSrvRecords(target, Result);
+}
+
+bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
+{
+   unsigned char answer[PACKETSZ];
+   int answer_len, compressed_name_len;
+   int answer_count;
+
+   if (res_init() != 0)
+      return _error->Errno("res_init", "Failed to init resolver");
+
+   answer_len = res_query(name.c_str(), C_IN, T_SRV, answer, sizeof(answer));
+   if (answer_len == -1)
+      return false;
+   if (answer_len < (int)sizeof(HEADER))
+      return _error->Warning("Not enough data from res_query (%i)", answer_len);
+
+   // check the header
+   HEADER *header = (HEADER*)answer;
+   if (header->rcode != NOERROR)
+      return _error->Warning("res_query returned rcode %i", header->rcode);
+   answer_count = ntohs(header->ancount);
+   if (answer_count <= 0)
+      return _error->Warning("res_query returned no answers (%i) ", answer_count);
+
+   // skip the header
+   compressed_name_len = dn_skipname(answer+sizeof(HEADER), answer+answer_len);
+   if(compressed_name_len < 0)
+      return _error->Warning("dn_skipname failed %i", compressed_name_len);
+
+   // pt points to the first answer record, go over all of them now
+   unsigned char *pt = answer+sizeof(HEADER)+compressed_name_len+QFIXEDSZ;
+   while ((int)Result.size() < answer_count && pt < answer+answer_len)
+   {
+      SrvRec rec;
+      u_int16_t type, klass, priority, weight, port, dlen;
+      char buf[MAXDNAME];
+
+      compressed_name_len = dn_skipname(pt, answer+answer_len);
+      if (compressed_name_len < 0)
+         return _error->Warning("dn_skipname failed (2): %i",
+                                compressed_name_len);
+      pt += compressed_name_len;
+      if (((answer+answer_len) - pt) < 16)
+         return _error->Warning("packet too short");
+
+      // extract the data out of the result buffer
+      #define extract_u16(target, p) target = *p++ << 8; target |= *p++;
+
+      extract_u16(type, pt);
+      if(type != T_SRV)
+         return _error->Warning("Unexpected type excepted %x != %x",
+                                T_SRV, type);
+      extract_u16(klass, pt);
+      if(klass != C_IN)
+         return _error->Warning("Unexpected class excepted %x != %x",
+                                C_IN, klass);
+      pt += 4;  // ttl
+      extract_u16(dlen, pt);
+      extract_u16(priority, pt);
+      extract_u16(weight, pt);
+      extract_u16(port, pt);
+
+      #undef extract_u16
+
+      compressed_name_len = dn_expand(answer, answer+answer_len, pt, buf, sizeof(buf));
+      if(compressed_name_len < 0)
+         return _error->Warning("dn_expand failed %i", compressed_name_len);
+      pt += compressed_name_len;
+
+      // add it to our class
+      rec.priority = priority;
+      rec.weight = weight;
+      rec.port = port;
+      rec.target = buf;
+      Result.push_back(rec);
+   }
+
+   // implement load balancing as specified in RFC-2782
+
+   // sort them by priority
+   std::stable_sort(Result.begin(), Result.end());
+
+   for(std::vector<SrvRec>::iterator I = Result.begin();
+      I != Result.end(); ++I)
+   {
+      if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+      {
+         std::cerr << "SrvRecs: got " << I->target
+                   << " prio: " << I->priority
+                   << " weight: " << I->weight
+                   << std::endl;
+      }
+   }
+
+   return true;
+}
+
+SrvRec PopFromSrvRecs(std::vector<SrvRec> &Recs)
+{
+   // FIXME: instead of the simplistic shuffle below use the algorithm
+   //        described in rfc2782 (with weights)
+   //        and figure out how the weights need to be adjusted if
+   //        a host refuses connections
+
+#if 0  // all code below is only needed for the weight adjusted selection 
+   // assign random number ranges
+   int prev_weight = 0;
+   int prev_priority = 0;
+   for(std::vector<SrvRec>::iterator I = Result.begin();
+      I != Result.end(); ++I)
+   {
+      if(prev_priority != I->priority)
+         prev_weight = 0;
+      I->random_number_range_start = prev_weight;
+      I->random_number_range_end = prev_weight + I->weight;
+      prev_weight = I->random_number_range_end;
+      prev_priority = I->priority;
+
+      if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+         std::cerr << "SrvRecs: got " << I->target
+                   << " prio: " << I->priority
+                   << " weight: " << I->weight
+                   << std::endl;
+   }
+
+   // go over the code in reverse order and note the max random range
+   int max = 0;
+   prev_priority = 0;
+   for(std::vector<SrvRec>::iterator I = Result.end();
+      I != Result.begin(); --I)
+   {
+      if(prev_priority != I->priority)
+         max = I->random_number_range_end;
+      I->random_number_range_max = max;
+   }
+#endif
+
+   // shuffle in a very simplistic way for now (equal weights)
+   std::vector<SrvRec>::iterator I, J;
+   I = J = Recs.begin();
+   for(;I != Recs.end(); ++I)
+   {
+      if(I->priority != J->priority)
+         break;
+   }
+
+   // FIXME: meeeeh, where to init this properly
+   unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
+   std::shuffle(J, I, std::default_random_engine(seed));
+
+   // meh, no pop_front() in std::vector?
+   SrvRec selected = *Recs.begin();
+   Recs.erase(Recs.begin());
+
+   if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+      std::cerr << "PopFromSrvRecs: selecting " << selected.target << std::endl;
+
+   return selected;
+}
diff --git a/apt-pkg/contrib/srvrec.h b/apt-pkg/contrib/srvrec.h
new file mode 100644
index 000000000..e07edc683
--- /dev/null
+++ b/apt-pkg/contrib/srvrec.h
@@ -0,0 +1,47 @@
+// -*- mode: cpp; mode: fold -*-
+// Description								/*{{{*/
+/* ######################################################################
+
+   SRV record support
+
+   ##################################################################### */
+									/*}}}*/
+#ifndef SRVREC_H
+#define SRVREC_H
+
+#include <arpa/nameser.h>
+#include <vector>
+#include <string>
+
+class SrvRec
+{
+ public:
+   std::string target;
+   u_int16_t priority;
+   u_int16_t weight;
+   u_int16_t port;
+
+   // each server is assigned a interval [start, end] in the space of [0, max]
+   int random_number_range_start;
+   int random_number_range_end;
+   int random_number_range_max;
+
+   bool operator<(SrvRec const &other) const { 
+      return this->priority < other.priority; 
+   }
+};
+
+/** \brief Get SRV records from host/port (builds the query string internally) 
+ */
+bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result);
+
+/** \brief Get SRV records for query string like: _http._tcp.example.com
+ */
+bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result);
+
+/** \brief Pop a single SRV record from the vector of SrvRec taking
+ *         priority and weight into account
+ */
+SrvRec PopFromSrvRecs(std::vector<SrvRec> &Recs);
+
+#endif
diff --git a/apt-pkg/makefile b/apt-pkg/makefile
index 5603b51ed..45dcddfe3 100644
--- a/apt-pkg/makefile
+++ b/apt-pkg/makefile
@@ -15,7 +15,7 @@ include ../buildlib/libversion.mak
 LIBRARY=apt-pkg
 MAJOR=$(LIBAPTPKG_MAJOR)
 MINOR=$(LIBAPTPKG_RELEASE)
-SLIBS=$(PTHREADLIB) $(INTLLIBS) -lutil -ldl
+SLIBS=$(PTHREADLIB) $(INTLLIBS) -lutil -ldl -lresolv
 ifeq ($(HAVE_ZLIB),yes)
 SLIBS+= -lz
 endif
diff --git a/cmdline/apt-helper.cc b/cmdline/apt-helper.cc
index a97fc903f..d235ac315 100644
--- a/cmdline/apt-helper.cc
+++ b/cmdline/apt-helper.cc
@@ -22,6 +22,7 @@
 #include <apt-private/private-output.h>
 #include <apt-private/private-download.h>
 #include <apt-private/private-cmndline.h>
+#include <apt-pkg/srvrec.h>
 
 #include <iostream>
 #include <string>
@@ -81,6 +82,31 @@ static bool DoDownloadFile(CommandLine &CmdL)
    return true;
 }
 
+static bool DoSrvLookup(CommandLine &CmdL)
+{
+   if (CmdL.FileSize() < 1)
+      return _error->Error(_("Must specifc at least one srv record"));
+
+   std::vector<SrvRec> srv_records;
+   c1out << "# target priority weight port" << std::endl;
+   for(int i=1; CmdL.FileList[i] != NULL; i++)
+   {
+      if(GetSrvRecords(CmdL.FileList[i], srv_records) == false)
+         _error->Warning(_("GetSrvRec failed for %s"), CmdL.FileList[i]);
+      for (std::vector<SrvRec>::const_iterator I = srv_records.begin();
+           I != srv_records.end(); ++I)
+      {
+         c1out << (*I).target.c_str() << " "
+               << (*I).priority << " "
+               << (*I).weight << " "
+               << (*I).port << " "
+               << std::endl;
+      }
+   }
+
+   return true;
+}
+
 static bool ShowHelp(CommandLine &)
 {
    ioprintf(std::cout, "%s %s (%s)\n", PACKAGE, PACKAGE_VERSION, COMMON_ARCH);
@@ -96,6 +122,7 @@ static bool ShowHelp(CommandLine &)
       "\n"
       "Commands:\n"
       "   download-file - download the given uri to the target-path\n"
+      "   srv-lookup - lookup a SRV record (e.g. _http._tcp.ftp.debian.org)\n"
       "   auto-detect-proxy - detect proxy using apt.conf\n"
       "\n"
       "                       This APT helper has Super Meep Powers.\n");
@@ -107,6 +134,7 @@ int main(int argc,const char *argv[])					/*{{{*/
 {
    CommandLine::Dispatch Cmds[] = {{"help",&ShowHelp},
 				   {"download-file", &DoDownloadFile},
+				   {"srv-lookup", &DoSrvLookup},
 				   {"auto-detect-proxy", &DoAutoDetectProxy},
                                    {0,0}};
 
diff --git a/cmdline/makefile b/cmdline/makefile
index 816038c3b..6d21b0803 100644
--- a/cmdline/makefile
+++ b/cmdline/makefile
@@ -49,7 +49,7 @@ include $(PROGRAM_H)
 
 # The apt-helper
 PROGRAM=apt-helper
-SLIBS = -lapt-pkg -lapt-private $(INTLLIBS)
+SLIBS = -lapt-pkg -lapt-private $(INTLLIBS) -lresolv
 LIB_MAKES = apt-pkg/makefile apt-private/makefile
 SOURCE = apt-helper.cc
 include $(PROGRAM_H)
diff --git a/methods/connect.cc b/methods/connect.cc
index e2cbf4f5c..5612af6ec 100644
--- a/methods/connect.cc
+++ b/methods/connect.cc
@@ -18,6 +18,7 @@
 #include <apt-pkg/strutl.h>
 #include <apt-pkg/acquire-method.h>
 #include <apt-pkg/configuration.h>
+#include <apt-pkg/srvrec.h>
 
 #include <stdio.h>
 #include <errno.h>
@@ -43,6 +44,9 @@ static int LastPort = 0;
 static struct addrinfo *LastHostAddr = 0;
 static struct addrinfo *LastUsed = 0;
 
+static std::vector<SrvRec> SrvRecords;
+static int LastSrvRecord = 0;
+
 // Set of IP/hostnames that we timed out before or couldn't resolve
 static std::set<std::string> bad_addr;
 
@@ -130,15 +134,12 @@ static bool DoConnect(struct addrinfo *Addr,std::string Host,
    return true;
 }
 									/*}}}*/
-// Connect - Connect to a server					/*{{{*/
-// ---------------------------------------------------------------------
-/* Performs a connection to the server */
-bool Connect(std::string Host,int Port,const char *Service,int DefPort,int &Fd,
-	     unsigned long TimeOut,pkgAcqMethod *Owner)
-{
-   if (_error->PendingError() == true)
-      return false;
 
+// Connect to a given Hostname 
+bool ConnectToHostname(std::string Host,int Port,const char *Service,
+                            int DefPort,int &Fd,
+                            unsigned long TimeOut,pkgAcqMethod *Owner)
+{
    // Convert the port name/number
    char ServStr[300];
    if (Port != 0)
@@ -258,3 +259,37 @@ bool Connect(std::string Host,int Port,const char *Service,int DefPort,int &Fd,
    return _error->Error(_("Unable to connect to %s:%s:"),Host.c_str(),ServStr);
 }
 									/*}}}*/
+// Connect - Connect to a server					/*{{{*/
+// ---------------------------------------------------------------------
+/* Performs a connection to the server (including SRV record lookup) */
+bool Connect(std::string Host,int Port,const char *Service,
+                            int DefPort,int &Fd,
+                            unsigned long TimeOut,pkgAcqMethod *Owner)
+{
+   if (_error->PendingError() == true)
+      return false;
+
+   if(LastHost != Host || LastPort != Port)
+   {
+      SrvRecords.clear();
+      if (_config->FindB("Acquire::EnableSrvRecods", true) == true)
+         GetSrvRecords(Host, DefPort, SrvRecords);
+   }
+   // we have no SrvRecords for this host, connect right away
+   if(SrvRecords.size() == 0)
+      return ConnectToHostname(Host, Port, Service, DefPort, Fd, 
+                                    TimeOut, Owner);
+
+   // try to connect in the priority order of the srv records
+   while(SrvRecords.size() > 0)
+   {
+      Host = PopFromSrvRecs(SrvRecords).target;
+      if(ConnectToHostname(Host, Port, Service, DefPort, Fd, TimeOut, Owner))
+         return true;
+
+      // we couldn't connect to this one, use the next
+      SrvRecords.erase(SrvRecords.begin());
+   }
+
+   return false;
+}
diff --git a/methods/makefile b/methods/makefile
index 6b7781294..868c52a40 100644
--- a/methods/makefile
+++ b/methods/makefile
@@ -46,21 +46,21 @@ include $(PROGRAM_H)
 
 # The http method
 PROGRAM=http
-SLIBS = -lapt-pkg $(SOCKETLIBS) $(INTLLIBS)
+SLIBS = -lapt-pkg $(SOCKETLIBS) $(INTLLIBS) -lresolv
 LIB_MAKES = apt-pkg/makefile
 SOURCE = http.cc http_main.cc rfc2553emu.cc connect.cc server.cc
 include $(PROGRAM_H)
 
 # The https method
 PROGRAM=https
-SLIBS = -lapt-pkg -lcurl $(INTLLIBS)
+SLIBS = -lapt-pkg -lcurl $(INTLLIBS) -lresolv
 LIB_MAKES = apt-pkg/makefile
 SOURCE = https.cc server.cc
 include $(PROGRAM_H)
 
 # The ftp method
 PROGRAM=ftp
-SLIBS = -lapt-pkg $(SOCKETLIBS) $(INTLLIBS)
+SLIBS = -lapt-pkg $(SOCKETLIBS) $(INTLLIBS) -lresolv
 LIB_MAKES = apt-pkg/makefile
 SOURCE = ftp.cc rfc2553emu.cc connect.cc
 include $(PROGRAM_H)
@@ -81,7 +81,7 @@ include $(PROGRAM_H)
 
 # The mirror method
 PROGRAM=mirror
-SLIBS = -lapt-pkg $(SOCKETLIBS)
+SLIBS = -lapt-pkg $(SOCKETLIBS) -lresolv
 LIB_MAKES = apt-pkg/makefile
 SOURCE = mirror.cc http.cc rfc2553emu.cc connect.cc server.cc
 include $(PROGRAM_H)
diff --git a/test/libapt/srvrecs_test.cc b/test/libapt/srvrecs_test.cc
new file mode 100644
index 000000000..7e43cc757
--- /dev/null
+++ b/test/libapt/srvrecs_test.cc
@@ -0,0 +1,33 @@
+#include <config.h>
+
+#include <apt-pkg/srvrec.h>
+
+#include <string>
+#include <iostream>
+
+#include <gtest/gtest.h>
+
+TEST(SrvRecTest, PopFromSrvRecs)
+{
+   // the PopFromSrvRecs() is using a random number so we must
+   // run it a bunch of times to ensure we are not fooled by randomness
+   std::set<std::string> selected;
+   for(int i=0;i<100;i++)
+   {
+      std::vector<SrvRec> Meep;
+      SrvRec foo = {target:"foo", priority: 20, weight: 0, port: 80};
+      Meep.push_back(foo);
+      
+      SrvRec bar = {target:"bar", priority: 20, weight: 0, port: 80};
+      Meep.push_back(bar);
+
+      EXPECT_EQ(Meep.size(), 2);
+      SrvRec result = PopFromSrvRecs(Meep);
+      selected.insert(result.target);
+      // ensure that pop removed one element
+      EXPECT_EQ(Meep.size(), 1);
+   }
+
+   // ensure that after enough runs we end up with both selected
+   EXPECT_EQ(selected.size(), 2);
+}