Replace go-cache with golang-lru in dnssec (#336)

* Replace go-cache with golang-lru

This fix replace go-cache with golang-lru, as is specified in 335.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Move cache initialization to setup

This commit move cache initialization to setup in dnssec middleware.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

5 files changed, 27 insertions, 16 deletions

diff --git a/middleware/dnssec/cache_test.go b/middleware/dnssec/cache_test.go
index 6624f65ec..a3b7b818d 100644
--- a/middleware/dnssec/cache_test.go
+++ b/middleware/dnssec/cache_test.go
@@ -6,6 +6,8 @@ import (
 
 	"github.com/miekg/coredns/middleware/test"
 	"github.com/miekg/coredns/request"
+
+	"github.com/hashicorp/golang-lru"
 )
 
 func TestCacheSet(t *testing.T) {
@@ -19,10 +21,11 @@ func TestCacheSet(t *testing.T) {
 		t.Fatalf("failed to parse key: %v\n", err)
 	}
 
+	cache, _ := lru.New(defaultCap)
 	m := testMsg()
 	state := request.Request{Req: m}
 	k := key(m.Answer) // calculate *before* we add the sig
-	d := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, nil)
+	d := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, nil, cache)
 	m = d.Sign(state, "miek.nl.", time.Now().UTC())
 
 	_, ok := d.get(k)
diff --git a/middleware/dnssec/dnssec.go b/middleware/dnssec/dnssec.go
index 1433fb1a3..de907c985 100644
--- a/middleware/dnssec/dnssec.go
+++ b/middleware/dnssec/dnssec.go
@@ -10,8 +10,8 @@ import (
 	"github.com/miekg/coredns/middleware/pkg/singleflight"
 	"github.com/miekg/coredns/request"
 
+	"github.com/hashicorp/golang-lru"
 	"github.com/miekg/dns"
-	gcache "github.com/patrickmn/go-cache"
 )
 
 // Dnssec signs the reply on-the-fly.
@@ -21,15 +21,15 @@ type Dnssec struct {
 	zones    []string
 	keys     []*DNSKEY
 	inflight *singleflight.Group
-	cache    *gcache.Cache
+	cache    *lru.Cache
 }
 
 // New returns a new Dnssec.
-func New(zones []string, keys []*DNSKEY, next middleware.Handler) Dnssec {
+func New(zones []string, keys []*DNSKEY, next middleware.Handler, cache *lru.Cache) Dnssec {
 	return Dnssec{Next: next,
 		zones:    zones,
 		keys:     keys,
-		cache:    gcache.New(defaultDuration, purgeDuration),
+		cache:    cache,
 		inflight: new(singleflight.Group),
 	}
 }
@@ -110,9 +110,7 @@ func (d Dnssec) sign(rrs []dns.RR, signerName string, ttl, incep, expir uint32)
 }
 
 func (d Dnssec) set(key string, sigs []dns.RR) {
-	// we insert the sigs with a duration that is 24 hours less then the expiration, as these
-	// sigs have *just* been made the duration is 7 days.
-	d.cache.Set(key, sigs, eightDays-24*time.Hour)
+	d.cache.Add(key, sigs)
 }
 
 func (d Dnssec) get(key string) ([]dns.RR, bool) {
@@ -129,7 +127,6 @@ func incepExpir(now time.Time) (uint32, uint32) {
 }
 
 const (
-	purgeDuration   = 3 * time.Hour
-	defaultDuration = 24 * time.Hour
-	eightDays       = 8 * 24 * time.Hour
+	eightDays  = 8 * 24 * time.Hour
+	defaultCap = 10000 // default capacity of the cache.
 )
diff --git a/middleware/dnssec/dnssec_test.go b/middleware/dnssec/dnssec_test.go
index a3eeec6af..d8f20bc00 100644
--- a/middleware/dnssec/dnssec_test.go
+++ b/middleware/dnssec/dnssec_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/miekg/coredns/middleware/test"
 	"github.com/miekg/coredns/request"
 
+	"github.com/hashicorp/golang-lru"
 	"github.com/miekg/dns"
 )
 
@@ -68,7 +69,8 @@ func TestSigningDifferentZone(t *testing.T) {
 
 	m := testMsgEx()
 	state := request.Request{Req: m}
-	d := New([]string{"example.org."}, []*DNSKEY{key}, nil)
+	cache, _ := lru.New(defaultCap)
+	d := New([]string{"example.org."}, []*DNSKEY{key}, nil, cache)
 	m = d.Sign(state, "example.org.", time.Now().UTC())
 	if !section(m.Answer, 1) {
 		t.Errorf("answer section should have 1 sig")
@@ -157,7 +159,8 @@ func testDelegationMsg() *dns.Msg {
 
 func newDnssec(t *testing.T, zones []string) (Dnssec, func(), func()) {
 	k, rm1, rm2 := newKey(t)
-	d := New(zones, []*DNSKEY{k}, nil)
+	cache, _ := lru.New(defaultCap)
+	d := New(zones, []*DNSKEY{k}, nil, cache)
 	return d, rm1, rm2
 }
 
diff --git a/middleware/dnssec/handler_test.go b/middleware/dnssec/handler_test.go
index a490c3744..c1ed240c3 100644
--- a/middleware/dnssec/handler_test.go
+++ b/middleware/dnssec/handler_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/miekg/coredns/middleware/pkg/dnsrecorder"
 	"github.com/miekg/coredns/middleware/test"
 
+	"github.com/hashicorp/golang-lru"
 	"github.com/miekg/dns"
 	"golang.org/x/net/context"
 )
@@ -77,7 +78,8 @@ func TestLookupZone(t *testing.T) {
 	dnskey, rm1, rm2 := newKey(t)
 	defer rm1()
 	defer rm2()
-	dh := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, fm)
+	cache, _ := lru.New(defaultCap)
+	dh := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, fm, cache)
 	ctx := context.TODO()
 
 	for _, tc := range dnsTestCases {
@@ -115,7 +117,8 @@ func TestLookupDNSKEY(t *testing.T) {
 	dnskey, rm1, rm2 := newKey(t)
 	defer rm1()
 	defer rm2()
-	dh := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, test.ErrorHandler())
+	cache, _ := lru.New(defaultCap)
+	dh := New([]string{"miek.nl."}, []*DNSKEY{dnskey}, test.ErrorHandler(), cache)
 	ctx := context.TODO()
 
 	for _, tc := range dnssecTestCases {
diff --git a/middleware/dnssec/setup.go b/middleware/dnssec/setup.go
index 78d285f09..a61a42453 100644
--- a/middleware/dnssec/setup.go
+++ b/middleware/dnssec/setup.go
@@ -6,6 +6,7 @@ import (
 	"github.com/miekg/coredns/core/dnsserver"
 	"github.com/miekg/coredns/middleware"
 
+	"github.com/hashicorp/golang-lru"
 	"github.com/mholt/caddy"
 )
 
@@ -22,8 +23,12 @@ func setup(c *caddy.Controller) error {
 		return middleware.Error("dnssec", err)
 	}
 
+	cache, err := lru.New(defaultCap)
+	if err != nil {
+		return err
+	}
 	dnsserver.GetConfig(c).AddMiddleware(func(next middleware.Handler) middleware.Handler {
-		return New(zones, keys, next)
+		return New(zones, keys, next, cache)
 	})
 
 	return nil