From: Cmdr Riker <riker@ted.com>
Date: Sun, 16 Jun 2019 20:21:30 +0000 (-0500)
Subject: dnsdist: Add support for custom DoH headers
X-Git-Tag: dnsdist-1.4.0-rc1~16^2~1
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=ee01507fbb337e2030493861240620b268e79560;p=pdns

dnsdist: Add support for custom DoH headers
---

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index fabbe5132..6a0e5211d 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -87,7 +87,7 @@ void resetLuaSideEffect()
   g_noLuaSideEffect = boost::logic::indeterminate;
 }
 
-typedef std::unordered_map<std::string, boost::variant<bool, int, std::string, std::vector<std::pair<int,int> > > > localbind_t;
+typedef std::unordered_map<std::string, boost::variant<bool, int, std::string, std::vector<std::pair<int,int> >, std::map<std::string,std::string> > > localbind_t;
 
 static void parseLocalBindVars(boost::optional<localbind_t> vars, bool& reusePort, int& tcpFastOpenQueueSize, std::string& interface, std::set<int>& cpus)
 {
@@ -1726,6 +1726,12 @@ void setupLuaConfig(bool client)
       if (vars->count("serverTokens")) {
         frontend->d_serverTokens = boost::get<const string>((*vars)["serverTokens"]);
       }
+      if (vars->count("customResponseHeaders")) {
+        for (auto const& headerMap : boost::get<std::map<std::string,std::string>>((*vars)["customResponseHeaders"])) {
+          std::pair<std::string,std::string> headerResponse = std::make_pair(headerMap.first, headerMap.second);
+          frontend->d_customResponseHeaders.push_back(headerResponse);
+        }
+      }
     }
     g_dohlocals.push_back(frontend);
     auto cs = std::unique_ptr<ClientState>(new ClientState(frontend->d_local, true, reusePort, tcpFastOpenQueueSize, interface, cpus));
diff --git a/pdns/dnsdistdist/doh.cc b/pdns/dnsdistdist/doh.cc
index 55ed64059..f932c49f3 100644
--- a/pdns/dnsdistdist/doh.cc
+++ b/pdns/dnsdistdist/doh.cc
@@ -445,6 +445,12 @@ try
     return 0;
   }
 
+  constexpr int overwrite_if_exists = 1;
+  constexpr int maybe_token = 1;
+  for (auto const& headerPair : dsc->df->d_customResponseHeaders) {
+    h2o_set_header_by_str(&req->pool, &req->res.headers, headerPair.first.c_str(), headerPair.first.size(), maybe_token, headerPair.second.c_str(), headerPair.second.size(), overwrite_if_exists);
+  }
+
   if(auto tlsversion = h2o_socket_get_ssl_protocol_version(sock)) {
     if(!strcmp(tlsversion, "TLSv1.0"))
       ++dsc->df->d_tls10queries;
diff --git a/pdns/doh.hh b/pdns/doh.hh
index 4e1acb0b5..b7b9bfe89 100644
--- a/pdns/doh.hh
+++ b/pdns/doh.hh
@@ -10,6 +10,7 @@ struct DOHFrontend
   std::string d_ciphers;
   std::string d_ciphers13;
   std::string d_serverTokens{"h2o/dnsdist"};
+  std::vector<std::pair<std::string, std::string>> d_customResponseHeaders;
   ComboAddress d_local;
 
   uint32_t d_idleTimeout{30};             // HTTP idle timeout in seconds
diff --git a/regression-tests.dnsdist/test_DOH.py b/regression-tests.dnsdist/test_DOH.py
index 4943c96ec..75586b92e 100644
--- a/regression-tests.dnsdist/test_DOH.py
+++ b/regression-tests.dnsdist/test_DOH.py
@@ -5,7 +5,7 @@ import clientsubnetoption
 from dnsdisttests import DNSDistTest
 
 import pycurl
-
+from io import StringIO
 #from hyper import HTTP20Connection
 #from hyper.ssl_compat import SSLContext, PROTOCOL_TLSv1_2
 
@@ -33,12 +33,14 @@ class DNSDistDOHTest(DNSDistTest):
     def sendDOHQuery(cls, port, servername, baseurl, query, response=None, timeout=2.0, caFile=None, useQueue=True, rawQuery=False, customHeaders=[]):
         url = cls.getDOHGetURL(baseurl, query, rawQuery)
         conn = cls.openDOHConnection(port, caFile=caFile, timeout=timeout)
+        response_headers = StringIO()
         #conn.setopt(pycurl.VERBOSE, True)
         conn.setopt(pycurl.URL, url)
         conn.setopt(pycurl.RESOLVE, ["%s:%d:127.0.0.1" % (servername, port)])
         conn.setopt(pycurl.SSL_VERIFYPEER, 1)
         conn.setopt(pycurl.SSL_VERIFYHOST, 2)
         conn.setopt(pycurl.HTTPHEADER, customHeaders)
+        conn.setopt(pycurl.HEADERFUNCTION, response_headers.write)
         if caFile:
             conn.setopt(pycurl.CAINFO, caFile)
 
@@ -47,6 +49,7 @@ class DNSDistDOHTest(DNSDistTest):
 
         receivedQuery = None
         message = None
+        cls._response_headers = ''
         data = conn.perform_rb()
         rcode = conn.getinfo(pycurl.RESPONSE_CODE)
         if rcode == 200:
@@ -55,6 +58,7 @@ class DNSDistDOHTest(DNSDistTest):
         if useQueue and not cls._fromResponderQueue.empty():
             receivedQuery = cls._fromResponderQueue.get(True, timeout)
 
+        cls._response_headers = response_headers.getvalue()
         return (receivedQuery, message)
 
     @classmethod
@@ -128,11 +132,13 @@ class TestDOH(DNSDistDOHTest):
     _serverName = 'tls.tests.dnsdist.org'
     _caCert = 'ca.pem'
     _dohServerPort = 8443
-    _serverName = 'tls.tests.dnsdist.org'
+    _customResponseHeader1 = 'access-control-allow-origin: *'
+    _customResponseHeader2 = 'user-agent: derp'
     _dohBaseURL = ("https://%s:%d/" % (_serverName, _dohServerPort))
     _config_template = """
     newServer{address="127.0.0.1:%s"}
-    addDOHLocal("127.0.0.1:%s", "%s", "%s", { "/" })
+
+    addDOHLocal("127.0.0.1:%s", "%s", "%s", { "/" }, {customResponseHeaders={["access-control-allow-origin"]="*",["user-agent"]="derp"}})
 
     addAction("drop.doh.tests.powerdns.com.", DropAction())
     addAction("refused.doh.tests.powerdns.com.", RCodeAction(DNSRCode.REFUSED))
@@ -164,6 +170,8 @@ class TestDOH(DNSDistDOHTest):
         self.assertTrue(receivedResponse)
         receivedQuery.id = expectedQuery.id
         self.assertEquals(expectedQuery, receivedQuery)
+        self.assertTrue((self._customResponseHeader1) in self._response_headers)
+        self.assertTrue((self._customResponseHeader2) in self._response_headers)
         self.checkQueryEDNSWithoutECS(expectedQuery, receivedQuery)
         self.assertEquals(response, receivedResponse)