From 839f302153d16bca051bba35f4df2feb9e39e80e Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Thu, 19 Mar 2015 13:53:04 +0100
Subject: [PATCH] isolate boost options parser from dnsdist-lua, allow servers
 to be part of multiple pools simultaneously

---
 pdns/dnsdist-lua.cc  | 30 +++++++++++++++++++-----------
 pdns/dnsdist.cc      |  7 +++++--
 pdns/dnsdist.hh      |  7 +------
 pdns/dnsdistconf.lua |  5 ++---
 4 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index fa6d926ac..6982f4458 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -10,11 +10,12 @@ using std::thread;
 
 static vector<std::function<void(void)>>* g_launchWork;
 
-vector<std::function<void(void)>> setupLua(bool client)
+vector<std::function<void(void)>> setupLua(bool client, const std::string& config)
 {
   g_launchWork= new vector<std::function<void(void)>>();
+  typedef std::unordered_map<std::string, boost::variant<std::string, vector<pair<int, std::string> > > > newserver_t;
   g_lua.writeFunction("newServer", 
-		      [client](boost::variant<string,std::unordered_map<std::string, std::string>> pvars, boost::optional<int> qps)
+		      [client](boost::variant<string,newserver_t> pvars, boost::optional<int> qps)
 		      { 
 			if(client) {
 			  return std::make_shared<DownstreamState>(ComboAddress());
@@ -44,23 +45,30 @@ vector<std::function<void(void)>> setupLua(bool client)
 
 			  return ret;
 			}
-			auto vars=boost::get<std::unordered_map<std::string, std::string>>(pvars);
-			auto ret=std::make_shared<DownstreamState>(ComboAddress(vars["address"], 53));
+			auto vars=boost::get<newserver_t>(pvars);
+			auto ret=std::make_shared<DownstreamState>(ComboAddress(boost::get<string>(vars["address"]), 53));
 			
 			if(vars.count("qps")) {
-			  ret->qps=QPSLimiter(boost::lexical_cast<int>(vars["qps"]),boost::lexical_cast<int>(vars["qps"]));
+			  int qps=boost::lexical_cast<int>(boost::get<string>(vars["qps"]));
+			  ret->qps=QPSLimiter(qps, qps);
 			}
 
 			if(vars.count("pool")) {
-			  ret->pools.insert(vars["pool"]);
+			  if(auto* pool = boost::get<string>(&vars["pool"]))
+			    ret->pools.insert(*pool);
+			  else {
+			    auto* pools = boost::get<vector<pair<int, string> > >(&vars["pool"]);
+			    for(auto& p : *pools)
+			      ret->pools.insert(p.second);
+			  }
 			}
 
 			if(vars.count("order")) {
-			  ret->order=boost::lexical_cast<int>(vars["order"]);
+			  ret->order=boost::lexical_cast<int>(boost::get<string>(vars["order"]));
 			}
 
 			if(vars.count("weight")) {
-			  ret->weight=boost::lexical_cast<int>(vars["weight"]);
+			  ret->weight=boost::lexical_cast<int>(boost::get<string>(vars["weight"]));
 			}
 
 			if(g_launchWork) {
@@ -661,11 +669,11 @@ vector<std::function<void(void)>> setupLua(bool client)
      }});
 
   
-  std::ifstream ifs(g_vm["config"].as<string>());
+  std::ifstream ifs(config);
   if(!ifs) 
-    warnlog("Unable to read configuration from '%s'", g_vm["config"].as<string>());
+    warnlog("Unable to read configuration from '%s'", config);
   else
-    infolog("Read configuration from '%s'", g_vm["config"].as<string>());
+    infolog("Read configuration from '%s'", config);
 
   g_lua.executeCode(ifs);
   auto ret=*g_launchWork;
diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index ad8498511..61850447b 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -34,7 +34,9 @@
 #include <fstream>
 #include "sodcrypto.hh"
 #include "dnsrulactions.hh"
+#include <boost/program_options.hpp>
 #undef L
+namespace po = boost::program_options;
 
 /* Known sins:
    No centralized statistics
@@ -84,6 +86,7 @@ vector<ComboAddress> g_locals;
    Round-robin with basic uptime checks
    Send to least loaded server (least outstanding)
    Send it to the first server that is not overloaded
+   Hashed weighted random
 */
 
 /* Idea:
@@ -1064,12 +1067,12 @@ try
 
   g_policy.setState(leastOutstandingPol);
   if(g_vm.count("client") || g_vm.count("command")) {
-    setupLua(true);
+    setupLua(true, g_vm["config"].as<string>());
     doClient(g_serverControl);
     exit(EXIT_SUCCESS);
   }
 
-  auto todo=setupLua(false);
+  auto todo=setupLua(false, g_vm["config"].as<string>());
 
   if(g_vm.count("local")) {
     g_locals.clear();
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index b042558e2..abf6cdd7f 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -6,7 +6,6 @@
 #include "dnsname.hh"
 #include <atomic>
 #include <boost/circular_buffer.hpp>
-#include <boost/program_options.hpp>
 #include <mutex>
 #include <thread>
 #include "sholder.hh"
@@ -245,11 +244,7 @@ extern std::string g_key; // in theory needs locking
 struct dnsheader;
 
 void controlThread(int fd, ComboAddress local);
-vector<std::function<void(void)>> setupLua(bool client);
-
-
-namespace po = boost::program_options;
-extern po::variables_map g_vm;
+vector<std::function<void(void)>> setupLua(bool client, const std::string& config);
 NumberedServerVector getDownstreamCandidates(const servers_t& servers, const std::string& pool);
 
 std::shared_ptr<DownstreamState> firstAvailable(const NumberedServerVector& servers, const ComboAddress& remote, const DNSName& qname, uint16_t qtype, dnsheader* dh);
diff --git a/pdns/dnsdistconf.lua b/pdns/dnsdistconf.lua
index a44484015..7fd9c554e 100644
--- a/pdns/dnsdistconf.lua
+++ b/pdns/dnsdistconf.lua
@@ -55,9 +55,8 @@ end
 
 -- setServerPolicyLua("luaroundrobin", luaroundrobin)
 
-xs=newServer{address="2001:888:2000:1d::2", pool="auth"}
-newServer{address="2a01:4f8:110:4389::2", pool="auth"}
-xs:addPool("dnssec")
+newServer{address="2001:888:2000:1d::2", pool={"auth", "dnssec"}}
+newServer{address="2a01:4f8:110:4389::2", pool={"auth", "dnssec"}}
 setDNSSECPool("dnssec")
 topRule()
 
-- 
2.40.0