From 080e4e806fa4f8d98e48348c2b5bb761fa97bd71 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 30 May 2025 11:42:32 +0200
Subject: [PATCH] dnsdist: Make `setEDNSOption` usable for responses

Without this change dnsdist removed the EDNS payload later in the
processing of the response.
---
 pdns/dnsdistdist/dnsdist-actions-factory.cc |  6 +++---
 pdns/dnsdistdist/dnsdist-ecs.cc             | 11 +++++++----
 pdns/dnsdistdist/dnsdist-ecs.hh             |  2 +-
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/pdns/dnsdistdist/dnsdist-actions-factory.cc b/pdns/dnsdistdist/dnsdist-actions-factory.cc
index c4806ff7d6..ec14364129 100644
--- a/pdns/dnsdistdist/dnsdist-actions-factory.cc
+++ b/pdns/dnsdistdist/dnsdist-actions-factory.cc
@@ -1018,7 +1018,7 @@ public:
   DNSAction::Action operator()(DNSQuestion* dnsquestion, std::string* ruleresult) const override
   {
     (void)ruleresult;
-    setEDNSOption(*dnsquestion, d_code, d_data);
+    setEDNSOption(*dnsquestion, d_code, d_data, true);
     return Action::None;
   }
 
@@ -1041,9 +1041,9 @@ public:
   {
   }
 
-  DNSResponseAction::Action operator()(DNSResponse* response, std::string* ruleresult) const override
+  DNSResponseAction::Action operator()(DNSResponse* response, [[maybe_unused]] std::string* ruleresult) const override
   {
-    setEDNSOption(*response, d_code, d_data);
+    setEDNSOption(*response, d_code, d_data, false);
     return Action::None;
   }
 
diff --git a/pdns/dnsdistdist/dnsdist-ecs.cc b/pdns/dnsdistdist/dnsdist-ecs.cc
index dc6b31cb88..c8cc4f6585 100644
--- a/pdns/dnsdistdist/dnsdist-ecs.cc
+++ b/pdns/dnsdistdist/dnsdist-ecs.cc
@@ -1163,7 +1163,7 @@ bool getEDNS0Record(const PacketBuffer& packet, EDNS0Record& edns0)
   return true;
 }
 
-bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::string& ednsData)
+bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::string& ednsData, bool isQuery)
 {
   std::string optRData;
   generateEDNSOption(ednsCode, ednsData, optRData);
@@ -1183,7 +1183,7 @@ bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::strin
     }
 
     dnsQuestion.getMutableData() = std::move(newContent);
-    if (!dnsQuestion.ids.ednsAdded && ednsAdded) {
+    if (isQuery && !dnsQuestion.ids.ednsAdded && ednsAdded) {
       dnsQuestion.ids.ednsAdded = true;
     }
 
@@ -1196,8 +1196,11 @@ bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::strin
       header.arcount = htons(1);
       return true;
     });
-    // make sure that any EDNS sent by the backend is removed before forwarding the response to the client
-    dnsQuestion.ids.ednsAdded = true;
+
+    if (isQuery) {
+      // make sure that any EDNS sent by the backend is removed before forwarding the response to the client
+      dnsQuestion.ids.ednsAdded = true;
+    }
   }
 
   return true;
diff --git a/pdns/dnsdistdist/dnsdist-ecs.hh b/pdns/dnsdistdist/dnsdist-ecs.hh
index 0c6a4780ed..e48da6d677 100644
--- a/pdns/dnsdistdist/dnsdist-ecs.hh
+++ b/pdns/dnsdistdist/dnsdist-ecs.hh
@@ -51,7 +51,7 @@ bool parseEDNSOptions(const DNSQuestion& dnsQuestion);
 bool queryHasEDNS(const DNSQuestion& dnsQuestion);
 bool getEDNS0Record(const PacketBuffer& packet, EDNS0Record& edns0);
 
-bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::string& data);
+bool setEDNSOption(DNSQuestion& dnsQuestion, uint16_t ednsCode, const std::string& data, bool isQuery = true);
 
 struct InternalQueryState;
 namespace dnsdist
-- 
2.47.2