From 8add550721c4f16671e60dca13950c01129f109c Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 12 Jan 2021 15:59:26 +0100
Subject: [PATCH] dnsdist: Add SkipCacheResponseAction

---
 pdns/dnsdist-console.cc                  |  1 +
 pdns/dnsdist-lua-actions.cc              | 18 +++++++++++++
 pdns/dnsdistdist/docs/rules-actions.rst  |  6 +++++
 regression-tests.dnsdist/test_Caching.py | 33 ++++++++++++++++++++++++
 4 files changed, 58 insertions(+)

diff --git a/pdns/dnsdist-console.cc b/pdns/dnsdist-console.cc
index c2ac58bed0..232c46b099 100644
--- a/pdns/dnsdist-console.cc
+++ b/pdns/dnsdist-console.cc
@@ -606,6 +606,7 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "shutdown", true, "", "shut down `dnsdist`" },
   { "SetProxyProtocolValuesAction", true, "values", "Set the Proxy-Protocol values for this queries to 'values'" },
   { "SkipCacheAction", true, "", "Don’t lookup the cache for this query, don’t store the answer" },
+  { "SkipCacheResponseAction", true, "", "Don’t store this response into the cache" },
   { "SNIRule", true, "name", "Create a rule which matches on the incoming TLS SNI value, if any (DoT or DoH)" },
   { "snmpAgent", true, "enableTraps [, masterSocket]", "enable `SNMP` support. `enableTraps` is a boolean indicating whether traps should be sent and `masterSocket` an optional string specifying how to connect to the master agent"},
   { "SNMPTrapAction", true, "[reason]", "send an SNMP trap, adding the optional `reason` string as the query description"},
diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc
index 45fa666337..4ae206fc72 100644
--- a/pdns/dnsdist-lua-actions.cc
+++ b/pdns/dnsdist-lua-actions.cc
@@ -856,6 +856,20 @@ public:
   }
 };
 
+class SkipCacheResponseAction : public DNSResponseAction
+{
+public:
+  DNSResponseAction::Action operator()(DNSResponse* dr, std::string* ruleresult) const override
+  {
+    dr->skipCache = true;
+    return Action::None;
+  }
+  std::string toString() const override
+  {
+    return "skip cache";
+  }
+};
+
 class TempFailureCacheTTLAction : public DNSAction
 {
 public:
@@ -1704,6 +1718,10 @@ void setupLuaActions(LuaContext& luaCtx)
       return std::shared_ptr<DNSAction>(new SkipCacheAction);
     });
 
+  luaCtx.writeFunction("SkipCacheResponseAction", []() {
+      return std::shared_ptr<DNSResponseAction>(new SkipCacheResponseAction);
+    });
+
   luaCtx.writeFunction("TempFailureCacheTTLAction", [](int maxTTL) {
       return std::shared_ptr<DNSAction>(new TempFailureCacheTTLAction(maxTTL));
     });
diff --git a/pdns/dnsdistdist/docs/rules-actions.rst b/pdns/dnsdistdist/docs/rules-actions.rst
index 6e6679cf5f..930979e057 100644
--- a/pdns/dnsdistdist/docs/rules-actions.rst
+++ b/pdns/dnsdistdist/docs/rules-actions.rst
@@ -1359,6 +1359,12 @@ The following actions exist.
 
   Don't lookup the cache for this query, don't store the answer.
 
+.. function:: SkipCacheResponseAction()
+
+  .. versionadded:: 1.6.0
+
+  Don't store this answer into the cache.
+
 .. function:: SNMPTrapAction([message])
 
   Send an SNMP trap, adding the optional ``message`` string as the query description.
diff --git a/regression-tests.dnsdist/test_Caching.py b/regression-tests.dnsdist/test_Caching.py
index d3f91753c2..be9927d4c1 100644
--- a/regression-tests.dnsdist/test_Caching.py
+++ b/regression-tests.dnsdist/test_Caching.py
@@ -12,6 +12,7 @@ class TestCaching(DNSDistTest):
     pc = newPacketCache(100, {maxTTL=86400, minTTL=1})
     getPool(""):setCache(pc)
     addAction(makeRule("nocache.cache.tests.powerdns.com."), SkipCacheAction())
+    addResponseAction(makeRule("nocache-response.cache.tests.powerdns.com."), SkipCacheResponseAction())
     function skipViaLua(dq)
         dq.skipCache = true
         return DNSAction.None, ""
@@ -200,6 +201,38 @@ class TestCaching(DNSDistTest):
             value = self._responsesCounter[key]
             self.assertEquals(value, numberOfQueries)
 
+    def testSkipCacheResponse(self):
+        """
+        Cache: SkipCacheResponseAction
+
+        dnsdist is configured to not cache entries for answer matching nocache-response.cache.tests.powerdns.com.
+         we are sending several requests and checking that the backend get them all.
+        """
+        name = 'nocache-response.cache.tests.powerdns.com.'
+        numberOfQueries = 10
+        query = dns.message.make_query(name, 'AAAA', 'IN')
+        response = dns.message.make_response(query)
+        rrset = dns.rrset.from_text(name,
+                                    3600,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.AAAA,
+                                    '::1')
+        response.answer.append(rrset)
+
+        for _ in range(numberOfQueries):
+            for method in ("sendUDPQuery", "sendTCPQuery"):
+                sender = getattr(self, method)
+                (receivedQuery, receivedResponse) = sender(query, response)
+                self.assertTrue(receivedQuery)
+                self.assertTrue(receivedResponse)
+                receivedQuery.id = query.id
+                self.assertEquals(query, receivedQuery)
+                self.assertEquals(receivedResponse, response)
+
+        for key in self._responsesCounter:
+            value = self._responsesCounter[key]
+            self.assertEquals(value, numberOfQueries)
+
     def testCacheExpiration(self):
         """
         Cache: Cache expiration
-- 
2.47.2