From fde296a33609c89a5414b4cbb363f5b6b834ee00 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Mon, 19 Sep 2016 18:09:19 +0200
Subject: [PATCH] rec: Add TCP management options described in section 10 of
 rfc7766

* max-tcp-queries-per-connection
---
 docs/markdown/recursor/settings.md |  6 ++++++
 pdns/pdns_recursor.cc              | 17 +++++++++++++----
 pdns/syncres.hh                    |  1 +
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/docs/markdown/recursor/settings.md b/docs/markdown/recursor/settings.md
index 61572b8a1a..0901eab509 100644
--- a/docs/markdown/recursor/settings.md
+++ b/docs/markdown/recursor/settings.md
@@ -613,6 +613,12 @@ Maximum number of simultaneous incoming TCP connections allowed.
 Maximum number of simultaneous incoming TCP connections allowed per client
 (remote IP address).
 
+## `max-tcp-queries-per-connection`
+* Integer
+* Default: 0 (unlimited)
+
+Maximum number of DNS queries in a TCP connection.
+
 ## `max-total-msec`
 * Integer
 * Default: 7000
diff --git a/pdns/pdns_recursor.cc b/pdns/pdns_recursor.cc
index 487fabfe78..bfdeac602f 100644
--- a/pdns/pdns_recursor.cc
+++ b/pdns/pdns_recursor.cc
@@ -93,6 +93,7 @@ extern SortList g_sortlist;
 __thread FDMultiplexer* t_fdm;
 __thread unsigned int t_id;
 unsigned int g_maxTCPPerClient;
+size_t g_tcpMaxQueriesPerConn;
 unsigned int g_networkTimeoutMsec;
 uint64_t g_latencyStatSize;
 bool g_logCommonErrors;
@@ -1139,10 +1140,16 @@ void startDoResolve(void *p)
         dc->d_socket = -1;
       }
       else {
-        dc->d_tcpConnection->state=TCPConnection::BYTE0;
-        Utility::gettimeofday(&g_now, 0); // needs to be updated
-        t_fdm->addReadFD(dc->d_socket, handleRunningTCPQuestion, dc->d_tcpConnection);
-        t_fdm->setReadTTD(dc->d_socket, g_now, g_tcpTimeout);
+        dc->d_tcpConnection->queriesCount++;
+        if (g_tcpMaxQueriesPerConn && dc->d_tcpConnection->queriesCount >= g_tcpMaxQueriesPerConn) {
+          dc->d_socket = -1;
+        }
+        else {
+          dc->d_tcpConnection->state=TCPConnection::BYTE0;
+          Utility::gettimeofday(&g_now, 0); // needs to be updated
+          t_fdm->addReadFD(dc->d_socket, handleRunningTCPQuestion, dc->d_tcpConnection);
+          t_fdm->setReadTTD(dc->d_socket, g_now, g_tcpTimeout);
+        }
       }
     }
 
@@ -2789,6 +2796,7 @@ int serviceMain(int argc, char*argv[])
 
   g_tcpTimeout=::arg().asNum("client-tcp-timeout");
   g_maxTCPPerClient=::arg().asNum("max-tcp-per-client");
+  g_tcpMaxQueriesPerConn=::arg().asNum("max-tcp-queries-per-connection");
 
   if(g_numThreads == 1) {
     L<<Logger::Warning<<"Operating unthreaded"<<endl;
@@ -3041,6 +3049,7 @@ int main(int argc, char **argv)
     ::arg().set("entropy-source", "If set, read entropy from this file")="/dev/urandom";
     ::arg().set("dont-query", "If set, do not query these netmasks for DNS data")=DONT_QUERY;
     ::arg().set("max-tcp-per-client", "If set, maximum number of TCP sessions per client (IP address)")="0";
+    ::arg().set("max-tcp-queries-per-connection", "If set, maximum number of TCP queries in a TCP connection")="0";
     ::arg().set("spoof-nearmiss-max", "If non-zero, assume spoofing after this many near misses")="20";
     ::arg().set("single-socket", "If set, only use a single socket for outgoing queries")="off";
     ::arg().set("auth-zones", "Zones for which we have authoritative data, comma separated domain=file pairs ")="";
diff --git a/pdns/syncres.hh b/pdns/syncres.hh
index 8820d35968..8621169780 100644
--- a/pdns/syncres.hh
+++ b/pdns/syncres.hh
@@ -659,6 +659,7 @@ public:
   uint16_t bytesread{0};
   const ComboAddress d_remote;
   char data[65535]; // damn
+  size_t queriesCount{0};
 
   static unsigned int getCurrentConnections() { return s_currentConnections; }
 private:
-- 
2.47.2