dnsdist: Randomized IDs over UDP

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 367de3b6dc0f69a820e09596c5929f0cc524a43a..8b7bceb1b35d4a5807395892ff90d73cb765ad55 100644 (file)
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -2783,6 +2783,10 @@ static void setupLuaConfig(LuaContext& luaCtx, bool client, bool configCheck)
     DownstreamState::s_randomizeSockets = randomized;
   });
 
+  luaCtx.writeFunction("setRandomizedIdsOverUDP", [](bool randomized) {
+    DownstreamState::s_randomizeIDs = randomized;
+  });
+
 #if defined(HAVE_LIBSSL)
   luaCtx.writeFunction("loadTLSEngine", [client](const std::string& engineName, boost::optional<std::string> defaultString) {
     if (client) {

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index aa8d52d9a5d6df87628e8f4d54a66c8b54df8f44..9695d1619de95487bf04010a315f3459c9628000 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -1524,33 +1524,9 @@ static void processUDPQuery(ClientState& cs, LocalHolders& holders, const struct
       return;
     }
 
-    unsigned int idOffset = (ss->idOffset++) % ss->idStates.size();
-    IDState* ids = &ss->idStates[idOffset];
-    ids->age = 0;
-    DOHUnitUniquePtr du(nullptr, DOHUnit::release);
-
-    /* that means that the state was in use, possibly with an allocated
-       DOHUnit that we will need to handle, but we can't touch it before
-       confirming that we now own this state */
-    if (ids->isInUse()) {
-      du = DOHUnitUniquePtr(ids->du, DOHUnit::release);
-    }
-
-    /* we atomically replace the value, we now own this state */
-    if (!ids->markAsUsed()) {
-      /* the state was not in use.
-         we reset 'du' because it might have still been in use when we read it. */
-      du.release();
-      ++ss->outstanding;
-    }
-    else {
-      /* we are reusing a state, no change in outstanding but if there was an existing DOHUnit we need
-         to handle it because it's about to be overwritten. */
-      ids->du = nullptr;
-      ++ss->reuseds;
-      ++g_stats.downstreamTimeouts;
-      handleDOHTimeout(std::move(du));
-    }
+    unsigned int idOffset = 0;
+    int64_t generation;
+    IDState* ids = ss->getIDState(idOffset, generation);
 
     ids->cs = &cs;
     ids->origFD = cs.udpFD;

diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 95db2e8550aeaf0ef67df5f6e1be1516ffa91c6d..6205895df2b92930942c9e60a6dd5cec29aecfe4 100644 (file)
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -882,6 +882,7 @@ public:
   bool passCrossProtocolQuery(std::unique_ptr<CrossProtocolQuery>&& cpq);
   int pickSocketForSending();
   void pickSocketsReadyForReceiving(std::vector<int>& ready);
+  IDState* getIDState(unsigned int& id, int64_t& generation);
 
   dnsdist::Protocol getProtocol() const
   {
@@ -898,6 +899,7 @@ public:
   }
 
   static bool s_randomizeSockets;
+  static bool s_randomizeIDs;
 };
 using servers_t =vector<std::shared_ptr<DownstreamState>>;
 

diff --git a/pdns/dnsdistdist/dnsdist-backend.cc b/pdns/dnsdistdist/dnsdist-backend.cc
index cd8c2be21740c3c817b8e07e67c68ff55ed5238b..e005595aaaf9ef3d193f59eb81557fc2c3791c3b 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-backend.cc
+++ b/pdns/dnsdistdist/dnsdist-backend.cc
@@ -244,6 +244,57 @@ void DownstreamState::pickSocketsReadyForReceiving(std::vector<int>& ready)
 }
 
 bool DownstreamState::s_randomizeSockets{false};
+bool DownstreamState::s_randomizeIDs{false};
+
+IDState* DownstreamState::getIDState(unsigned int& selectedID, int64_t& generation)
+{
+  DOHUnitUniquePtr du(nullptr, DOHUnit::release);
+  IDState* ids = nullptr;
+  if (s_randomizeIDs) {
+    /* if the state is already in use we will retry,
+       up to 5 five times. The last selected one is used
+       even if it was already in use */
+    size_t remainingAttempts = 5;
+    do {
+      selectedID = dnsdist::getRandomValue(idStates.size());
+      ids = &idStates[selectedID];
+      remainingAttempts--;
+    }
+    while (ids->isInUse() && remainingAttempts > 0);
+  }
+  else {
+    selectedID = (idOffset++) % idStates.size();
+    ids = &idStates[selectedID];
+  }
+
+  ids->age = 0;
+
+  /* that means that the state was in use, possibly with an allocated
+     DOHUnit that we will need to handle, but we can't touch it before
+     confirming that we now own this state */
+  if (ids->isInUse()) {
+    du = DOHUnitUniquePtr(ids->du, DOHUnit::release);
+  }
+
+  /* we atomically replace the value, we now own this state */
+  generation = ids->generation++;
+  if (!ids->markAsUsed(generation)) {
+    /* the state was not in use.
+       we reset 'du' because it might have still been in use when we read it. */
+    du.release();
+    ++outstanding;
+  }
+  else {
+    /* we are reusing a state, no change in outstanding but if there was an existing DOHUnit we need
+       to handle it because it's about to be overwritten. */
+    ids->du = nullptr;
+    ++reuseds;
+    ++g_stats.downstreamTimeouts;
+    handleDOHTimeout(std::move(du));
+  }
+
+  return ids;
+}
 
 size_t ServerPool::countServers(bool upOnly)
 {

diff --git a/pdns/dnsdistdist/doh.cc b/pdns/dnsdistdist/doh.cc
index e06d64737f6a38bd6e7ab47c94ec5b266fed0dc0..dc8ef1400110c319f9b4bd17991a17896e92b145 100644 (file)
--- a/pdns/dnsdistdist/doh.cc
+++ b/pdns/dnsdistdist/doh.cc
@@ -661,33 +661,9 @@ static void processDOHQuery(DOHUnitUniquePtr&& du)
     }
 
     ComboAddress dest = du->ids.origDest;
-    unsigned int idOffset = (du->downstream->idOffset++) % du->downstream->idStates.size();
-    IDState* ids = &du->downstream->idStates[idOffset];
-    ids->age = 0;
-    DOHUnit* oldDU = nullptr;
-    if (ids->isInUse()) {
-      /* that means that the state was in use, possibly with an allocated
-         DOHUnit that we will need to handle, but we can't touch it before
-         confirming that we now own this state */
-      oldDU = ids->du;
-    }
-
-    /* we atomically replace the value, we now own this state */
-    int64_t generation = ids->generation++;
-    if (!ids->markAsUsed(generation)) {
-      /* the state was not in use.
-         we reset 'oldDU' because it might have still been in use when we read it. */
-      oldDU = nullptr;
-      ++du->downstream->outstanding;
-    }
-    else {
-      ids->du = nullptr;
-      /* we are reusing a state, no change in outstanding but if there was an existing DOHUnit we need
-         to handle it because it's about to be overwritten. */
-      ++du->downstream->reuseds;
-      ++g_stats.downstreamTimeouts;
-      handleDOHTimeout(DOHUnitUniquePtr(oldDU, DOHUnit::release));
-    }
+    unsigned int idOffset = 0;
+    int64_t generation;
+    IDState* ids = du->downstream->getIDState(idOffset, generation);
 
     ids->origFD = 0;
     /* increase the ref count since we are about to store the pointer */

diff --git a/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc b/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
index d1f5080836098b852516d2eba7145c7032564d54..32d6f2eddaee32d6049689da3608f9fb3d99f383 100644 (file)
--- a/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistlbpolicies_cc.cc
@@ -64,6 +64,10 @@ void DOHUnit::setHTTPResponse(uint16_t statusCode, PacketBuffer&& body_, const s
 }
 #endif /* HAVE_DNS_OVER_HTTPS */
 
+void handleDOHTimeout(DOHUnitUniquePtr&& oldDU)
+{
+}
+
 std::string DNSQuestion::getTrailingData() const
 {
   return "";