From f5062066366874f42c7a307cdc93bb67ad44da2d Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Mon, 28 Dec 2015 10:12:22 +0000
Subject: [PATCH] make ipfilter get passed the dnsheader, make dnsheader useful
 for lua

---
 pdns/lua-recursor4.cc            | 17 ++++++++++++++++-
 pdns/lua-recursor4.hh            |  2 +-
 pdns/powerdns-example-script.lua |  5 +++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/pdns/lua-recursor4.cc b/pdns/lua-recursor4.cc
index 12916218ab..64fb427338 100644
--- a/pdns/lua-recursor4.cc
+++ b/pdns/lua-recursor4.cc
@@ -167,6 +167,21 @@ struct DynMetric
 RecursorLua4::RecursorLua4(const std::string& fname)
 {
   d_lw = new LuaContext;
+
+  d_lw->registerFunction<int(dnsheader::*)()>("getID", [](dnsheader& dh) { return dh.id; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getCD", [](dnsheader& dh) { return dh.cd; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getTC", [](dnsheader& dh) { return dh.tc; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getRA", [](dnsheader& dh) { return dh.ra; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getAD", [](dnsheader& dh) { return dh.ad; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getAA", [](dnsheader& dh) { return dh.aa; });
+  d_lw->registerFunction<bool(dnsheader::*)()>("getRD", [](dnsheader& dh) { return dh.rd; });
+  d_lw->registerFunction<int(dnsheader::*)()>("getRCODE", [](dnsheader& dh) { return dh.rcode; });
+  d_lw->registerFunction<int(dnsheader::*)()>("getOPCODE", [](dnsheader& dh) { return dh.opcode; });
+  d_lw->registerFunction<int(dnsheader::*)()>("getQDCOUNT", [](dnsheader& dh) { return ntohs(dh.qdcount); });
+  d_lw->registerFunction<int(dnsheader::*)()>("getANCOUNT", [](dnsheader& dh) { return ntohs(dh.ancount); });
+  d_lw->registerFunction<int(dnsheader::*)()>("getNSCOUNT", [](dnsheader& dh) { return ntohs(dh.nscount); });
+  d_lw->registerFunction<int(dnsheader::*)()>("getARCOUNT", [](dnsheader& dh) { return ntohs(dh.arcount); });
+
   d_lw->writeFunction("newDN", [](const std::string& dom){ return DNSName(dom); });  
   d_lw->registerFunction("isPartOf", &DNSName::isPartOf);
   d_lw->registerFunction<bool(DNSName::*)(const std::string&)>("equal",
@@ -306,7 +321,7 @@ bool RecursorLua4::preoutquery(const ComboAddress& ns, const ComboAddress& reque
 bool RecursorLua4::ipfilter(const ComboAddress& remote, const ComboAddress& local, const struct dnsheader& dh)
 {
   if(d_ipfilter)
-    return d_ipfilter({remote}, {local});
+    return d_ipfilter(remote, local, dh);
   return false; // don't block
 }
 
diff --git a/pdns/lua-recursor4.hh b/pdns/lua-recursor4.hh
index 6af51a5094..ee8b56c632 100644
--- a/pdns/lua-recursor4.hh
+++ b/pdns/lua-recursor4.hh
@@ -49,7 +49,7 @@ private:
   typedef std::function<bool(std::shared_ptr<DNSQuestion>)> luacall_t;
   luacall_t d_preresolve, d_nxdomain, d_nodata, d_postresolve, d_preoutquery, d_postoutquery;
   bool genhook(luacall_t& func, const ComboAddress& remote,const ComboAddress& local, const DNSName& query, const QType& qtype, vector<DNSRecord>& res, int& ret, bool* variable);
-  typedef std::function<bool(ComboAddress,ComboAddress)> ipfilter_t;
+  typedef std::function<bool(ComboAddress,ComboAddress, struct dnsheader)> ipfilter_t;
   ipfilter_t d_ipfilter;
 };
 
diff --git a/pdns/powerdns-example-script.lua b/pdns/powerdns-example-script.lua
index 8177b0ab34..23a1ab2c15 100644
--- a/pdns/powerdns-example-script.lua
+++ b/pdns/powerdns-example-script.lua
@@ -91,8 +91,9 @@ badips = newNMG()
 badips:addMask("127.1.0.0/16")
 
 -- this check is applied before any packet parsing is done
-function ipfilter(loc, rem)
-	print("ipfilter called, rem: ", rem:toString(), badips:match(rem))
+function ipfilter(rem, loc, dh)
+	print("ipfilter called, rem: ", rem:toString(), "loc: ",loc:toString(),"match:", badips:match(rem))
+	print("id: ",dh:getID(), "aa: ", dh:getAA(), "ad: ", dh:getAD(), "arcount: ", dh:getARCOUNT())
 	return badips:match(rem)
 end
 
-- 
2.47.2