feat(dnsdist): centralize the TRACEPARENT packet mangling

diff --git a/pdns/dnsdistdist/dnsdist-opentelemetry.cc b/pdns/dnsdistdist/dnsdist-opentelemetry.cc
index d2fd43eaa7085954cb81b0fcc235e5bf1591b224..c0983a2e4a54689c8cf88347816e5bc4f2a75629 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-opentelemetry.cc
+++ b/pdns/dnsdistdist/dnsdist-opentelemetry.cc
@@ -22,7 +22,9 @@
 
 #include "dnsdist-opentelemetry.hh"
 #include "misc.hh"
+#include "dnsdist-ecs.hh"
 
+#include <memory>
 #include <vector>
 
 #ifndef DISABLE_PROTOBUF
@@ -339,4 +341,41 @@ std::vector<uint8_t> makeEDNSTraceParentOption(std::shared_ptr<Tracer> tracer)
 #endif
   return ret;
 }
+
+bool addTraceparentEdnsOptionToPacketBuffer(PacketBuffer& origBuf, const std::shared_ptr<Tracer>& tracer, const size_t qnameWireLength, const size_t proxyProtocolPayloadSize, const uint16_t traceparentOptionCode, const bool isTCP)
+{
+#ifndef DISABLE_PROTOBUF
+  // buf contains the whole DNS query without PROXY protocol and TCP length header
+  PacketBuffer buf{origBuf.begin() + proxyProtocolPayloadSize + (isTCP ? 2 : 0), origBuf.end()};
+
+  uint16_t optRDPosition;
+  size_t remaining;
+  bool queryHadEdns = ::dnsdist::getEDNSOptionsStart(buf, qnameWireLength, &optRDPosition, &remaining) == 0;
+  if (queryHadEdns) {
+    size_t optLen = buf.size() - optRDPosition - remaining;
+    removeEDNSOptionFromOPT(reinterpret_cast<char*>(buf.data() + optRDPosition), &optLen, traceparentOptionCode);
+  }
+
+  auto opt = pdns::trace::dnsdist::makeEDNSTraceParentOption(tracer);
+  bool ednsAdded{false};
+  bool optionAdded{false};
+  uint16_t maxEdnsSize = queryHadEdns ? (uint16_t)(buf.at(optRDPosition - 6) << 8) + buf.at(optRDPosition - 5) : 512;
+  setEDNSOption(buf, traceparentOptionCode, std::string(opt.begin(), opt.end()), isTCP ? std::numeric_limits<uint16_t>().max() : maxEdnsSize, ednsAdded, optionAdded);
+
+  if (isTCP) {
+    const std::array<uint8_t, 2> sizeBytes{static_cast<uint8_t>(buf.size() / 256), static_cast<uint8_t>(buf.size() % 256)};
+    buf.insert(buf.begin(), sizeBytes.begin(), sizeBytes.end());
+  }
+
+  // Resize the buffer to remove the existing packet, but keep any PROXYv2 data
+  origBuf.resize(proxyProtocolPayloadSize);
+  // Insert the new query into the buffer
+  origBuf.insert(origBuf.end(), buf.begin(), buf.end());
+
+  return ednsAdded;
+#else
+  return false;
+#endif
+}
+
 } // namespace pdns::trace::dnsdist

diff --git a/pdns/dnsdistdist/dnsdist-opentelemetry.hh b/pdns/dnsdistdist/dnsdist-opentelemetry.hh
index 74def92ef845b059f016948cab8d80fcf25741d3..138f4a7efb2a463c747d4c3d92b3207bc622bcc9 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-opentelemetry.hh
+++ b/pdns/dnsdistdist/dnsdist-opentelemetry.hh
@@ -25,6 +25,8 @@
 #include <string>
 #include <vector>
 
+#include "ednsoptions.hh"
+
 #ifndef DISABLE_PROTOBUF
 #include "protozero-trace.hh"
 using TraceID = pdns::trace::TraceID;
@@ -342,4 +344,5 @@ private:
 };
 
 std::vector<uint8_t> makeEDNSTraceParentOption(std::shared_ptr<Tracer> tracer);
+bool addTraceparentEdnsOptionToPacketBuffer(PacketBuffer& origBuf, const std::shared_ptr<Tracer>& tracer, const size_t qnameWireLength, const size_t proxyProtocolPayloadSize, const uint16_t traceparentOptionCode = EDNSOptionCode::TRACEPARENT, const bool isTCP = false);
 } // namespace pdns::trace::dnsdist

diff --git a/pdns/dnsdistdist/dnsdist-tcp-downstream.cc b/pdns/dnsdistdist/dnsdist-tcp-downstream.cc
index 173bae31b3a78b139a13f89e0941a0885b6871bf..01fad2f72ddf22e106fc554739cd1dc48859be40 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-tcp-downstream.cc
+++ b/pdns/dnsdistdist/dnsdist-tcp-downstream.cc
@@ -290,33 +290,14 @@ IOState TCPConnectionToBackend::sendQuery(std::shared_ptr<TCPConnectionToBackend
 
 #ifndef DISABLE_PROTOBUF
   if (auto tracer = conn->d_currentQuery.d_query.d_idstate.getTracer(); conn->d_currentQuery.d_query.d_idstate.sendTraceParentToDownstreamID != 0 && tracer != nullptr) {
-    // TODO: set a flag in ID State whether or not we should update the SpanID
-    uint16_t optRDPosition;
-    size_t remaining;
-    PacketBuffer buf{conn->d_currentQuery.d_query.d_buffer.begin() + conn->d_currentQuery.d_query.d_idstate.d_proxyProtocolPayloadSize + 2 /* TCP length */, conn->d_currentQuery.d_query.d_buffer.end()};
-    // clear any existing option
-    if (dnsdist::getEDNSOptionsStart(buf, conn->d_currentQuery.d_query.d_idstate.qname.wirelength(), &optRDPosition, &remaining) == 0) {
-      size_t optLen = buf.size() - optRDPosition - remaining;
-      removeEDNSOptionFromOPT(reinterpret_cast<char*>(buf.data() + optRDPosition), &optLen, conn->d_currentQuery.d_query.d_idstate.sendTraceParentToDownstreamID);
-    }
-
-    auto opt = pdns::trace::dnsdist::makeEDNSTraceParentOption(tracer);
-    bool ednsAdded, optionAdded;
-    setEDNSOption(buf, conn->d_currentQuery.d_query.d_idstate.sendTraceParentToDownstreamID, std::string(opt.begin(), opt.end()), std::numeric_limits<uint16_t>().max(), ednsAdded, optionAdded);
-    if (ednsAdded) {
-      conn->d_currentQuery.d_query.d_idstate.ednsAdded = ednsAdded;
-    }
-
-    // Calculate new TCP size
-    const std::array<uint8_t, 2> sizeBytes{static_cast<uint8_t>(buf.size() / 256), static_cast<uint8_t>(buf.size() % 256)};
-    buf.insert(buf.begin(), sizeBytes.begin(), sizeBytes.end());
-
-    // Resize the buffer twice to remove the existing packet, but keep any PROXYv2 data
-    conn->d_currentQuery.d_query.d_buffer.resize(conn->d_currentQuery.d_query.d_idstate.d_proxyProtocolPayloadSize);
-    conn->d_currentQuery.d_query.d_buffer.resize(conn->d_currentQuery.d_query.d_idstate.d_proxyProtocolPayloadSize + buf.size());
-
-    // Insert the new query into the buffer
-    conn->d_currentQuery.d_query.d_buffer.insert(conn->d_currentQuery.d_query.d_buffer.begin(), buf.begin(), buf.end());
+    auto ednsAdded = pdns::trace::dnsdist::addTraceparentEdnsOptionToPacketBuffer(
+      conn->d_currentQuery.d_query.d_buffer,
+      tracer,
+      conn->d_currentQuery.d_query.d_idstate.qname.wirelength(),
+      conn->d_currentQuery.d_query.d_idstate.d_proxyProtocolPayloadSize,
+      conn->d_currentQuery.d_query.d_idstate.sendTraceParentToDownstreamID,
+      true);
+    conn->d_currentQuery.d_query.d_idstate.ednsAdded = conn->d_currentQuery.d_query.d_idstate.ednsAdded || ednsAdded;
   }
 #endif
 

diff --git a/pdns/dnsdistdist/dnsdist.cc b/pdns/dnsdistdist/dnsdist.cc
index 5339db061a217f270294ac46adb9bac16f1b567c..16252a8a3185462a08c441448d7e1699e2f07e39 100644 (file)
--- a/pdns/dnsdistdist/dnsdist.cc
+++ b/pdns/dnsdistdist/dnsdist.cc
@@ -1891,26 +1891,14 @@ bool assignOutgoingUDPQueryToBackend(std::shared_ptr<DownstreamState>& downstrea
 
 #ifndef DISABLE_PROTOBUF
     if (auto tracer = dnsQuestion.ids.getTracer(); dnsQuestion.ids.sendTraceParentToDownstreamID != 0 && tracer != nullptr) {
-      // TODO: set a flag in ID State whether or not we should update the SpanID
-      uint16_t optRDPosition;
-      size_t remaining;
-      PacketBuffer buf{dnsQuestion.getData().begin() + dnsQuestion.ids.d_proxyProtocolPayloadSize, dnsQuestion.getData().end()};
-      // clear any existing option
-      if (dnsdist::getEDNSOptionsStart(buf, dnsQuestion.ids.qname.wirelength(), &optRDPosition, &remaining) == 0) {
-        size_t optLen = buf.size() - optRDPosition - remaining;
-        removeEDNSOptionFromOPT(reinterpret_cast<char*>(buf.data() + optRDPosition), &optLen, dnsQuestion.ids.sendTraceParentToDownstreamID);
-      }
-      auto opt = pdns::trace::dnsdist::makeEDNSTraceParentOption(tracer);
-
-      bool ednsAdded, optionAdded;
-      setEDNSOption(buf, dnsQuestion.ids.sendTraceParentToDownstreamID, std::string(opt.begin(), opt.end()), std::numeric_limits<uint16_t>().max(), ednsAdded, optionAdded);
-      if (ednsAdded) {
-        dnsQuestion.ids.ednsAdded = ednsAdded;
-      }
-
-      // Remove the query, but keep any PROXYv2 data
-      dnsQuestion.getMutableData().resize(dnsQuestion.ids.d_proxyProtocolPayloadSize);
-      dnsQuestion.getMutableData().insert(dnsQuestion.getMutableData().end(), buf.begin(), buf.end());
+      auto ednsAdded = pdns::trace::dnsdist::addTraceparentEdnsOptionToPacketBuffer(
+        dnsQuestion.getMutableData(),
+        tracer,
+        dnsQuestion.ids.qname.wirelength(),
+        dnsQuestion.ids.d_proxyProtocolPayloadSize,
+        dnsQuestion.ids.sendTraceParentToDownstreamID,
+        false);
+      dnsQuestion.ids.ednsAdded = dnsQuestion.ids.ednsAdded || ednsAdded;
     }
 #endif