bucketeer/internal/pipeline/pipeline_test.go

package pipeline

import (
	"context"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"

	"software.sslmate.com/src/certspotter"

	"tea.karlcom.de/karl/bucketeer/internal/dns"
	"tea.karlcom.de/karl/bucketeer/internal/metrics"
	"tea.karlcom.de/karl/bucketeer/internal/model"
)

func makeHosts(prefix, suffix string) []string {
	hosts := make([]string, 256)
	for i := range hosts {
		hosts[i] = prefix + strconv.Itoa(i) + suffix
	}
	return hosts
}

// TestHandleHostNoAllocs audits the sink hot path: after warmup, classifying and
// enqueueing a hostname must not allocate.
func TestHandleHostNoAllocs(t *testing.T) {
	p := New(Config{DNSQueueSize: 1 << 20, DedupSize: 1 << 20})
	nonMatch := makeHosts("api", ".example.com")
	direct := makeHosts("bucket", ".s3.eu-central-1.amazonaws.com")
	for _, h := range nonMatch {
		p.HandleHost(h)
	}
	for _, h := range direct {
		p.HandleHost(h)
	}
	allocs := testing.AllocsPerRun(2000, func() {
		for i := 0; i < 64; i++ {
			p.HandleHost(nonMatch[i])
			p.HandleHost(direct[i])
		}
	})
	if allocs != 0 {
		t.Fatalf("HandleHost allocates %.3f allocs/op, want 0", allocs)
	}
}

func BenchmarkHandleHostNonMatch(b *testing.B) {
	p := New(Config{DNSQueueSize: 1 << 20, DedupSize: 1 << 20})
	hosts := makeHosts("api", ".example.com")
	for _, h := range hosts {
		p.HandleHost(h)
	}
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		p.HandleHost(hosts[i&255])
	}
}

func BenchmarkHandleHostDirectMatch(b *testing.B) {
	p := New(Config{DirectQueueSize: 1 << 20, DedupSize: 1 << 20})
	hosts := makeHosts("bucket", ".s3.eu-central-1.amazonaws.com")
	for _, h := range hosts {
		p.HandleHost(h)
	}
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		p.HandleHost(hosts[i&255])
	}
}

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

func (f *fakeResolver) LookupCNAME(ctx context.Context, host string) (string, error) {
	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) {
	return f.ips[host], nil
}

func newEnricher(r dns.Resolver) *dns.Enricher {
	return &dns.Enricher{
		Resolver:    r,
		Cache:       dns.NewTTLCache[dns.Result](1024, 8),
		PositiveTTL: time.Minute,
		NegativeTTL: time.Minute,
		MaxDepth:    8,
	}
}

type collector struct {
	mu    sync.Mutex
	items []Finding
}

func (c *collector) write(ctx context.Context, batch []Finding) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.items = append(c.items, batch...)
	return nil
}

func (c *collector) all() []Finding {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]Finding(nil), c.items...)
}

func TestDirectMatch(t *testing.T) {
	col := &collector{}
	m := metrics.New()
	p := New(Config{
		Enricher:  newEnricher(&fakeResolver{}),
		Write:     col.write,
		Metrics:   m,
		BatchSize: 10,
	})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Start(ctx)
	p.HandleIdentifiers(ctx, &certspotter.Identifiers{DNSNames: []string{"bucket.s3.eu-central-1.amazonaws.com"}})
	time.Sleep(50 * time.Millisecond)
	p.Close()

	items := col.all()
	if len(items) != 1 {
		t.Fatalf("got %d findings, want 1 (%+v)", len(items), items)
	}
	f := items[0]
	if f.Provider != model.ProviderAWS || f.Service != model.ServiceS3 || f.Source != model.SourceDirectCT {
		t.Fatalf("unexpected finding %+v", f)
	}
	if m.DirectMatches.Load() != 1 {
		t.Fatalf("direct matches = %d, want 1", m.DirectMatches.Load())
	}
}

func TestCNAMEMatch(t *testing.T) {
	col := &collector{}
	r := &fakeResolver{
		cnames: map[string]string{"files.example.com": "files.s3.amazonaws.com"},
		ips:    map[string][]net.IP{"files.s3.amazonaws.com": {net.ParseIP("1.2.3.4")}},
	}
	m := metrics.New()
	p := New(Config{Enricher: newEnricher(r), Write: col.write, Metrics: m, BatchSize: 10})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Start(ctx)
	p.HandleIdentifiers(ctx, &certspotter.Identifiers{DNSNames: []string{"files.example.com"}})
	time.Sleep(50 * time.Millisecond)
	p.Close()

	items := col.all()
	if len(items) != 1 {
		t.Fatalf("got %d findings, want 1", len(items))
	}
	f := items[0]
	if f.Source != model.SourceCNAME || f.Hostname != "files.s3.amazonaws.com" || f.Original != "files.example.com" {
		t.Fatalf("unexpected finding %+v", f)
	}
	if len(f.IPv4) != 1 || f.IPv4[0] != "1.2.3.4" {
		t.Fatalf("unexpected IPs %+v", f.IPv4)
	}
	if m.CNAMEMatches.Load() != 1 {
		t.Fatalf("cname matches = %d, want 1", m.CNAMEMatches.Load())
	}
}

func TestCNAMEMatchIPv6(t *testing.T) {
	col := &collector{}
	r := &fakeResolver{
		cnames: map[string]string{"db.example.com": "mydb.cluster-abc.us-east-1.rds.amazonaws.com"},
		ips:    map[string][]net.IP{"mydb.cluster-abc.us-east-1.rds.amazonaws.com": {net.ParseIP("2606:4700::1111")}},
	}
	p := New(Config{Enricher: newEnricher(r), Write: col.write, BatchSize: 10})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Start(ctx)
	p.HandleHost("db.example.com")
	time.Sleep(50 * time.Millisecond)
	p.Close()
	items := col.all()
	if len(items) != 1 || len(items[0].IPv6) != 1 || items[0].IPv6[0] != "2606:4700::1111" {
		t.Fatalf("unexpected IPv6 finding: %+v", items)
	}
}

type blockingResolver struct{ release chan struct{} }

func (b *blockingResolver) LookupCNAME(ctx context.Context, host string) (string, error) {
	select {
	case <-b.release:
	case <-ctx.Done():
	}
	return host, nil
}

func (b *blockingResolver) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	return nil, nil
}

// TestDirectNotStarvedByDNSBacklog verifies direct provider matches are written
// even while the DNS queue is saturated with optional CNAME work.
func TestDirectNotStarvedByDNSBacklog(t *testing.T) {
	col := &collector{}
	release := make(chan struct{})
	p := New(Config{
		DNSQueueSize:  1,
		DNSWorkers:    1,
		Enricher:      newEnricher(&blockingResolver{release: release}),
		Write:         col.write,
		BatchSize:     10,
		FlushInterval: 10 * time.Millisecond,
	})
	ctx, cancel := context.WithCancel(context.Background())
	p.Start(ctx)
	for i := 0; i < 100; i++ {
		p.HandleHost("api" + strconv.Itoa(i) + ".example.com") // fills/blocks DNS queue
	}
	p.HandleHost("bucket.s3.eu-central-1.amazonaws.com")

	found := false
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		for _, f := range col.all() {
			if f.Hostname == "bucket.s3.eu-central-1.amazonaws.com" {
				found = true
			}
		}
		if found {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}
	close(release)
	cancel()
	p.Close()
	if !found {
		t.Fatal("direct match was not written while the DNS queue was saturated")
	}
}

func TestQueueSaturationDrops(t *testing.T) {
	m := metrics.New()
	p := New(Config{DirectQueueSize: 1, DNSQueueSize: 1, Metrics: m})
	// Intentionally do not Start: the queue stays full.
	p.HandleHost("bucket.s3.amazonaws.com")
	p.HandleHost("other.s3.amazonaws.com")
	if m.QueueDrops.Load() == 0 {
		t.Fatal("expected a drop when the direct queue is full")
	}
	if m.Candidates.Load() != 2 {
		t.Fatalf("candidates = %d, want 2", m.Candidates.Load())
	}
}

func TestDNSDedup(t *testing.T) {
	m := metrics.New()
	p := New(Config{DNSQueueSize: 8, Metrics: m})
	p.HandleHost("files.example.com")
	p.HandleHost("files.example.com")
	if got := p.dnsQ.Len(); got != 1 {
		t.Fatalf("dns queue len = %d, want 1", got)
	}
}

func TestGenericEndpointNotStored(t *testing.T) {
	col := &collector{}
	p := New(Config{Write: col.write, BatchSize: 10})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	p.Start(ctx)
	p.HandleIdentifiers(ctx, &certspotter.Identifiers{DNSNames: []string{
		"objectstorage.us-ashburn-1.oraclecloud.com",
	}})
	time.Sleep(30 * time.Millisecond)
	p.Close()
	if items := col.all(); len(items) != 0 {
		t.Fatalf("generic endpoint produced findings: %+v", items)
	}
}