Skip to content

Commit

Permalink
Remove MXLookup cache, add concept of caching lookups specific to a n…
Browse files Browse the repository at this point in the history
…ameserver (#447)

* add nameserver for cache

* added basic cache unit tests

* remove comments

* use cache for both external/iterative lookups, vary whether to care about nses

* added cache integration test for external lookups

* remove mxlookup cache, use cache on external lookups specific to nameservers

* added unit test, cleaned up lint errors

* fix tests

* commented cachedRetryingLookup

* add details to print

* adjusted constant to compare cached lookups

* populate protocol/resolver on cache retrieval
  • Loading branch information
phillip-stephens authored Sep 18, 2024
1 parent 5dc707f commit 6aa8a81
Show file tree
Hide file tree
Showing 6 changed files with 240 additions and 56 deletions.
25 changes: 2 additions & 23 deletions src/modules/mxlookup/mx_lookup.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,12 @@ package mxlookup

import (
"strings"
"sync"

"github.com/pkg/errors"
log "github.com/sirupsen/logrus"
"github.com/zmap/dns"

"github.com/zmap/zdns/src/cli"
"github.com/zmap/zdns/src/internal/cachehash"
"github.com/zmap/zdns/src/zdns"
)

Expand Down Expand Up @@ -51,11 +49,8 @@ type MXResult struct {
}

type MXLookupModule struct {
IPv4Lookup bool `long:"ipv4-lookup" description:"perform A lookups for each MX server"`
IPv6Lookup bool `long:"ipv6-lookup" description:"perform AAAA record lookups for each MX server"`
MXCacheSize int `long:"mx-cache-size" default:"1000" description:"number of records to store in MX -> A/AAAA cache"`
CacheHash *cachehash.CacheHash
CHmu sync.Mutex
IPv4Lookup bool `long:"ipv4-lookup" description:"perform A lookups for each MX server"`
IPv6Lookup bool `long:"ipv6-lookup" description:"perform AAAA record lookups for each MX server"`
cli.BasicLookupModule
}

Expand All @@ -77,31 +72,15 @@ func (mxMod *MXLookupModule) Init() {
if !mxMod.IPv4Lookup && !mxMod.IPv6Lookup {
log.Fatal("At least one of ipv4-lookup or ipv6-lookup must be true")
}
if mxMod.MXCacheSize <= 0 {
log.Fatal("mxCacheSize must be greater than 0, got ", mxMod.MXCacheSize)
}
mxMod.CacheHash = new(cachehash.CacheHash)
mxMod.CacheHash.Init(mxMod.MXCacheSize)
}

func (mxMod *MXLookupModule) lookupIPs(r *zdns.Resolver, name string, nameServer *zdns.NameServer, ipMode zdns.IPVersionMode) (CachedAddresses, zdns.Trace) {
mxMod.CHmu.Lock()
// TODO - Phillip this comment V is present in the original code and has been there since 2017 IIRC, so ask Zakir what to do
// XXX this should be changed to a miekglookup
res, found := mxMod.CacheHash.Get(name)
mxMod.CHmu.Unlock()
if found {
return res.(CachedAddresses), zdns.Trace{}
}
retv := CachedAddresses{}
result, trace, status, _ := r.DoTargetedLookup(name, nameServer, mxMod.IsIterative, mxMod.IPv4Lookup, mxMod.IPv6Lookup)
if status == zdns.StatusNoError && result != nil {
retv.IPv4Addresses = result.IPv4Addresses
retv.IPv6Addresses = result.IPv6Addresses
}
mxMod.CHmu.Lock()
mxMod.CacheHash.Upsert(name, retv)
mxMod.CHmu.Unlock()
return retv, trace
}

Expand Down
64 changes: 39 additions & 25 deletions src/zdns/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@ type TimedAnswer struct {
ExpiresAt time.Time
}

type CachedKey struct {
Question Question
NameServer string // optional
}

type CachedResult struct {
Answers map[interface{}]TimedAnswer
}
Expand All @@ -46,7 +51,7 @@ func (s *Cache) VerboseLog(depth int, args ...interface{}) {
log.Debug(makeVerbosePrefix(depth), args)
}

func (s *Cache) AddCachedAnswer(answer interface{}, depth int) {
func (s *Cache) AddCachedAnswer(answer interface{}, ns *NameServer, depth int) {
a, ok := answer.(Answer)
if !ok {
// we can't cache this entry because we have no idea what to name it
Expand All @@ -64,13 +69,17 @@ func (s *Cache) AddCachedAnswer(answer interface{}, depth int) {
return
}
expiresAt := time.Now().Add(time.Duration(a.TTL) * time.Second)
s.IterativeCache.Lock(q)
defer s.IterativeCache.Unlock(q)
// don't bother to move this to the top of the linked list. we're going
// to add this record back in momentarily and that will take care of this
ca := CachedResult{}
ca.Answers = make(map[interface{}]TimedAnswer)
i, ok := s.IterativeCache.GetNoMove(q)
cacheKey := CachedKey{q, ""}
if ns != nil {
cacheKey.NameServer = ns.String()
}
s.IterativeCache.Lock(cacheKey)
defer s.IterativeCache.Unlock(cacheKey)
// don't bother to move this to the top of the linked list. we're going
// to add this record back in momentarily and that will take care of this
i, ok := s.IterativeCache.GetNoMove(cacheKey)
if ok {
// record found, check type on interface
ca, ok = i.(CachedResult)
Expand All @@ -83,18 +92,24 @@ func (s *Cache) AddCachedAnswer(answer interface{}, depth int) {
Answer: answer,
ExpiresAt: expiresAt}
ca.Answers[a] = ta
s.IterativeCache.Add(q, ca)
s.IterativeCache.Add(cacheKey, ca)
s.VerboseLog(depth+1, "Upsert cached answer ", q, " ", ca)
}

func (s *Cache) GetCachedResult(q Question, isAuthCheck bool, depth int) (SingleQueryResult, bool) {
s.VerboseLog(depth+1, "Cache request for: ", q.Name, " (", q.Type, ")")
func (s *Cache) GetCachedResult(q Question, ns *NameServer, depth int) (SingleQueryResult, bool) {
var retv SingleQueryResult
s.IterativeCache.Lock(q)
unres, ok := s.IterativeCache.Get(q)
cacheKey := CachedKey{q, ""}
if ns != nil {
cacheKey.NameServer = ns.String()
s.VerboseLog(depth+1, "Cache request for: ", q.Name, " (", q.Type, ") @", cacheKey.NameServer)
} else {
s.VerboseLog(depth+1, "Cache request for: ", q.Name, " (", q.Type, ")")
}
s.IterativeCache.Lock(cacheKey)
unres, ok := s.IterativeCache.Get(cacheKey)
if !ok { // nothing found
s.VerboseLog(depth+2, "-> no entry found in cache")
s.IterativeCache.Unlock(q)
s.IterativeCache.Unlock(cacheKey)
return retv, false
}
retv.Authorities = make([]interface{}, 0)
Expand All @@ -116,26 +131,25 @@ func (s *Cache) GetCachedResult(q Question, isAuthCheck bool, depth int) (Single
delete(cachedRes.Answers, k)
} else {
// this result is valid. append it to the SingleQueryResult we're going to hand to the user
if isAuthCheck {
retv.Authorities = append(retv.Authorities, cachedAnswer.Answer)
} else {
retv.Answers = append(retv.Answers, cachedAnswer.Answer)
}
retv.Answers = append(retv.Answers, cachedAnswer.Answer)
}
}
s.IterativeCache.Unlock(q)
s.IterativeCache.Unlock(cacheKey)
// Don't return an empty response.
if len(retv.Answers) == 0 && len(retv.Authorities) == 0 && len(retv.Additional) == 0 {
s.VerboseLog(depth+2, "-> no entry found in cache, after expiration")
var emptyRetv SingleQueryResult
return emptyRetv, false
}
if ns != nil {
retv.Resolver = ns.String()
}

s.VerboseLog(depth+2, "Cache hit: ", retv)
return retv, true
}

func (s *Cache) SafeAddCachedAnswer(a interface{}, layer string, debugType string, depth int) {
func (s *Cache) SafeAddCachedAnswer(a interface{}, ns *NameServer, layer string, debugType string, depth int) {
ans, ok := a.(Answer)
if !ok {
s.VerboseLog(depth+1, "unable to cast ", debugType, ": ", layer, ": ", a)
Expand All @@ -145,19 +159,19 @@ func (s *Cache) SafeAddCachedAnswer(a interface{}, layer string, debugType strin
log.Info("detected poison ", debugType, ": ", ans.Name, "(", ans.Type, "): ", layer, ": ", a)
return
}
s.AddCachedAnswer(a, depth)
s.AddCachedAnswer(a, ns, depth)
}

func (s *Cache) CacheUpdate(layer string, result SingleQueryResult, depth int) {
func (s *Cache) CacheUpdate(layer string, result SingleQueryResult, ns *NameServer, depth int, cacheNonAuthoritativeAns bool) {
for _, a := range result.Additional {
s.SafeAddCachedAnswer(a, layer, "additional", depth)
s.SafeAddCachedAnswer(a, ns, layer, "additional", depth)
}
for _, a := range result.Authorities {
s.SafeAddCachedAnswer(a, layer, "authority", depth)
s.SafeAddCachedAnswer(a, ns, layer, "authority", depth)
}
if result.Flags.Authoritative {
if result.Flags.Authoritative || cacheNonAuthoritativeAns {
for _, a := range result.Answers {
s.SafeAddCachedAnswer(a, layer, "answer", depth)
s.SafeAddCachedAnswer(a, ns, layer, "answer", depth)
}
}
}
148 changes: 148 additions & 0 deletions src/zdns/cache_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
/*
* ZDNS Copyright 2024 Regents of the University of Michigan
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
* use this file except in compliance with the License. You may obtain a copy
* of the License at http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
* implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package zdns

import (
"net"
"testing"

"github.com/stretchr/testify/assert"
)

func TestCheckForNonExistentKey(t *testing.T) {
cache := Cache{}
cache.Init(4096)
_, found := cache.GetCachedResult(Question{1, 1, "google.com"}, nil, 0)
assert.False(t, found, "Expected no cache entry")
}

func TestNoNameServerLookupSuccess(t *testing.T) {
res := SingleQueryResult{
Answers: []interface{}{Answer{
TTL: 3600,
RrType: 1,
RrClass: 1,
Name: "google.com.",
Answer: "192.0.2.1",
}},
Additional: nil,
Authorities: nil,
Protocol: "",
Flags: DNSFlags{Authoritative: true},
}
cache := Cache{}
cache.Init(4096)
cache.CacheUpdate(".", res, nil, 0, false)
_, found := cache.GetCachedResult(Question{1, 1, "google.com."}, nil, 0)
assert.True(t, found, "Expected cache entry")
}

func TestNoNameServerLookupForNamedNameServer(t *testing.T) {
res := SingleQueryResult{
Answers: []interface{}{Answer{
TTL: 3600,
RrType: 1,
RrClass: 1,
Name: "google.com.",
Answer: "192.0.2.1",
}},
Additional: nil,
Authorities: nil,
Protocol: "",
Flags: DNSFlags{Authoritative: true},
}
cache := Cache{}
cache.Init(4096)
cache.CacheUpdate(".", res, nil, 0, false)
_, found := cache.GetCachedResult(Question{1, 1, "google.com."}, &NameServer{
IP: net.ParseIP("1.1.1.1"),
Port: 53,
}, 0)
assert.False(t, found, "Cache has an answer from a generic nameserver, we wanted a specific one. Shouldn't be found.")
}

func TestNamedServerLookupForNonNamedNameServer(t *testing.T) {
res := SingleQueryResult{
Answers: []interface{}{Answer{
TTL: 3600,
RrType: 1,
RrClass: 1,
Name: "google.com.",
Answer: "192.0.2.1",
}},
Additional: nil,
Authorities: nil,
Protocol: "",
Flags: DNSFlags{Authoritative: true},
}
cache := Cache{}
cache.Init(4096)
cache.CacheUpdate(".", res, &NameServer{
IP: net.ParseIP("1.1.1.1"),
Port: 53,
}, 0, false)
_, found := cache.GetCachedResult(Question{1, 1, "google.com."}, nil, 0)
assert.False(t, found, "Cache has an answer from a named nameserver, we wanted a generic one. Shouldn't be found.")
}

func TestNamedServerLookupForNamedNameServer(t *testing.T) {
res := SingleQueryResult{
Answers: []interface{}{Answer{
TTL: 3600,
RrType: 1,
RrClass: 1,
Name: "google.com.",
Answer: "192.0.2.1",
}},
Additional: nil,
Authorities: nil,
Protocol: "",
Flags: DNSFlags{Authoritative: true},
}
cache := Cache{}
cache.Init(4096)
cache.CacheUpdate(".", res, &NameServer{
IP: net.ParseIP("1.1.1.1"),
Port: 53,
}, 0, false)
_, found := cache.GetCachedResult(Question{1, 1, "google.com."}, &NameServer{
IP: net.ParseIP("1.1.1.1"),
Port: 53,
}, 0)
assert.True(t, found, "Should be found")
}

func TestNoNameServerLookupNotAuthoritative(t *testing.T) {
res := SingleQueryResult{
Answers: []interface{}{Answer{
TTL: 3600,
RrType: 1,
RrClass: 1,
Name: "google.com.",
Answer: "192.0.2.1",
}},
Additional: nil,
Authorities: nil,
Protocol: "",
Flags: DNSFlags{Authoritative: false},
}
cache := Cache{}
cache.Init(4096)
cache.CacheUpdate(".", res, nil, 0, false)
_, found := cache.GetCachedResult(Question{1, 1, "google.com."}, nil, 0)
assert.False(t, found, "shouldn't cache non-authoritative answers")
cache.CacheUpdate(".", res, nil, 0, true)
_, found = cache.GetCachedResult(Question{1, 1, "google.com."}, nil, 0)
assert.True(t, found, "should cache non-authoritative answers")
}
Loading

0 comments on commit 6aa8a81

Please sign in to comment.