From a5d105352057bf24a59a08a2695f1d48f033cb17 Mon Sep 17 00:00:00 2001
From: Eugene Bujak <hmage@hmage.net>
Date: Wed, 10 Oct 2018 19:10:38 +0300
Subject: [PATCH] Cache DNS lookups when resolving safebrowsing or parental
 servers, also cache replacement hostnames as well.

---
 coredns_plugin/coredns_plugin.go | 45 +++++++++++++++++++++++---------
 dnsfilter/dnsfilter.go           |  1 +
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/coredns_plugin/coredns_plugin.go b/coredns_plugin/coredns_plugin.go
index 195a9218..4b00eca6 100644
--- a/coredns_plugin/coredns_plugin.go
+++ b/coredns_plugin/coredns_plugin.go
@@ -45,6 +45,16 @@ func init() {
 	})
 }
 
+type cacheEntry struct {
+	answer      []dns.RR
+	lastUpdated time.Time
+}
+
+var (
+	lookupCacheTime = time.Minute * 30
+	lookupCache     = map[string]cacheEntry{}
+)
+
 type plugSettings struct {
 	SafeBrowsingBlockHost string
 	ParentalBlockHost     string
@@ -324,20 +334,29 @@ func (p *plug) replaceHostWithValAndReply(ctx context.Context, w dns.ResponseWri
 		records = append(records, result)
 	} else {
 		// this is a domain name, need to look it up
-		req := new(dns.Msg)
-		req.SetQuestion(dns.Fqdn(val), question.Qtype)
-		req.RecursionDesired = true
-		reqstate := request.Request{W: w, Req: req, Context: ctx}
-		result, err := p.upstream.Lookup(reqstate, dns.Fqdn(val), reqstate.QType())
-		if err != nil {
-			log.Printf("Got error %s\n", err)
-			return dns.RcodeServerFailure, fmt.Errorf("plugin/dnsfilter: %s", err)
-		}
-		if result != nil {
-			for _, answer := range result.Answer {
-				answer.Header().Name = question.Name
+		cacheentry := lookupCache[val]
+		if time.Since(cacheentry.lastUpdated) > lookupCacheTime {
+			req := new(dns.Msg)
+			req.SetQuestion(dns.Fqdn(val), question.Qtype)
+			req.RecursionDesired = true
+			reqstate := request.Request{W: w, Req: req, Context: ctx}
+			result, err := p.upstream.Lookup(reqstate, dns.Fqdn(val), reqstate.QType())
+			if err != nil {
+				log.Printf("Got error %s\n", err)
+				return dns.RcodeServerFailure, fmt.Errorf("plugin/dnsfilter: %s", err)
 			}
-			records = result.Answer
+			if result != nil {
+				for _, answer := range result.Answer {
+					answer.Header().Name = question.Name
+				}
+				records = result.Answer
+				cacheentry.answer = result.Answer
+				cacheentry.lastUpdated = time.Now()
+				lookupCache[val] = cacheentry
+			}
+		} else {
+			// get from cache
+			records = cacheentry.answer
 		}
 	}
 	m := new(dns.Msg)
diff --git a/dnsfilter/dnsfilter.go b/dnsfilter/dnsfilter.go
index 1344ea9f..ec2ad5f7 100644
--- a/dnsfilter/dnsfilter.go
+++ b/dnsfilter/dnsfilter.go
@@ -16,6 +16,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	_ "github.com/benburkert/dns/init"
 	"github.com/bluele/gcache"
 	"golang.org/x/net/publicsuffix"
 )