Fix DNS cache

Author: btomaev

internal/dns/cache.go

@@ -32,6 +32,9 @@ func NewDNSCache(forwardTTL int, reverseTTL int, backend store.Interface) *DnsCa
 }
 
 func (d *Dns) getCachedForward(host string) ([]string, bool) {
+	if d.cache == nil {
+		return nil, false
+	}
 	if cached, err := d.cache.forward.Get(d.ctx, host); err == nil {
 		slog.Debug("DNS: forward cache hit", "name", host, "ips", cached)
 		return cached, true
@@ -41,6 +44,9 @@ func (d *Dns) getCachedForward(host string) ([]string, bool) {
 }
 
 func (d *Dns) getCachedReverse(addr string) ([]string, bool) {
+	if d.cache == nil {
+		return nil, false
+	}
 	if cached, err := d.cache.reverse.Get(d.ctx, addr); err == nil {
 		slog.Debug("DNS: reverse cache hit", "addr", addr, "names", cached)
 		return cached, true
@@ -50,9 +56,15 @@ func (d *Dns) getCachedReverse(addr string) ([]string, bool) {
 }
 
 func (d *Dns) forwardCachePut(host string, entries []string) {
+	if d.cache == nil {
+		return
+	}
 	d.cache.forward.Set(d.ctx, host, entries, d.cache.forwardTTL)
 }
 
 func (d *Dns) reverseCachePut(addr string, entries []string) {
+	if d.cache == nil {
+		return
+	}
 	d.cache.reverse.Set(d.ctx, addr, entries, d.cache.reverseTTL)
 }

lib/policy/expressions/environment_test.go

@@ -1,18 +1,28 @@
 package expressions
 
 import (
+	"context"
 	"errors"
 	"net"
 	"strings"
 	"testing"
 
-	"github.com/TecharoHQ/anubis/internal"
+	"github.com/TecharoHQ/anubis/internal/dns"
+	"github.com/TecharoHQ/anubis/lib/store/memory"
 	"github.com/google/cel-go/common/types"
 	"github.com/google/cel-go/common/types/ref"
 )
 
+// newTestDNS is a helper function to create a new Dns object with an in-memory cache for testing.
+func newTestDNS(forwardTTL int, reverseTTL int) *dns.Dns {
+	ctx := context.Background()
+	memStore := memory.New(ctx)
+	cache := dns.NewDNSCache(forwardTTL, reverseTTL, memStore)
+	return dns.New(ctx, cache)
+}
+
 func TestBotEnvironment(t *testing.T) {
-	dnsObj := internal.NewDNS(300, 300)
+	dnsObj := newTestDNS(300, 300)
 	env, err := BotEnvironment(dnsObj)
 	if err != nil {
 		t.Fatalf("failed to create bot environment: %v", err)
@@ -291,11 +301,11 @@ func TestBotEnvironment(t *testing.T) {
 	})
 
 	t.Run("dnsFunctions", func(t *testing.T) {
-		originalDNSLookupAddr := internal.DNSLookupAddr
-		originalDNSLookupHost := internal.DNSLookupHost
+		originalDNSLookupAddr := dns.DNSLookupAddr
+		originalDNSLookupHost := dns.DNSLookupHost
 		defer func() {
-			internal.DNSLookupAddr = originalDNSLookupAddr
-			internal.DNSLookupHost = originalDNSLookupHost
+			dns.DNSLookupAddr = originalDNSLookupAddr
+			dns.DNSLookupHost = originalDNSLookupHost
 		}()
 
 		t.Run("reverseDNS", func(t *testing.T) {
@@ -337,7 +347,7 @@ func TestBotEnvironment(t *testing.T) {
 
 			for _, tt := range tests {
 				t.Run(tt.name, func(t *testing.T) {
-					internal.DNSLookupAddr = func(addr string) ([]string, error) {
+					dns.DNSLookupAddr = func(addr string) ([]string, error) {
 						if addr == tt.addr {
 							return tt.mockReturn, tt.mockError
 						}
@@ -399,7 +409,7 @@ func TestBotEnvironment(t *testing.T) {
 
 			for _, tt := range tests {
 				t.Run(tt.name, func(t *testing.T) {
-					internal.DNSLookupHost = func(host string) ([]string, error) {
+					dns.DNSLookupHost = func(host string) ([]string, error) {
 						if host == tt.host {
 							return tt.mockReturn, tt.mockError
 						}
@@ -482,13 +492,13 @@ func TestBotEnvironment(t *testing.T) {
 
 			for _, tt := range tests {
 				t.Run(tt.name, func(t *testing.T) {
-					internal.DNSLookupAddr = func(addr string) ([]string, error) {
+					dns.DNSLookupAddr = func(addr string) ([]string, error) {
 						if addr == tt.addr {
 							return tt.reverseMockReturn, tt.reverseMockError
 						}
 						return nil, errors.New("unexpected address for reverse lookup")
 					}
-					internal.DNSLookupHost = func(host string) ([]string, error) {
+					dns.DNSLookupHost = func(host string) ([]string, error) {
 						host = strings.TrimSuffix(host, ".")
 						if ips, ok := tt.forwardMockReturn[host]; ok {
 							return ips, nil

lib/policy/policy.go

@@ -43,6 +43,7 @@ type ParsedConfig struct {
 	StatusCodes       config.StatusCodes
 	DefaultDifficulty int
 	DNSBL             bool
+	DnsCache          *dns.DnsCache
 	Dns               *dns.Dns
 	Logger            *slog.Logger
 }
@@ -68,6 +69,22 @@ func ParseConfig(ctx context.Context, fin io.Reader, fname string, defaultDiffic
 	result := newParsedConfig(c)
 	result.DefaultDifficulty = defaultDifficulty
 
+	stFac, ok := store.Get(c.Store.Backend)
+	switch ok {
+	case true:
+		store, err := stFac.Build(ctx, c.Store.Parameters)
+		if err != nil {
+			validationErrs = append(validationErrs, err)
+		} else {
+			result.Store = store
+		}
+	case false:
+		validationErrs = append(validationErrs, config.ErrUnknownStoreBackend)
+	}
+
+	result.DnsCache = dns.NewDNSCache(result.orig.DNSTTL.Forward, result.orig.DNSTTL.Reverse, result.Store)
+	result.Dns = dns.New(ctx, result.DnsCache)
+
 	for _, b := range c.Bots {
 		if berr := b.Valid(); berr != nil {
 			validationErrs = append(validationErrs, berr)
@@ -196,22 +213,6 @@ func ParseConfig(ctx context.Context, fin io.Reader, fname string, defaultDiffic
 		result.Thresholds = append(result.Thresholds, threshold)
 	}
 
-	stFac, ok := store.Get(c.Store.Backend)
-	switch ok {
-	case true:
-		store, err := stFac.Build(ctx, c.Store.Parameters)
-		if err != nil {
-			validationErrs = append(validationErrs, err)
-		} else {
-			result.Store = store
-		}
-	case false:
-		validationErrs = append(validationErrs, config.ErrUnknownStoreBackend)
-	}
-
-	dnsCache := dns.NewDNSCache(result.orig.DNSTTL.Forward, result.orig.DNSTTL.Reverse, result.Store)
-	result.Dns = dns.New(ctx, dnsCache)
-
 	if c.Logging.Level != nil {
 		logLevel = c.Logging.Level.String()
 	}