From 1ecf6f40306270263de5c781ed6b357f2bff7311 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 | 141 ++++++++++++++++++++++++++-------------------
 1 file changed, 81 insertions(+), 60 deletions(-)

diff --git a/agents/dnsagent.go b/agents/dnsagent.go
index b21aec6a1..2ae6d1268 100644
--- a/agents/dnsagent.go
+++ b/agents/dnsagent.go
@@ -84,70 +84,23 @@ func (da *DNSAgent) Reload() (err error) {
 // requests are reaching here asynchronously
 func (da *DNSAgent) handleMessage(w dns.ResponseWriter, req *dns.Msg) {
 	dnsDP := newDnsDP(req)
-	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()),
-		},
-	}
+
 	rply := new(dns.Msg)
 	rply.SetReply(req)
-	// message preprocesing
-	if req.Question[0].Qtype == dns.TypeNAPTR {
-		reqVars.Map[QueryName] = utils.NewLeafNode(req.Question[0].Name)
-	}
-	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(
-			reqProcessor,
-			NewAgentRequest(
-				dnsDP, reqVars, cgrRplyNM, rplyNM,
-				opts, reqProcessor.Tenant,
-				da.cgrCfg.GeneralCfg().DefaultTenant,
-				utils.FirstNonEmpty(da.cgrCfg.DNSAgentCfg().Timezone,
-					da.cgrCfg.GeneralCfg().DefaultTimezone),
-				da.fltrS, nil),
-			utils.DNSAgent, da.connMgr,
-			da.cgrCfg.DNSAgentCfg().SessionSConns,
-			nil, da.fltrS)
-		if lclProcessed {
-			processed = lclProcessed
-		}
-		if err != nil ||
-			(lclProcessed && !reqProcessor.Flags.GetBool(utils.MetaContinue)) {
-			break
+
+	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 != 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 {
-		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)
-		return
-	}
-	if err = updateDNSMsgFromNM(rply, rplyNM); 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
+
+	if err := dnsWriteMsg(w, rply); err != nil { // failed sending, most probably content issue
 		rply = new(dns.Msg)
 		rply.SetReply(req)
 		rply.Rcode = dns.RcodeServerFailure
@@ -159,3 +112,71 @@ func (da *DNSAgent) handleMessage(w dns.ResponseWriter, req *dns.Msg) {
 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[q.Qtype]),
+			QueryName:        utils.NewLeafNode(q.Name),
+			utils.RemoteHost: utils.NewLeafNode(rmtAddr),
+		},
+	}
+	// message preprocesing
+	if q.Qtype == dns.TypeNAPTR {
+		/*
+			e164, err := e164FromNAPTR(req.Question[0].Name)
+			if err != nil {
+				utils.Logger.Warning(
+					fmt.Sprintf("<%s> decoding NAPTR query: <%s>, err: %s",
+						utils.DNSAgent, req.Question[0].Name, err.Error()))
+				rply.Rcode = dns.RcodeServerFailure
+				dnsWriteMsg(w, rply)
+				return
+			}
+			reqVars.Map[E164Address] = utils.NewLeafNode(e164)
+			reqVars.Map[DomainName] = utils.NewLeafNode(domainNameFromNAPTR(req.Question[0].Name))
+		*/
+	}
+	cgrRplyNM := &utils.DataNode{Type: utils.NMMapType, Map: make(map[string]*utils.DataNode)}
+	rplyNM := utils.NewOrderedNavigableMap() // share it among different processors
+	opts := utils.MapStorage{}
+	for _, reqProcessor := range da.cgrCfg.DNSAgentCfg().RequestProcessors {
+		var lclProcessed bool
+		if lclProcessed, err = processRequest(
+			reqProcessor,
+			NewAgentRequest(
+				dnsDP, reqVars, cgrRplyNM, rplyNM,
+				opts, reqProcessor.Tenant,
+				da.cgrCfg.GeneralCfg().DefaultTenant,
+				utils.FirstNonEmpty(da.cgrCfg.DNSAgentCfg().Timezone,
+					da.cgrCfg.GeneralCfg().DefaultTimezone),
+				da.fltrS, nil),
+			utils.DNSAgent, da.connMgr,
+			da.cgrCfg.DNSAgentCfg().SessionSConns,
+			nil, da.fltrS); err != nil {
+			utils.Logger.Warning(
+				fmt.Sprintf("<%s> error: %s processing message: %s from %s",
+					utils.DNSAgent, err.Error(), dnsDP, rmtAddr))
+			return
+		}
+		processed = processed || lclProcessed
+		if lclProcessed && !reqProcessor.Flags.GetBool(utils.MetaContinue) {
+			break
+		}
+	}
+	if !processed {
+		utils.Logger.Warning(
+			fmt.Sprintf("<%s> no request processor enabled, ignoring message %s from %s",
+				utils.DNSAgent, dnsDP, rmtAddr))
+		return
+	}
+	if err = updateDNSMsgFromNM(rply, rplyNM); 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)))
+	}
+	return
+}