From ad4858964b0672584a2e13f676521169b862accf Mon Sep 17 00:00:00 2001
From: Remi Gacogne <rgacogne-github@coredump.fr>
Date: Tue, 3 Nov 2015 11:54:15 +0100
Subject: [PATCH] Allow custom DNSName and QType for dnsdist healthchecks

---
 pdns/dnsdist-lua.cc |  8 ++++++++
 pdns/dnsdist.cc     | 13 +++++++------
 pdns/dnsdist.hh     |  2 ++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 21ac9931b..679b538eb 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -168,6 +168,14 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 			  ret->name=boost::get<string>(vars["name"]);
 			}
 
+			if(vars.count("checkName")) {
+			  ret->checkName=DNSName(boost::get<string>(vars["checkName"]));
+			}
+
+			if(vars.count("checkType")) {
+			  ret->checkType=boost::get<string>(vars["checkType"]);
+			}
+
 			if(g_launchWork) {
 			  g_launchWork->push_back([ret]() {
 			      ret->tid = move(thread(responderThread, ret));
diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index b9a2e1c5d..c35281f9b 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -214,7 +214,8 @@ bool operator<(const struct timespec&a, const struct timespec& b)
 }
 
 
-DownstreamState::DownstreamState(const ComboAddress& remote_)
+DownstreamState::DownstreamState(const ComboAddress& remote_):
+checkName("a.root-servers.net."), checkType(QType::A)
 {
   remote = remote_;
   
@@ -556,11 +557,11 @@ catch(...)
 }
 
 
-bool upCheck(const ComboAddress& remote)
+bool upCheck(const ComboAddress& remote, const DNSName& checkName, const QType& checkType)
 try
 {
   vector<uint8_t> packet;
-  DNSPacketWriter dpw(packet, DNSName("a.root-servers.net."), QType::A);
+  DNSPacketWriter dpw(packet, checkName, checkType.getCode());
   dnsheader * requestHeader = dpw.getHeader();
   requestHeader->rd=true;
 
@@ -612,7 +613,7 @@ void* maintThread()
 
     for(auto& dss : g_dstates.getCopy()) { // this points to the actual shared_ptrs!
       if(dss->availability==DownstreamState::Availability::Auto) {
-	bool newState=upCheck(dss->remote);
+	bool newState=upCheck(dss->remote, dss->checkName, dss->checkType);
 	if(newState != dss->upStatus) {
 	  warnlog("Marking downstream %s as '%s'", dss->getName(), newState ? "up" : "down");
 	}
@@ -1124,8 +1125,8 @@ try
 
   for(auto& dss : g_dstates.getCopy()) { // it is a copy, but the internal shared_ptrs are the real deal
     if(dss->availability==DownstreamState::Availability::Auto) {
-      bool newState=upCheck(dss->remote);
-      warnlog("Marking downstream %s as '%s'", dss->getName(), newState ? "up" : "down");
+      bool newState=upCheck(dss->remote, dss->checkName, dss->checkType);
+      warnlog("Marking downstream %s as '%s'", dss->remote.toStringWithPort(), newState ? "up" : "down");
       dss->upStatus = newState;
     }
   }
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index cc6502ad7..271364bf3 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -231,6 +231,8 @@ struct DownstreamState
   ComboAddress remote;
   QPSLimiter qps;
   vector<IDState> idStates;
+  DNSName checkName;
+  QType checkType;
   std::atomic<uint64_t> idOffset{0};
   std::atomic<uint64_t> sendErrors{0};
   std::atomic<uint64_t> outstanding{0};
-- 
2.40.0