From: bert hubert <bert.hubert@netherlabs.nl>
Date: Mon, 31 Oct 2016 23:06:43 +0000 (+0100)
Subject: interim commit that goes with https://gist.github.com/ahupowerdns/1e8bfbba95a277a4fac... 
X-Git-Tag: dnsdist-1.3.1~136^2~61
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=5dbd408caafca2cd9e455757045954b5abcf02c9;p=thirdparty%2Fpdns.git

interim commit that goes with https://gist.github.com/ahupowerdns/1e8bfbba95a277a4fac09cb3654eb2ac
---

diff --git a/pdns/dnsrecords.cc b/pdns/dnsrecords.cc
index fbd05a8356..2feebc990a 100644
--- a/pdns/dnsrecords.cc
+++ b/pdns/dnsrecords.cc
@@ -141,7 +141,7 @@ boilerplate_conv(DNAME, QType::DNAME, conv.xfrName(d_content));
 boilerplate_conv(MR, QType::MR, conv.xfrName(d_alias, true));
 boilerplate_conv(MINFO, QType::MINFO, conv.xfrName(d_rmailbx, true); conv.xfrName(d_emailbx, true));
 boilerplate_conv(TXT, QType::TXT, conv.xfrText(d_text, true));
-boilerplate_conv(LUA, QType::LUA, conv.xfrText(d_code, true));
+boilerplate_conv(LUA, QType::LUA, conv.xfrType(d_type); conv.xfrText(d_code, true));
 boilerplate_conv(ENT, 0, );
 boilerplate_conv(SPF, 99, conv.xfrText(d_text, true));
 boilerplate_conv(HINFO, QType::HINFO,  conv.xfrText(d_cpu);   conv.xfrText(d_host));
diff --git a/pdns/dnsrecords.hh b/pdns/dnsrecords.hh
index fbaff717da..8e9ec7b016 100644
--- a/pdns/dnsrecords.hh
+++ b/pdns/dnsrecords.hh
@@ -190,6 +190,7 @@ class LUARecordContent : public DNSRecordContent
 public:
   includeboilerplate(LUA)
 
+  uint16_t d_type;
   string d_code;
 };
 
diff --git a/pdns/lua-auth4.cc b/pdns/lua-auth4.cc
index 56d05389bc..e7912080c5 100644
--- a/pdns/lua-auth4.cc
+++ b/pdns/lua-auth4.cc
@@ -6,6 +6,8 @@
 #include "namespaces.hh"
 #include "ednssubnet.hh"
 #include <unordered_set>
+#include "sstuff.hh"
+#include <thread>
 
 #if !defined(HAVE_LUA)
 
@@ -283,8 +285,50 @@ bool AuthLua4::updatePolicy(const DNSName &qname, QType qtype, const DNSName &zo
 AuthLua4::~AuthLua4() { }
 
 
+class IsUpOracle
+{
+public:
+  bool isUp(const ComboAddress& remote);
+  
+  
+private:
+  void checkThread(const ComboAddress& rem) {
+    d_statuses[rem].second=false;
+    for(;;) {
+      try {
+        Socket s(rem.sin4.sin_family, SOCK_STREAM);
+        s.setNonBlocking();
+        s.connect(rem, 1);
+        if(!d_statuses[rem].second)
+          cout<<"Declaring "<<rem.toStringWithPort()<<" UP!"<<endl;
+        d_statuses[rem].second=true;
+      }
+      catch(NetworkError& ne) {
+        if(d_statuses[rem].second)
+          cout<<"Failed to connect to "<<rem.toStringWithPort()<<", setting DOWN"<<endl;
+        d_statuses[rem].second=false;
+      }
+      sleep(1);
+    }
+  }
+  map<ComboAddress, pair<std::thread*, bool>> d_statuses;
+};
+
+bool IsUpOracle::isUp(const ComboAddress& remote)
+{
+  auto iter = d_statuses.find(remote);
+  if(iter == d_statuses.end()) {
+    cout<<"First ever query for "<<remote.toStringWithPort()<<", launching checker"<<endl;
+    std::thread* checker = new std::thread(&IsUpOracle::checkThread, this, remote);
+    d_statuses[remote]={checker, false};
+    return false;
+  }
+  return iter->second.second;
+}
+
+IsUpOracle g_up;
 
-std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, uint16_t qtype) 
+std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, const DNSName& query, const ComboAddress& who, uint16_t qtype) 
 {
   std::vector<shared_ptr<DNSRecordContent>> ret;
   cout<<"Code: "<<code<<endl;
@@ -292,7 +336,52 @@ std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, uint
   strip.resize(strip.size()-1);
   
   LuaContext lua;
-  ret.push_back(std::shared_ptr<DNSRecordContent>(DNSRecordContent::mastermake(qtype, 1, lua.executeCode<string>("return "+strip))));
+  lua.writeVariable("qname", query.toString());
+  lua.writeVariable("who", who.toString());
+
+  lua.writeFunction("ifportup", [](int port, const vector<pair<int, string> >& ips) {
+      vector<ComboAddress> candidates;
+      for(const auto& i : ips) {
+        ComboAddress rem(i.second, port);
+        if(g_up.isUp(rem))
+          candidates.push_back(rem);
+      }
+      cout<<"Have "<<candidates.size()<<" candidate IP addresses: ";
+      for(const auto& c : candidates)
+        cout<<c.toString()<<" ";
+      cout<<endl;
+      if(candidates.empty()) {
+        cout<<"Picking a random one, everything is down. Should return all of them"<<endl;
+        return ips[random() % ips.size()].second;
+      }
+      return candidates[random() % candidates.size()].toString();
+    });
+  
+  lua.writeFunction("pickRandom", [](const vector<pair<int, string> >& ips) {
+      return ips[random()%ips.size()].second;
+    });
+
+  // wrandom({ {100, '1.2.3.4'}, {50, '5.4.3.2'}, {1, '192.168.1.0'}})"
+
+  lua.writeFunction("wrandom", [](std::unordered_map<int, std::unordered_map<int, string> > ips) {
+      int sum=0;
+      vector<pair<int, string> > pick;
+      for(auto& i : ips) {
+        sum += atoi(i.second[1].c_str());
+        pick.push_back({sum, i.second[2]});
+      }
+      int r = random() % sum;
+      auto p = upper_bound(pick.begin(), pick.end(),r, [](int r, const decltype(pick)::value_type& a) { return  r < a.first;});
+      return p->second;
+      
+    });
+
+  
+  string content=lua.executeCode<string>(strip);
+  if(qtype==QType::TXT)
+    content = '"'+content+'"';
+  
+  ret.push_back(std::shared_ptr<DNSRecordContent>(DNSRecordContent::mastermake(qtype, 1, content )));
 
   return ret;
 }
diff --git a/pdns/lua-auth4.hh b/pdns/lua-auth4.hh
index fcdd545f80..e7aec2fb13 100644
--- a/pdns/lua-auth4.hh
+++ b/pdns/lua-auth4.hh
@@ -42,4 +42,4 @@ private:
   luacall_update_policy_t d_update_policy;
   luacall_axfr_filter_t d_axfr_filter;
 };
-std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, uint16_t qtype);
+std::vector<shared_ptr<DNSRecordContent>> luaSynth(const std::string& code, const DNSName& qname, const ComboAddress& who, uint16_t qtype);
diff --git a/pdns/packethandler.cc b/pdns/packethandler.cc
index e50247d257..0ba3e64a2d 100644
--- a/pdns/packethandler.cc
+++ b/pdns/packethandler.cc
@@ -1284,10 +1284,12 @@ DNSPacket *PacketHandler::doQuestion(DNSPacket *p)
     
     while(B.get(rr)) {
       if(rr.dr.d_type == QType::LUA) {
-        cerr<<"We got LUA: "<<endl;
-        auto recvec=luaSynth(getRR<LUARecordContent>(rr.dr)->d_code, p->qtype.getCode());
+        auto rec=getRR<LUARecordContent>(rr.dr);
+        if(rec->d_type != p->qtype.getCode())
+          continue;
+        
+        auto recvec=luaSynth(rec->d_code, target, p->getRemote(), p->qtype.getCode());
         if(!recvec.empty()) {
-          cerr<<"We got answers from Lua"<<endl;
           for(const auto& r : recvec) {
             rr.dr.d_type = p->qtype.getCode();
             rr.dr.d_content = r;