From e4ddc6e9daae3d79c0c28a6391f1380e378238c8 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 1 Sep 2020 09:53:49 +0200
Subject: [PATCH] sdig: Increment the DNS message IDs when pipelining

As stated in section 6.2.1 of rfc7766:
  "When sending multiple queries over a TCP connection, clients MUST NOT
   reuse the DNS Message ID of an in-flight query on that connection in
   order to avoid Message ID collisions.  This is especially important
   if the server could be performing out-of-order processing"
---
 pdns/sdig.cc | 29 ++++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/pdns/sdig.cc b/pdns/sdig.cc
index e6e3c71e34..ac89a2a6d7 100644
--- a/pdns/sdig.cc
+++ b/pdns/sdig.cc
@@ -57,11 +57,13 @@ static const string nameForClass(uint16_t qclass, uint16_t qtype)
   }
 }
 
+static std::unordered_set<uint16_t> s_expectedIDs;
+
 static void fillPacket(vector<uint8_t>& packet, const string& q, const string& t,
-  bool dnssec, const boost::optional<Netmask> ednsnm,
-  bool recurse, uint16_t xpfcode, uint16_t xpfversion,
-  uint64_t xpfproto, char* xpfsrc, char* xpfdst,
-  uint16_t qclass)
+                       bool dnssec, const boost::optional<Netmask> ednsnm,
+                       bool recurse, uint16_t xpfcode, uint16_t xpfversion,
+                       uint64_t xpfproto, char* xpfsrc, char* xpfdst,
+                       uint16_t qclass, uint16_t qid)
 {
   DNSPacketWriter pw(packet, DNSName(q), DNSRecordContent::TypeToNumber(t), qclass);
 
@@ -100,14 +102,18 @@ static void fillPacket(vector<uint8_t>& packet, const string& q, const string& t
   if (recurse) {
     pw.getHeader()->rd = true;
   }
+
+  pw.getHeader()->id = htons(qid);
 }
 
 static void printReply(const string& reply, bool showflags, bool hidesoadetails)
 {
   MOADNSParser mdp(false, reply);
-  if (mdp.d_header.id) {
-    cout << "ID is not zero, this response was not meant for us!"<<endl;
+  if (!s_expectedIDs.count(ntohs(mdp.d_header.id))) {
+    cout << "ID " << ntohs(mdp.d_header.id) << " was not expected, this response was not meant for us!"<<endl;
   }
+  s_expectedIDs.erase(ntohs(mdp.d_header.id));
+
   cout << "Reply to question for qname='" << mdp.d_qname.toString()
        << "', qtype=" << DNSRecordContent::NumberToType(mdp.d_qtype) << endl;
   cout << "Rcode: " << mdp.d_header.rcode << " ("
@@ -304,8 +310,9 @@ try {
   if (doh) {
 #ifdef HAVE_LIBCURL
     vector<uint8_t> packet;
+    s_expectedIDs.insert(0);
     fillPacket(packet, name, type, dnssec, ednsnm, recurse, xpfcode, xpfversion,
-      xpfproto, xpfsrc, xpfdst, qclass);
+               xpfproto, xpfsrc, xpfdst, qclass, 0);
     MiniCurl mc;
     MiniCurl::MiniCurlHeaders mch;
     mch.insert(std::make_pair("Content-Type", "application/dns-message"));
@@ -337,13 +344,16 @@ try {
 
     printReply(reply, showflags, hidesoadetails);
   } else if (tcp) {
+    uint16_t counter = 0;
     Socket sock(dest.sin4.sin_family, SOCK_STREAM);
     sock.connect(dest);
     sock.writen(proxyheader);
     for (const auto& it : questions) {
       vector<uint8_t> packet;
+      s_expectedIDs.insert(counter);
       fillPacket(packet, it.first, it.second, dnssec, ednsnm, recurse, xpfcode,
-        xpfversion, xpfproto, xpfsrc, xpfdst, qclass);
+                 xpfversion, xpfproto, xpfsrc, xpfdst, qclass, counter);
+      counter++;
 
       uint16_t len = htons(packet.size());
       if (sock.write((const char *)&len, 2) != 2)
@@ -374,8 +384,9 @@ try {
   } else // udp
   {
     vector<uint8_t> packet;
+    s_expectedIDs.insert(0);
     fillPacket(packet, name, type, dnssec, ednsnm, recurse, xpfcode, xpfversion,
-      xpfproto, xpfsrc, xpfdst, qclass);
+               xpfproto, xpfsrc, xpfdst, qclass, 0);
     string question(packet.begin(), packet.end());
     Socket sock(dest.sin4.sin_family, SOCK_DGRAM);
     question = proxyheader + question;
-- 
2.47.2