From 7b5803482c4e06b9ff327daa2fcd2280a8871a70 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Tue, 29 Sep 2020 09:02:49 +0200
Subject: [PATCH] Support non-standard format on dns query name

---
 nameserver/nameserver.go | 24 +++++++++++++++++++-----
 resolver/rrcache.go      | 32 ++++++++++++++++++++------------
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/nameserver/nameserver.go b/nameserver/nameserver.go
index 21b81151..0d780d87 100644
--- a/nameserver/nameserver.go
+++ b/nameserver/nameserver.go
@@ -82,10 +82,19 @@ func handleRequestAsWorker(w dns.ResponseWriter, query *dns.Msg) {
 
 func handleRequest(ctx context.Context, w dns.ResponseWriter, request *dns.Msg) error { //nolint:gocognit // TODO
 	// Only process first question, that's how everyone does it.
-	question := request.Question[0]
+	originalQuestion := request.Question[0]
+
+	// Check if we are handling a non-standard query name.
+	var nonStandardQuestionFormat bool
+	lowerCaseQuestion := strings.ToLower(originalQuestion.Name)
+	if lowerCaseQuestion != originalQuestion.Name {
+		nonStandardQuestionFormat = true
+	}
+
+	// Create query for the resolver.
 	q := &resolver.Query{
-		FQDN:  question.Name,
-		QType: dns.Type(question.Qtype),
+		FQDN:  lowerCaseQuestion,
+		QType: dns.Type(originalQuestion.Qtype),
 	}
 
 	// Get remote address of request.
@@ -118,9 +127,9 @@ func handleRequest(ctx context.Context, w dns.ResponseWriter, request *dns.Msg)
 	}
 
 	// Check the Query Class.
-	if question.Qclass != dns.ClassINET {
+	if originalQuestion.Qclass != dns.ClassINET {
 		// we only serve IN records, return nxdomain
-		tracer.Warningf("nameserver: only IN record requests are supported but received Qclass %d, returning NXDOMAIN", question.Qclass)
+		tracer.Warningf("nameserver: only IN record requests are supported but received QClass %d, returning NXDOMAIN", originalQuestion.Qclass)
 		return reply(nsutil.Refused("unsupported qclass"))
 	}
 
@@ -243,6 +252,11 @@ func handleRequest(ctx context.Context, w dns.ResponseWriter, request *dns.Msg)
 	// Save dns request as open.
 	defer network.SaveOpenDNSRequest(conn)
 
+	// Revert back to non-standard question format, if we had to convert.
+	if nonStandardQuestionFormat {
+		rrCache.ReplaceAnswerNames(originalQuestion.Name)
+	}
+
 	// Reply with successful response.
 	tracer.Infof("nameserver: returning %s response for %s to %s", conn.Verdict.Verb(), q.ID(), conn.Process())
 	return reply(rrCache, conn, rrCache)
diff --git a/resolver/rrcache.go b/resolver/rrcache.go
index 9411e9db..5a277987 100644
--- a/resolver/rrcache.go
+++ b/resolver/rrcache.go
@@ -24,7 +24,7 @@ type RRCache struct {
 	Question dns.Type // constant
 	RCode    int      // constant
 
-	Answer []dns.RR // constant
+	Answer []dns.RR // mutable (mixing, non-standard formats)
 	Ns     []dns.RR // constant
 	Extra  []dns.RR // constant
 	TTL    int64    // constant
@@ -273,25 +273,33 @@ func (rrCache *RRCache) ShallowCopy() *RRCache {
 	}
 }
 
+// ReplaceAnswerNames is a helper function that replaces all answer names, that
+// match the query domain, with another value. This is used to support handling
+// non-standard query names, which are resolved normalized, but have to be
+// reverted back for the origin non-standard query name in order for the
+// clients to recognize the response.
+func (rrCache *RRCache) ReplaceAnswerNames(fqdn string) {
+	for _, answer := range rrCache.Answer {
+		if answer.Header().Name == rrCache.Domain {
+			answer.Header().Name = fqdn
+		}
+	}
+}
+
 // ReplyWithDNS creates a new reply to the given query with the data from the RRCache, and additional informational records.
 func (rrCache *RRCache) ReplyWithDNS(ctx context.Context, request *dns.Msg) *dns.Msg {
 	// reply to query
 	reply := new(dns.Msg)
 	reply.SetRcode(request, rrCache.RCode)
+	reply.Answer = rrCache.Answer
 	reply.Ns = rrCache.Ns
 	reply.Extra = rrCache.Extra
 
-	if len(rrCache.Answer) > 0 {
-		// Copy answers, as we randomize their order a little.
-		reply.Answer = make([]dns.RR, len(rrCache.Answer))
-		copy(reply.Answer, rrCache.Answer)
-
-		// Randomize the order of the answer records a little to allow dumb clients
-		// (who only look at the first record) to reliably connect.
-		for i := range reply.Answer {
-			j := rand.Intn(i + 1)
-			reply.Answer[i], reply.Answer[j] = reply.Answer[j], reply.Answer[i]
-		}
+	// Randomize the order of the answer records a little to allow dumb clients
+	// (who only look at the first record) to reliably connect.
+	for i := range reply.Answer {
+		j := rand.Intn(i + 1)
+		reply.Answer[i], reply.Answer[j] = reply.Answer[j], reply.Answer[i]
 	}
 
 	return reply