From 650c88db1ba210ec03bb386f83aeaa25dd550468 Mon Sep 17 00:00:00 2001
From: Alex Schittko <alex4108@live.com>
Date: Sat, 13 Jul 2024 08:00:00 +0000
Subject: [PATCH] Refactor

---
 tunneldns/tunnel.go | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/tunneldns/tunnel.go b/tunneldns/tunnel.go
index 05f72ee8..e991f9e3 100644
--- a/tunneldns/tunnel.go
+++ b/tunneldns/tunnel.go
@@ -91,28 +91,37 @@ func CreateListener(address string, port uint16, upstreams []string, bootstraps
 		upstreamList = append(upstreamList, upstream)
 	}
 
-	
-	
-	// Create a local cache with HTTPS proxy plugin
-	proxyPlugin := ProxyPlugin{
-		Upstreams: upstreamList,
-	}
-	chain := cache.New()
-	chain.Next = proxyPlugin	
-
-	// Optionally disable http response caching
-	if os.Getenv("DISABLE_TUNNELDNS_CACHE") == "true" { 
-		chain = proxyPlugin 
-	}
-	
 	// Format an endpoint
 	endpoint := "dns://" + net.JoinHostPort(address, strconv.FormatUint(uint64(port), 10))
 
+	// get middlewares for DNS Server
+	middlewares := getMiddlewares(upstreamList)
+	chain := NewMetricsPlugin(middlewares)
+
 	// Create the actual middleware server
-	server, err := dnsserver.NewServer(endpoint, []*dnsserver.Config{createConfig(address, port, NewMetricsPlugin(chain))})
+	server, err := dnsserver.NewServer(endpoint, []*dnsserver.Config{createConfig(address, port, chain)})
 	if err != nil {
 		return nil, err
 	}
+	
 
 	return &Listener{server: server, log: log}, nil
 }
+
+// getMiddlewares() returns the middleware chain for DNS Server
+// 
+// Middleware includes features like...
+// * Response Cache
+// * HTTP Proxy Settings
+func getMiddlewares(upstreamList []Upstream) plugin.Handler { 
+	proxyPlugin := ProxyPlugin{
+		Upstreams: upstreamList,
+	}
+	cachePlugin := cache.New()
+	cachePlugin.Next = proxyPlugin
+
+	if os.Getenv("DISABLE_TUNNELDNS_CACHE") == "true" { 
+		return cachePlugin 
+	} 
+	return proxyPlugin
+}
\ No newline at end of file