From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 2 Apr 2019 14:16:25 +0000 (+0200)
Subject: dumresp: Add TCP support
X-Git-Tag: rec-4.2.0-rc1~41^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=65ae96092bece0be8122ce922e61f96fbdd02b54;p=thirdparty%2Fpdns.git

dumresp: Add TCP support
---

diff --git a/pdns/dumresp.cc b/pdns/dumresp.cc
index 8d12c5121f..354eba7461 100644
--- a/pdns/dumresp.cc
+++ b/pdns/dumresp.cc
@@ -30,9 +30,9 @@
 #include <thread>
 StatBag S;
 
-std::atomic<uint64_t>* g_counter;
+static std::atomic<uint64_t>* g_counter;
 
-void printStatus()
+static void printStatus()
 {
   auto prev= g_counter->load();
   for(;;) {
@@ -42,47 +42,151 @@ void printStatus()
   }
 }
 
-void usage() {
-  cerr<<"Syntax: dumresp LOCAL-ADDRESS LOCAL-PORT NUMBER-OF-PROCESSES"<<endl;
+static void usage() {
+  cerr<<"Syntax: dumresp LOCAL-ADDRESS LOCAL-PORT NUMBER-OF-PROCESSES [tcp]"<<endl;
+}
+
+static void turnQueryIntoResponse(dnsheader* dh)
+{
+  (*g_counter)++;
+
+  dh->qr=1;
+  dh->ad=0;
+}
+
+static void tcpConnectionHandler(int sock)
+try
+{
+  char buffer[1500];
+  auto dh = reinterpret_cast<struct dnsheader*>(buffer);
+
+  for (;;) {
+    uint16_t len = 0;
+    ssize_t got = read(sock, &len, sizeof(len));
+
+    if (got == 0) {
+      break;
+    }
+
+    if (got != sizeof(len))
+      unixDie("read 1");
+
+    len = ntohs(len);
+
+    if (len < sizeof(dnsheader))
+      unixDie("too small");
+
+    if (len > sizeof(buffer))
+      unixDie("too large");
+
+    got = read(sock, buffer, len);
+    if (got != len)
+      unixDie("read 2: " + std::to_string(got) + " / " + std::to_string(len));
+
+    if (dh->qr)
+      continue;
+
+    turnQueryIntoResponse(dh);
+
+    uint16_t wirelen = htons(len);
+    if (write(sock, &wirelen, sizeof(wirelen)) != sizeof(wirelen))
+      unixDie("send 1");
+
+    if (write(sock, buffer, len) < 0)
+      unixDie("send 2");
+  }
+
+  close(sock);
+}
+catch(const std::exception& e) {
+  cerr<<"TCP connection handler got an exception: "<<e.what()<<endl;
+}
+
+static void tcpAcceptor(const ComboAddress local)
+{
+  Socket tcpSocket(local.sin4.sin_family, SOCK_STREAM);
+#ifdef SO_REUSEPORT
+  int one=1;
+  if(setsockopt(tcpSocket.getHandle(), SOL_SOCKET, SO_REUSEPORT, &one, sizeof(one)) < 0)
+    unixDie("setsockopt for REUSEPORT");
+#endif
+
+  tcpSocket.bind(local);
+  tcpSocket.listen(1024);
+
+  ComboAddress rem("::1");
+  auto socklen = rem.getSocklen();
+
+  for (;;) {
+    int sock = accept(tcpSocket.getHandle(), reinterpret_cast<struct sockaddr*>(&rem), &socklen);
+    if (sock == -1) {
+        continue;
+    }
+
+    std::thread connectionHandler(tcpConnectionHandler, sock);
+    connectionHandler.detach();
+  }
 }
 
 int main(int argc, char** argv)
 try
 {
+  bool tcp = false;
+
   for(int i = 1; i < argc; i++) {
-    if((string) argv[i] == "--help"){
+    if(std::string(argv[i]) == "--help"){
       usage();
       return(EXIT_SUCCESS);
     }
 
-    if((string) argv[i] == "--version"){
+    if(std::string(argv[i]) == "--version"){
       cerr<<"dumresp "<<VERSION<<endl;
       return(EXIT_SUCCESS);
     }
   }
 
-  if(argc != 4) {
+  if(argc == 5) {
+    if (std::string(argv[4]) == "tcp") {
+      tcp = true;
+    }
+    else {
+      usage();
+      exit(EXIT_FAILURE);
+    }
+  }
+  else if(argc != 4) {
     usage();
     exit(EXIT_FAILURE);
   }
 
-  auto ptr = mmap(NULL, sizeof(std::atomic<uint64_t>), PROT_READ | PROT_WRITE,
+  auto ptr = mmap(nullptr, sizeof(std::atomic<uint64_t>), PROT_READ | PROT_WRITE,
 		  MAP_SHARED | MAP_ANONYMOUS, -1, 0);
 
   g_counter = new(ptr) std::atomic<uint64_t>();
-  
+
+  int numberOfListeners = atoi(argv[3]);
+  ComboAddress local(argv[1], atoi(argv[2]));
+
   int i=1;
-  for(; i < atoi(argv[3]); ++i) {
+  for(; i < numberOfListeners; ++i) {
     if(!fork())
       break;
   }
-  if(i==1) {
+
+  if (i==1) {
     std::thread t(printStatus);
     t.detach();
+
+    if (tcp) {
+      for (int j = 0; j < numberOfListeners; j++) {
+        cout<<"Listening to TCP "<<local.toStringWithPort()<<endl;
+        std::thread tcpAcceptorThread(tcpAcceptor, local);
+        tcpAcceptorThread.detach();
+      }
+    }
   }
-  
-  ComboAddress local(argv[1], atoi(argv[2]));
-  Socket s(local.sin4.sin_family, SOCK_DGRAM);  
+
+  Socket s(local.sin4.sin_family, SOCK_DGRAM);
 #ifdef SO_REUSEPORT
   int one=1;
   if(setsockopt(s.getHandle(), SOL_SOCKET, SO_REUSEPORT, &one, sizeof(one)) < 0)
@@ -90,28 +194,32 @@ try
 #endif
 
   s.bind(local);
-  cout<<"Bound to "<<local.toStringWithPort()<<endl;
-  char buffer[1500];
-  struct dnsheader* dh = (struct dnsheader*)buffer;
-  int len;
-  ComboAddress rem=local;
+  cout<<"Bound to UDP "<<local.toStringWithPort()<<endl;
+
+  ComboAddress rem = local;
   socklen_t socklen = rem.getSocklen();
+  char buffer[1500];
+  auto dh = reinterpret_cast<struct dnsheader*>(buffer);
+
   for(;;) {
-    len=recvfrom(s.getHandle(), buffer, sizeof(buffer), 0, (struct sockaddr*)&rem, &socklen);
-    (*g_counter)++;
+    uint16_t len = recvfrom(s.getHandle(), buffer, sizeof(buffer), 0, reinterpret_cast<struct sockaddr*>(&rem), &socklen);
+
     if(len < 0)
       unixDie("recvfrom");
 
+    if (len < sizeof(dnsheader))
+      unixDie("too small " + std::to_string(len));
+
     if(dh->qr)
       continue;
-    dh->qr=1;
-    dh->ad=0;
-    if(sendto(s.getHandle(), buffer, len, 0,  (struct sockaddr*)&rem, socklen) < 0)
-      unixDie("sendto");
 
+    turnQueryIntoResponse(dh);
+
+    if(sendto(s.getHandle(), buffer, len, 0,  reinterpret_cast<const struct sockaddr*>(&rem), socklen) < 0)
+      unixDie("sendto");
   }
 }
-catch(std::exception& e)
+catch(const std::exception& e)
 {
   cerr<<"Fatal error: "<<e.what()<<endl;
   exit(EXIT_FAILURE);