pdns/dumresp.cc

   1 /*
   2  * This file is part of PowerDNS or dnsdist.
   3  * Copyright -- PowerDNS.COM B.V. and its contributors
   4  *
   5  * This program is free software; you can redistribute it and/or modify
   6  * it under the terms of version 2 of the GNU General Public License as
   7  * published by the Free Software Foundation.
   8  *
   9  * In addition, for the avoidance of any doubt, permission is granted to
  10  * link this program with OpenSSL and to (re)distribute the binaries
  11  * produced as the result of such linking.
  12  *
  13  * This program is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16  * GNU General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License
  19  * along with this program; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  21  */
  22 #ifdef HAVE_CONFIG_H
  23 #include "config.h"
  24 #endif
  25 #include "iputils.hh"
  26 #include "sstuff.hh"
  27 #include "statbag.hh"
  28 #include <atomic>
  29 #include <sys/mman.h>
  30 #include <thread>
  31 StatBag S;
  32
  33 static std::atomic<uint64_t>* g_counter;
  34
  35 static void printStatus()
  36 {
  37   auto prev= g_counter->load();
  38   for(;;) {
  39     sleep(1);
  40     cout<<g_counter->load()-prev<<"\t"<<g_counter->load()<<endl;
  41     prev=g_counter->load();
  42   }
  43 }
  44
  45 static void usage() {
  46   cerr<<"Syntax: dumresp LOCAL-ADDRESS LOCAL-PORT NUMBER-OF-PROCESSES [tcp]"<<endl;
  47 }
  48
  49 static void turnQueryIntoResponse(dnsheader* dh)
  50 {
  51   (*g_counter)++;
  52
  53   dh->qr=1;
  54   dh->ad=0;
  55 }
  56
  57 static void tcpConnectionHandler(int sock)
  58 try
  59 {
  60   char buffer[1500];
  61   auto dh = reinterpret_cast<struct dnsheader*>(buffer);
  62
  63   for (;;) {
  64     uint16_t len = 0;
  65     ssize_t got = read(sock, &len, sizeof(len));
  66
  67     if (got == 0) {
  68       break;
  69     }
  70
  71     if (got != sizeof(len))
  72       unixDie("read 1");
  73
  74     len = ntohs(len);
  75
  76     if (len < sizeof(dnsheader))
  77       unixDie("too small");
  78
  79     if (len > sizeof(buffer))
  80       unixDie("too large");
  81
  82     got = read(sock, buffer, len);
  83     if (got != len)
  84       unixDie("read 2: " + std::to_string(got) + " / " + std::to_string(len));
  85
  86     if (dh->qr)
  87       continue;
  88
  89     turnQueryIntoResponse(dh);
  90
  91     uint16_t wirelen = htons(len);
  92     if (write(sock, &wirelen, sizeof(wirelen)) != sizeof(wirelen))
  93       unixDie("send 1");
  94
  95     if (write(sock, buffer, len) < 0)
  96       unixDie("send 2");
  97   }
  98
  99   close(sock);
 100 }
 101 catch(const std::exception& e) {
 102   cerr<<"TCP connection handler got an exception: "<<e.what()<<endl;
 103 }
 104
 105 static void tcpAcceptor(const ComboAddress local)
 106 {
 107   Socket tcpSocket(local.sin4.sin_family, SOCK_STREAM);
 108   setReusePort(tcpSocket.getHandle());
 109   tcpSocket.bind(local);
 110   tcpSocket.listen(1024);
 111
 112   ComboAddress rem("::1");
 113   auto socklen = rem.getSocklen();
 114
 115   for (;;) {
 116     int sock = accept(tcpSocket.getHandle(), reinterpret_cast<struct sockaddr*>(&rem), &socklen);
 117     if (sock == -1) {
 118         continue;
 119     }
 120
 121     std::thread connectionHandler(tcpConnectionHandler, sock);
 122     connectionHandler.detach();
 123   }
 124 }
 125
 126 int main(int argc, char** argv)
 127 try
 128 {
 129   bool tcp = false;
 130
 131   for(int i = 1; i < argc; i++) {
 132     if(std::string(argv[i]) == "--help"){
 133       usage();
 134       return(EXIT_SUCCESS);
 135     }
 136
 137     if(std::string(argv[i]) == "--version"){
 138       cerr<<"dumresp "<<VERSION<<endl;
 139       return(EXIT_SUCCESS);
 140     }
 141   }
 142
 143   if(argc == 5) {
 144     if (std::string(argv[4]) == "tcp") {
 145       tcp = true;
 146     }
 147     else {
 148       usage();
 149       exit(EXIT_FAILURE);
 150     }
 151   }
 152   else if(argc != 4) {
 153     usage();
 154     exit(EXIT_FAILURE);
 155   }
 156
 157   auto ptr = mmap(nullptr, sizeof(std::atomic<uint64_t>), PROT_READ | PROT_WRITE,
 158                   MAP_SHARED | MAP_ANONYMOUS, -1, 0);
 159
 160   g_counter = new(ptr) std::atomic<uint64_t>();
 161
 162   int numberOfListeners = atoi(argv[3]);
 163   ComboAddress local(argv[1], atoi(argv[2]));
 164
 165   int i=1;
 166   for(; i < numberOfListeners; ++i) {
 167     if(!fork())
 168       break;
 169   }
 170
 171   if (i==1) {
 172     std::thread t(printStatus);
 173     t.detach();
 174
 175     if (tcp) {
 176       for (int j = 0; j < numberOfListeners; j++) {
 177         cout<<"Listening to TCP "<<local.toStringWithPort()<<endl;
 178         std::thread tcpAcceptorThread(tcpAcceptor, local);
 179         tcpAcceptorThread.detach();
 180       }
 181     }
 182   }
 183
 184   Socket s(local.sin4.sin_family, SOCK_DGRAM);
 185   setReusePort(s.getHandle());
 186   s.bind(local);
 187   cout<<"Bound to UDP "<<local.toStringWithPort()<<endl;
 188
 189   ComboAddress rem = local;
 190   socklen_t socklen = rem.getSocklen();
 191   char buffer[1500];
 192   auto dh = reinterpret_cast<struct dnsheader*>(buffer);
 193
 194   for(;;) {
 195     ssize_t len = recvfrom(s.getHandle(), buffer, sizeof(buffer), 0, reinterpret_cast<struct sockaddr*>(&rem), &socklen);
 196
 197     if(len < 0)
 198       unixDie("recvfrom");
 199
 200     if (static_cast<size_t>(len) < sizeof(dnsheader))
 201       unixDie("too small " + std::to_string(len));
 202
 203     if(dh->qr)
 204       continue;
 205
 206     turnQueryIntoResponse(dh);
 207
 208     if(sendto(s.getHandle(), buffer, len, 0,  reinterpret_cast<const struct sockaddr*>(&rem), socklen) < 0)
 209       unixDie("sendto");
 210   }
 211 }
 212 catch(const std::exception& e)
 213 {
 214   cerr<<"Fatal error: "<<e.what()<<endl;
 215   exit(EXIT_FAILURE);
 216 }