From c1ea322ab687c5e92bc3ff39506f43978aa5402b Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 6 Apr 2021 15:06:34 +0200
Subject: [PATCH] dnsdist: Add missing getEDNSOptions and getDO bindings for
 DNSResponse

(cherry picked from commit e292b91b89827450ca116a8fe7dbca146705ba3b)
---
 pdns/dnsdist-lua-bindings-dnsquestion.cc     | 10 ++++++
 regression-tests.dnsdist/test_EDNSOptions.py | 38 ++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/pdns/dnsdist-lua-bindings-dnsquestion.cc b/pdns/dnsdist-lua-bindings-dnsquestion.cc
index 9c87d129c8..8d194a6c30 100644
--- a/pdns/dnsdist-lua-bindings-dnsquestion.cc
+++ b/pdns/dnsdist-lua-bindings-dnsquestion.cc
@@ -146,6 +146,16 @@ void setupLuaBindingsDNSQuestion()
   g_lua.registerFunction<void(DNSResponse::*)(std::function<uint32_t(uint8_t section, uint16_t qclass, uint16_t qtype, uint32_t ttl)> editFunc)>("editTTLs", [](const DNSResponse& dr, std::function<uint32_t(uint8_t section, uint16_t qclass, uint16_t qtype, uint32_t ttl)> editFunc) {
         editDNSPacketTTL((char*) dr.dh, dr.len, editFunc);
       });
+  g_lua.registerFunction<bool(DNSResponse::*)()const>("getDO", [](const DNSResponse& dq) {
+      return getEDNSZ(dq) & EDNS_HEADER_FLAG_DO;
+    });
+  g_lua.registerFunction<std::map<uint16_t, EDNSOptionView>(DNSResponse::*)()>("getEDNSOptions", [](DNSResponse& dq) {
+      if (dq.ednsOptions == nullptr) {
+        parseEDNSOptions(dq);
+      }
+
+      return *dq.ednsOptions;
+    });
   g_lua.registerFunction<std::string(DNSResponse::*)(void)>("getTrailingData", [](const DNSResponse& dq) {
       return dq.getTrailingData();
     });
diff --git a/regression-tests.dnsdist/test_EDNSOptions.py b/regression-tests.dnsdist/test_EDNSOptions.py
index f1e386e825..3f00834f24 100644
--- a/regression-tests.dnsdist/test_EDNSOptions.py
+++ b/regression-tests.dnsdist/test_EDNSOptions.py
@@ -58,6 +58,43 @@ class EDNSOptionsBase(DNSDistTest):
 
       return DNSAction.None, ""
 
+    end
+
+    function testEDNSOptionsInResponses(dr)
+      local options = dr:getEDNSOptions()
+      local qname = tostring(dr.qname)
+
+      if string.match(qname, 'multiplecookies') then
+        return DNSAction.None, ""
+      elseif string.match(qname, 'cookie') then
+        if options[EDNSOptionCode.COOKIE] == nil then
+          return DNSAction.Spoof, "192.0.2.1"
+        end
+        if options[EDNSOptionCode.COOKIE]:count() ~= 1 or options[EDNSOptionCode.COOKIE]:getValues()[1]:len() ~= 16 then
+          return DNSAction.Spoof, "192.0.2.2"
+        end
+      end
+
+      if string.match(qname, 'ecs4') then
+        if options[EDNSOptionCode.ECS] == nil then
+          return DNSAction.Spoof, "192.0.2.51"
+        end
+        if options[EDNSOptionCode.ECS]:count() ~= 1 or options[EDNSOptionCode.ECS]:getValues()[1]:len() ~= 7 then
+          return DNSAction.Spoof, "192.0.2.52"
+        end
+      end
+
+      if string.match(qname, 'ecs6') then
+        if options[EDNSOptionCode.ECS] == nil then
+          return DNSAction.Spoof, "192.0.2.101"
+        end
+        if options[EDNSOptionCode.ECS]:count() ~= 1 or options[EDNSOptionCode.ECS]:getValues()[1]:len() ~= 11 then
+          return DNSAction.Spoof, "192.0.2.102"
+        end
+      end
+
+      return DNSAction.None, ""
+
     end
     """
 
@@ -224,6 +261,7 @@ class TestEDNSOptionsAddingECS(EDNSOptionsBase):
     %s
 
     addAction(AllRule(), LuaAction(testEDNSOptions))
+    addResponseAction("ednsoptions-ecs.tests.powerdns.com.", LuaResponseAction(testEDNSOptionsInResponses))
 
     newServer{address="127.0.0.1:%s", useClientSubnet=true}
     """
-- 
2.47.2