From 3d8a6b1d04f636492347a43b457d7b11293a5aa3 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@powerdns.com>
Date: Mon, 25 Apr 2016 12:58:07 +0200
Subject: [PATCH] make TeeAction() quit properly, add getAction() so you can
 get access to an action, add getStats() to give you raw stats

---
 pdns/dnsdist-lua2.cc              | 19 ++++++++++-
 pdns/dnsdistdist/dnsrulactions.cc | 56 ++++++++++++++++++++++++-------
 pdns/dnsrulactions.hh             | 11 +++++-
 3 files changed, 72 insertions(+), 14 deletions(-)

diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 5759ffe5f..c21fe1899 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -584,13 +584,30 @@ void moreLua(bool client)
       });
 
     g_lua.writeFunction("TeeAction", [](const std::string& remote) {
+        setLuaNoSideEffect();
         return std::shared_ptr<DNSAction>(new TeeAction(ComboAddress(remote, 53)));
       });
 
     g_lua.registerFunction<void(DNSAction::*)()>("printStats", [](const DNSAction& ta) {
+        setLuaNoSideEffect();
         auto stats = ta.getStats();
         for(const auto& s : stats) {
-          g_outputBuffer+=s.first+"\t"+std::to_string(s.second)+"\n";
+          g_outputBuffer+=s.first+"\t";
+          if((uint64_t)s.second == s.second)
+            g_outputBuffer += std::to_string((uint64_t)s.second)+"\n";
+          else
+            g_outputBuffer += std::to_string(s.second)+"\n";
         }
       });
+
+    g_lua.writeFunction("getAction", [](unsigned int num) {
+        setLuaNoSideEffect();
+        auto rulactions = g_rulactions.getCopy();
+        if(num >= rulactions.size())
+          return std::shared_ptr<DNSAction>();
+        return rulactions[num].second;
+      });
+
+    g_lua.registerFunction("getStats", &DNSAction::getStats);
+
 }
diff --git a/pdns/dnsdistdist/dnsrulactions.cc b/pdns/dnsdistdist/dnsrulactions.cc
index 263cd73d3..852f02816 100644
--- a/pdns/dnsdistdist/dnsrulactions.cc
+++ b/pdns/dnsdistdist/dnsrulactions.cc
@@ -5,16 +5,14 @@ using namespace std;
 
 TeeAction::TeeAction(const ComboAddress& ca) : d_remote(ca)
 {
-  cerr<<"Created!"<<endl;
   d_fd=SSocket(d_remote.sin4.sin_family, SOCK_DGRAM, 0);
   SConnect(d_fd, d_remote);
-  d_worker=std::thread(std::bind(&TeeAction::worker, this));
-  
+  setNonBlocking(d_fd);
+  d_worker=std::thread(std::bind(&TeeAction::worker, this));  
 }
 
 TeeAction::~TeeAction()
 {
-  cerr<<"Closding down!"<<endl;
   d_pleaseQuit=true;
   close(d_fd);
   d_worker.join();
@@ -22,8 +20,13 @@ TeeAction::~TeeAction()
 
 DNSAction::Action TeeAction::operator()(DNSQuestion* dq, string* ruleresult) const 
 {
-  d_queries++;
-  send(d_fd, (char*)dq->dh, dq->len, 0);
+  if(dq->tcp) 
+    d_tcpdrops++;
+  else {
+    d_queries++;
+    if(send(d_fd, (char*)dq->dh, dq->len, 0) <= 0) 
+      d_senderrors++;
+  }
   return DNSAction::Action::None;
 }
 
@@ -35,21 +38,50 @@ string TeeAction::toString() const
 std::unordered_map<string,double> TeeAction::getStats() const
 {
   return {{"queries", d_queries},
-      {"responses", d_responses},
-        {"socket-errors", d_errors}};
+          {"responses", d_responses},
+          {"recv-errors", d_recverrors},
+          {"send-errors", d_senderrors},
+          {"noerrors", d_noerrors},
+          {"nxdomains", d_nxdomains},
+          {"refuseds", d_refuseds},
+          {"servfails", d_servfails},
+          {"other-rcode", d_otherrcode},
+          {"tcp-drops", d_tcpdrops}
+  };
 }
 
 void TeeAction::worker()
 {
   char packet[1500];
   int res=0;
+  struct dnsheader* dh=(struct dnsheader*)packet;
   for(;;) {
+    res=waitForData(d_fd, 0, 250000);
+    if(d_pleaseQuit)
+      break;
+    if(res < 0) {
+      usleep(250000);
+      continue;
+    }
+    if(res==0)
+      continue;
     res=recv(d_fd, packet, sizeof(packet), 0);
-    if(res < 0) 
-      d_errors++;
+    if(res <= 0) 
+      d_recverrors++;
     else if(res > 0)
       d_responses++;
-    if(d_pleaseQuit)
-      break;
+
+    if(dh->rcode == RCode::NoError)
+      d_noerrors++;
+    else if(dh->rcode == RCode::ServFail)
+      d_servfails++;
+    else if(dh->rcode == RCode::NXDomain)
+      d_nxdomains++;
+    else if(dh->rcode == RCode::Refused)
+      d_refuseds++;
+    else if(dh->rcode == RCode::FormErr)
+      d_formerrs++;
+    else if(dh->rcode == RCode::NotImp)
+      d_notimps++;
   }
 }
diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index ba6c75831..a9a4a0407 100644
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -425,9 +425,18 @@ private:
   void worker();
 
   int d_fd;
-  unsigned long d_errors{0};
+  mutable unsigned long d_senderrors{0};
+  unsigned long d_recverrors{0};
   mutable unsigned long d_queries{0};
   unsigned long d_responses{0};
+  unsigned long d_nxdomains{0};
+  unsigned long d_servfails{0};
+  unsigned long d_refuseds{0};
+  unsigned long d_formerrs{0};
+  unsigned long d_notimps{0};
+  unsigned long d_noerrors{0};
+  mutable unsigned long d_tcpdrops{0};
+  unsigned long d_otherrcode{0};
   std::atomic<bool> d_pleaseQuit{false};
 };
 
-- 
2.40.0