bucketeer/internal/dns/resolver_test.go

package dns

import (
	"context"
	"errors"
	"net"
	"testing"
	"time"
)

type fakeResolver struct {
	cnames map[string]string
	ips    map[string][]net.IP
	err    error
}

func (f *fakeResolver) LookupCNAME(ctx context.Context, host string) (string, error) {
	if f.err != nil {
		return "", f.err
	}
	if c, ok := f.cnames[host]; ok {
		return c, nil
	}
	return host, nil
}

func (f *fakeResolver) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	if f.err != nil {
		return nil, f.err
	}
	return f.ips[host], nil
}

func TestFollowChain(t *testing.T) {
	r := &fakeResolver{cnames: map[string]string{
		"a.example.com":      "b.example.com",
		"b.example.com":      "c.s3.amazonaws.com",
		"c.s3.amazonaws.com": "c.s3.amazonaws.com",
	}}
	chain, err := Follow(context.Background(), r, "a.example.com", 8)
	if err != nil {
		t.Fatal(err)
	}
	want := []string{"a.example.com", "b.example.com", "c.s3.amazonaws.com"}
	if len(chain) != len(want) {
		t.Fatalf("chain = %v, want %v", chain, want)
	}
	for i := range want {
		if chain[i] != want[i] {
			t.Fatalf("chain = %v, want %v", chain, want)
		}
	}
}

func TestFollowLoop(t *testing.T) {
	r := &fakeResolver{cnames: map[string]string{
		"a.example.com": "b.example.com",
		"b.example.com": "a.example.com",
	}}
	chain, err := Follow(context.Background(), r, "a.example.com", 8)
	if !errors.Is(err, ErrCNAMELoop) {
		t.Fatalf("err = %v, want ErrCNAMELoop", err)
	}
	if len(chain) < 2 {
		t.Fatalf("chain = %v, want at least 2 entries", chain)
	}
}

func TestFollowDepth(t *testing.T) {
	r := &fakeResolver{cnames: map[string]string{
		"h1.example.com": "h2.example.com",
		"h2.example.com": "h3.example.com",
		"h3.example.com": "h4.example.com",
	}}
	chain, err := Follow(context.Background(), r, "h1.example.com", 2)
	if err != nil {
		t.Fatal(err)
	}
	if len(chain) != 3 {
		t.Fatalf("chain = %v, want 3 entries (host + 2 hops)", chain)
	}
}

func TestEnricherCaches(t *testing.T) {
	r := &fakeResolver{
		cnames: map[string]string{"a.example.com": "b.s3.amazonaws.com"},
		ips:    map[string][]net.IP{"a.example.com": {net.ParseIP("1.2.3.4")}},
	}
	e := &Enricher{Resolver: r, Cache: NewTTLCache[Result](16, 2), PositiveTTL: time.Minute, NegativeTTL: time.Minute, MaxDepth: 8}
	if _, ok := e.Resolve(context.Background(), "a.example.com"); !ok {
		t.Fatal("first resolve failed")
	}
	r.cnames = nil // a cached lookup must not hit the resolver
	res, ok := e.Resolve(context.Background(), "a.example.com")
	if !ok || len(res.IPs) != 1 {
		t.Fatalf("cached resolve = %+v ok=%v", res, ok)
	}
}

func TestEnricherNegativeCache(t *testing.T) {
	r := &fakeResolver{err: &net.DNSError{Err: "no such host", IsNotFound: true}}
	e := &Enricher{Resolver: r, Cache: NewTTLCache[Result](16, 2), PositiveTTL: time.Minute, NegativeTTL: time.Minute, MaxDepth: 8}
	if _, ok := e.Resolve(context.Background(), "missing.example.com"); ok {
		t.Fatal("expected failure")
	}
	if _, ok := e.Resolve(context.Background(), "missing.example.com"); ok {
		t.Fatal("expected cached failure")
	}
}

type ctxResolver struct{}

func (ctxResolver) LookupCNAME(ctx context.Context, host string) (string, error) {
	return "", ctx.Err()
}
func (ctxResolver) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	return nil, ctx.Err()
}

func TestResolveContextCancellation(t *testing.T) {
	e := &Enricher{Resolver: ctxResolver{}, Cache: NewTTLCache[Result](16, 2), PositiveTTL: time.Minute, NegativeTTL: time.Minute, MaxDepth: 8}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, ok := e.Resolve(ctx, "example.com"); ok {
		t.Fatal("expected cancellation to fail resolution")
	}
}

func TestCacheBoundAndExpiry(t *testing.T) {
	c := NewTTLCache[int](8, 2)
	for i := 0; i < 100; i++ {
		c.Set(string(rune('a'+i%26))+string(rune('0'+i%10)), i, time.Minute)
	}
	if c.Len() > 8 {
		t.Fatalf("cache len = %d, want <= 8", c.Len())
	}
	c.Set("k", 1, time.Millisecond)
	time.Sleep(5 * time.Millisecond)
	if _, ok := c.Get("k"); ok {
		t.Fatal("expected expired entry to be gone")
	}
}