From 3cea1600115aa73f7599b95dba86915f0c023707 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 8 Mar 2024 17:52:54 +0100
Subject: [PATCH] dnsdist: Add the ability to set tags from dynamic block rules

This still needs:
- ability to set the tag name and value from the Dynamic Block configuration
- tests
- documentation
---
 pdns/dnsdistdist/dnsdist.cc | 28 ++++++++++++++++++++++++++++
 pdns/dnsdistdist/dnsdist.hh | 18 ++++++++++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/pdns/dnsdistdist/dnsdist.cc b/pdns/dnsdistdist/dnsdist.cc
index d411b0b66c..2942319109 100644
--- a/pdns/dnsdistdist/dnsdist.cc
+++ b/pdns/dnsdistdist/dnsdist.cc
@@ -1022,6 +1022,8 @@ bool processRulesResult(const DNSAction::Action& action, DNSQuestion& dnsQuestio
   case DNSAction::Action::Delay:
     pdns::checked_stoi_into(dnsQuestion.ids.delayMsec, ruleresult); // sorry
     break;
+  case DNSAction::Action::SetTag:
+    /* unsupported for non-dynamic block */
   case DNSAction::Action::None:
     /* fall-through */
   case DNSAction::Action::NoOp:
@@ -1143,6 +1145,19 @@ static bool applyRulesToQuery(LocalHolders& holders, DNSQuestion& dnsQuestion, c
           return true;
         });
         return true;
+      case DNSAction::Action::SetTag:
+      {
+        if (!got->second.tagSettings) {
+          vinfolog("Skipping set tag dynamic block for query from %s because of missing options", dnsQuestion.ids.origRemote.toStringWithPort());
+          break;
+        }
+        updateBlockStats();
+        const auto& tagName = got->second.tagSettings->d_name;
+        const auto& tagValue = got->second.tagSettings->d_value;
+        dnsQuestion.setTag(tagName, tagValue);
+        vinfolog("Query from %s setting tag %s to %s because of dynamic block", dnsQuestion.ids.origRemote.toStringWithPort(), tagName, tagValue);
+        return true;
+      }
       default:
         updateBlockStats();
         vinfolog("Query from %s dropped because of dynamic block", dnsQuestion.ids.origRemote.toStringWithPort());
@@ -1205,6 +1220,19 @@ static bool applyRulesToQuery(LocalHolders& holders, DNSQuestion& dnsQuestion, c
           return true;
         });
         return true;
+      case DNSAction::Action::SetTag:
+      {
+        if (!got->tagSettings) {
+          vinfolog("Skipping set tag dynamic block for query from %s because of missing options", dnsQuestion.ids.origRemote.toStringWithPort());
+          break;
+        }
+        updateBlockStats();
+        const auto& tagName = got->tagSettings->d_name;
+        const auto& tagValue = got->tagSettings->d_value;
+        dnsQuestion.setTag(tagName, tagValue);
+        vinfolog("Query from %s setting tag %s to %s because of dynamic block", dnsQuestion.ids.origRemote.toStringWithPort(), tagName, tagValue);
+        return true;
+      }
       default:
         updateBlockStats();
         vinfolog("Query from %s for %s dropped because of dynamic block", dnsQuestion.ids.origRemote.toStringWithPort(), dnsQuestion.ids.qname.toLogString());
diff --git a/pdns/dnsdistdist/dnsdist.hh b/pdns/dnsdistdist/dnsdist.hh
index 12dca5706d..76444e6c64 100644
--- a/pdns/dnsdistdist/dnsdist.hh
+++ b/pdns/dnsdistdist/dnsdist.hh
@@ -239,7 +239,8 @@ public:
     NoOp,
     NoRecurse,
     SpoofRaw,
-    SpoofPacket
+    SpoofPacket,
+    SetTag,
   };
   static std::string typeToString(const Action& action)
   {
@@ -268,6 +269,8 @@ public:
       return "Truncate over UDP";
     case Action::ServFail:
       return "Send ServFail";
+    case Action::SetTag:
+      return "Set Tag";
     case Action::None:
     case Action::NoOp:
       return "Do nothing";
@@ -349,6 +352,9 @@ struct DynBlock
     blocks.store(rhs.blocks);
     warning = rhs.warning;
     bpf = rhs.bpf;
+    if (rhs.tagSettings != nullptr) {
+      tagSettings = std::make_unique<TagSettings>(*rhs.tagSettings);
+    }
     return *this;
   }
 
@@ -361,13 +367,21 @@ struct DynBlock
     blocks.store(rhs.blocks);
     warning = rhs.warning;
     bpf = rhs.bpf;
+    tagSettings = std::move(rhs.tagSettings);
     return *this;
   }
 
+  struct TagSettings
+  {
+    std::string d_name;
+    std::string d_value;
+  };
+
   string reason;
   DNSName domain;
   timespec until{};
-  mutable std::atomic<unsigned int> blocks{0};
+  std::unique_ptr<TagSettings> tagSettings{nullptr};
+  mutable std::atomic<uint32_t> blocks{0};
   DNSAction::Action action{DNSAction::Action::None};
   bool warning{false};
   bool bpf{false};
-- 
2.47.2