From 45a661edef3b7db54e103cd13c242623fff5a4e5 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 3 Jan 2020 18:06:54 +0100
Subject: [PATCH] dnsdist: Add direct Lua FFI interface for masking the remote
 addr

---
 pdns/dnsdistdist/dnsdist-lua-ffi-interface.h  | 1 +
 pdns/dnsdistdist/dnsdist-lua-ffi.cc           | 6 ++++++
 pdns/dnsdistdist/dnsdist-lua-ffi.hh           | 1 +
 pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc | 3 ++-
 4 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
index 327da5b498..daf4a50ebd 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
@@ -42,6 +42,7 @@ typedef struct dnsdist_tag {
 void dnsdist_ffi_dnsquestion_get_localaddr(const dnsdist_ffi_dnsquestion_t* dq, const void** addr, size_t* addrSize) __attribute__ ((visibility ("default")));
 uint16_t dnsdist_ffi_dnsquestion_get_local_port(const dnsdist_ffi_dnsquestion_t* dq) __attribute__ ((visibility ("default")));
 void dnsdist_ffi_dnsquestion_get_remoteaddr(const dnsdist_ffi_dnsquestion_t* dq, const void** addr, size_t* addrSize) __attribute__ ((visibility ("default")));
+void dnsdist_ffi_dnsquestion_get_masked_remoteaddr(dnsdist_ffi_dnsquestion_t* dq, const void** addr, size_t* addrSize, uint8_t bits) __attribute__ ((visibility ("default")));
 uint16_t dnsdist_ffi_dnsquestion_get_remote_port(const dnsdist_ffi_dnsquestion_t* dq) __attribute__ ((visibility ("default")));
 void dnsdist_ffi_dnsquestion_get_qname_raw(const dnsdist_ffi_dnsquestion_t* dq, const char** qname, size_t* qnameSize) __attribute__ ((visibility ("default")));
 size_t dnsdist_ffi_dnsquestion_get_qname_hash(const dnsdist_ffi_dnsquestion_t* dq, size_t init);
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
index 548357785b..efb7695854 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
@@ -55,6 +55,12 @@ void dnsdist_ffi_dnsquestion_get_remoteaddr(const dnsdist_ffi_dnsquestion_t* dq,
   dnsdist_ffi_comboaddress_to_raw(*dq->dq->remote, addr, addrSize);
 }
 
+void dnsdist_ffi_dnsquestion_get_masked_remoteaddr(dnsdist_ffi_dnsquestion_t* dq, const void** addr, size_t* addrSize, uint8_t bits)
+{
+  dq->maskedRemote = Netmask(*dq->dq->remote, bits).getMaskedNetwork();
+  dnsdist_ffi_comboaddress_to_raw(dq->maskedRemote, addr, addrSize);
+}
+
 uint16_t dnsdist_ffi_dnsquestion_get_local_port(const dnsdist_ffi_dnsquestion_t* dq)
 {
   return dq->dq->local->getPort();
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.hh b/pdns/dnsdistdist/dnsdist-lua-ffi.hh
index 118fd0ce8d..ca3e8d4edb 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.hh
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.hh
@@ -51,6 +51,7 @@ struct dnsdist_ffi_dnsquestion_t
   std::vector<dnsdist_tag_t> tagsVect;
   std::unordered_map<std::string, std::string> httpHeaders;
   std::string trailingData;
+  ComboAddress maskedRemote;
   boost::optional<std::string> result{boost::none};
   boost::optional<std::string> httpPath{boost::none};
   boost::optional<std::string> httpQueryString{boost::none};
diff --git a/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc b/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
index e04168b2e3..4f858d7606 100644
--- a/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
@@ -2,7 +2,6 @@
 #define BOOST_TEST_DYN_LINK
 #define BOOST_TEST_NO_MAIN
 
-#include <thread>
 #include <boost/test/unit_test.hpp>
 
 #include "dnsdist.hh"
@@ -505,6 +504,7 @@ BOOST_AUTO_TEST_CASE(test_lua) {
 }
 
 #ifdef LUAJIT_VERSION
+
 BOOST_AUTO_TEST_CASE(test_lua_ffi_rr) {
   std::vector<DNSName> names;
   names.reserve(1000);
@@ -726,6 +726,7 @@ BOOST_AUTO_TEST_CASE(test_lua_ffi_chashed) {
 
   g_verbose = existingVerboseValue;
 }
+
 #endif /* LUAJIT_VERSION */
 
 BOOST_AUTO_TEST_SUITE_END()
-- 
2.47.2