From 295167c5f0ff80917cb3bc5f51e56b92d95e9e6f Mon Sep 17 00:00:00 2001
From: Trial97 <alexandru.tripon@itsyscom.com>
Date: Fri, 1 Oct 2021 17:10:51 +0300
Subject: [PATCH] Updated dns agent

---
 agents/dnsagent.go | 82 +++++++++++++++++++++++++---------------------
 1 file changed, 45 insertions(+), 37 deletions(-)

diff --git a/agents/dnsagent.go b/agents/dnsagent.go
index a4eb96437..2cd453d16 100644
--- a/agents/dnsagent.go
+++ b/agents/dnsagent.go
@@ -85,24 +85,53 @@ func (da *DNSAgent) Reload() (err error) {
 // requests are reaching here asynchronously
 func (da *DNSAgent) handleMessage(w dns.ResponseWriter, req *dns.Msg) {
 	dnsDP := newDnsDP(req)
+
+	rply := new(dns.Msg)
+	rply.SetReply(req)
+
+	rmtAddr := w.RemoteAddr().String()
+	for _, q := range req.Question {
+		if processed, err := da.handleQuestion(dnsDP, rply, &q, rmtAddr); err != nil ||
+			!processed {
+			rply.Rcode = dns.RcodeServerFailure
+			rply = new(dns.Msg)
+			rply.SetReply(req)
+			dnsWriteMsg(w, rply)
+			return
+		}
+	}
+
+	if err := dnsWriteMsg(w, rply); err != nil { // failed sending, most probably content issue
+		rply = new(dns.Msg)
+		rply.SetReply(req)
+		rply.Rcode = dns.RcodeServerFailure
+		dnsWriteMsg(w, rply)
+	}
+}
+
+// Shutdown stops the DNS server
+func (da *DNSAgent) Shutdown() error {
+	return da.server.Shutdown()
+}
+
+// handleMessage is the entry point of all DNS requests
+// requests are reaching here asynchronously
+func (da *DNSAgent) handleQuestion(dnsDP utils.DataProvider, rply *dns.Msg, q *dns.Question, rmtAddr string) (processed bool, err error) {
 	reqVars := &utils.DataNode{
 		Type: utils.NMMapType,
 		Map: map[string]*utils.DataNode{
-			QueryType:        utils.NewLeafNode(dns.TypeToString[req.Question[0].Qtype]),
-			utils.RemoteHost: utils.NewLeafNode(w.RemoteAddr().String()),
+			QueryType:        utils.NewLeafNode(dns.TypeToString[q.Qtype]),
+			QueryName:        utils.NewLeafNode(q.Name),
+			utils.RemoteHost: utils.NewLeafNode(rmtAddr),
 		},
 	}
-	rply := new(dns.Msg)
-	rply.SetReply(req)
 	// message preprocesing
 	cgrRplyNM := &utils.DataNode{Type: utils.NMMapType, Map: make(map[string]*utils.DataNode)}
 	rplyNM := utils.NewOrderedNavigableMap() // share it among different processors
 	opts := utils.MapStorage{}
-	var processed bool
-	var err error
 	for _, reqProcessor := range da.cgrCfg.DNSAgentCfg().RequestProcessors {
 		var lclProcessed bool
-		lclProcessed, err = processRequest(
+		if lclProcessed, err = processRequest(
 			context.TODO(),
 			reqProcessor,
 			NewAgentRequest(
@@ -114,48 +143,27 @@ func (da *DNSAgent) handleMessage(w dns.ResponseWriter, req *dns.Msg) {
 				da.fltrS, nil),
 			utils.DNSAgent, da.connMgr,
 			da.cgrCfg.DNSAgentCfg().SessionSConns,
-			da.fltrS)
-		if lclProcessed {
-			processed = lclProcessed
+			da.fltrS); err != nil {
+			utils.Logger.Warning(
+				fmt.Sprintf("<%s> error: %s processing message: %s from %s",
+					utils.DNSAgent, err.Error(), dnsDP, rmtAddr))
+			return
 		}
-		if err != nil ||
-			(lclProcessed && !reqProcessor.Flags.GetBool(utils.MetaContinue)) {
+		processed = processed || lclProcessed
+		if lclProcessed && !reqProcessor.Flags.GetBool(utils.MetaContinue) {
 			break
 		}
 	}
-	if err != nil {
-		utils.Logger.Warning(
-			fmt.Sprintf("<%s> error: %s processing message: %s from %s",
-				utils.DNSAgent, err.Error(), req, w.RemoteAddr()))
-		rply.Rcode = dns.RcodeServerFailure
-		dnsWriteMsg(w, rply)
-		return
-	} else if !processed {
+	if !processed {
 		utils.Logger.Warning(
 			fmt.Sprintf("<%s> no request processor enabled, ignoring message %s from %s",
-				utils.DNSAgent, req, w.RemoteAddr()))
-		rply.Rcode = dns.RcodeServerFailure
-		dnsWriteMsg(w, rply)
+				utils.DNSAgent, dnsDP, rmtAddr))
 		return
 	}
 	if err = updateDNSMsgFromNM(rply, rplyNM, q.Qtype, q.Name); err != nil {
 		utils.Logger.Warning(
 			fmt.Sprintf("<%s> error: %s updating answer: %s from NM %s",
 				utils.DNSAgent, err.Error(), utils.ToJSON(rply), utils.ToJSON(rplyNM)))
-		rply.Rcode = dns.RcodeServerFailure
-		dnsWriteMsg(w, rply)
-		return
-	}
-	if err = dnsWriteMsg(w, rply); err != nil { // failed sending, most probably content issue
-		rply = new(dns.Msg)
-		rply.SetReply(req)
-		rply.Rcode = dns.RcodeServerFailure
-		dnsWriteMsg(w, rply)
 	}
 	return
 }
-
-// Shutdown stops the DNS server
-func (da *DNSAgent) Shutdown() error {
-	return da.server.Shutdown()
-}