From a860259d8cc971d5bfb2d50c0d341e9f45cf5cf1 Mon Sep 17 00:00:00 2001
From: Artem Glazychev <artem.glazychev@xored.com>
Date: Thu, 8 Sep 2022 00:12:51 +0700
Subject: [PATCH] Increase DNS timeout

Signed-off-by: Artem Glazychev <artem.glazychev@xored.com>
---
 pkg/tools/dnsutils/dnsutils.go       |  2 +-
 pkg/tools/dnsutils/memory/handler.go | 35 +++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/pkg/tools/dnsutils/dnsutils.go b/pkg/tools/dnsutils/dnsutils.go
index ca1d3a3c8..ef7cb6fbb 100644
--- a/pkg/tools/dnsutils/dnsutils.go
+++ b/pkg/tools/dnsutils/dnsutils.go
@@ -37,7 +37,7 @@ func ListenAndServe(ctx context.Context, handler Handler, listenOn string) {
 
 	for _, network := range networks {
 		var server = &dns.Server{Addr: listenOn, Net: network, Handler: dns.HandlerFunc(func(w dns.ResponseWriter, m *dns.Msg) {
-			var timeoutCtx, cancel = context.WithTimeout(context.Background(), time.Second)
+			var timeoutCtx, cancel = context.WithTimeout(context.Background(), time.Second*5)
 			defer cancel()
 
 			handler.ServeDNS(timeoutCtx, w, m)
diff --git a/pkg/tools/dnsutils/memory/handler.go b/pkg/tools/dnsutils/memory/handler.go
index 35dabea38..5aa042219 100644
--- a/pkg/tools/dnsutils/memory/handler.go
+++ b/pkg/tools/dnsutils/memory/handler.go
@@ -29,21 +29,41 @@ import (
 
 const defaultTTL = 3600
 
+// Since memory is supposed to be one of the targets that stores information, we have to keep track of whether something has been written to the writer.
+// We must write something into the writer, this is how the dns package works. Otherwise, we get a timeout error on the client side.
+
+type responseWriter struct {
+	dns.ResponseWriter
+	passed bool
+}
+
+func (r *responseWriter) WriteMsg(m *dns.Msg) error {
+	r.passed = true
+	return r.ResponseWriter.WriteMsg(m)
+}
+
 type memoryHandler struct {
-	recoreds *Map
+	records *Map
 }
 
 func (f *memoryHandler) ServeDNS(ctx context.Context, rw dns.ResponseWriter, msg *dns.Msg) {
 	if len(msg.Question) == 0 {
-		next.Handler(ctx).ServeDNS(ctx, rw, msg)
+		dns.HandleFailed(rw, msg)
 		return
 	}
 
+	rwWrapper := &responseWriter{
+		ResponseWriter: rw,
+	}
+
 	var name = dns.Name(msg.Question[0].Name).String()
-	var records, ok = f.recoreds.Load(name)
+	var records, ok = f.records.Load(name)
 
 	if !ok {
-		next.Handler(ctx).ServeDNS(ctx, rw, msg)
+		next.Handler(ctx).ServeDNS(ctx, rwWrapper, msg)
+		if !rwWrapper.passed {
+			dns.HandleFailed(rw, msg)
+		}
 		return
 	}
 
@@ -59,7 +79,10 @@ func (f *memoryHandler) ServeDNS(ctx context.Context, rw dns.ResponseWriter, msg
 	}
 
 	if len(resp.Answer) == 0 {
-		next.Handler(ctx).ServeDNS(ctx, rw, msg)
+		next.Handler(ctx).ServeDNS(ctx, rwWrapper, msg)
+		if !rwWrapper.passed {
+			dns.HandleFailed(rw, msg)
+		}
 		return
 	}
 
@@ -73,7 +96,7 @@ func NewDNSHandler(records *Map) dnsutils.Handler {
 	if records == nil {
 		panic("records cannot be nil")
 	}
-	return &memoryHandler{recoreds: records}
+	return &memoryHandler{records: records}
 }
 func a(domain string, ips []net.IP) []dns.RR {
 	answers := make([]dns.RR, len(ips))