bucketeer/internal/pipeline/pipeline.go

// Package pipeline wires the in-process CertSpotter identifier callback to the
// classifier, DNS enrichment and the batched PostgreSQL writer using bounded
// queues throughout.
package pipeline

import (
	"context"
	"net"
	"sync"
	"time"

	"software.sslmate.com/src/certspotter"

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

// Finding is a fully enriched candidate ready to be persisted.
type Finding struct {
	model.Candidate
	CNAMEChain []string
	IPv4       []string
	IPv6       []string
}

// WriteFunc persists a batch of findings.
type WriteFunc func(ctx context.Context, batch []Finding) error

type Config struct {
	DirectQueueSize int
	DNSQueueSize    int
	DBQueueSize     int
	DirectWorkers   int
	DNSWorkers      int
	BatchSize       int
	FlushInterval   time.Duration
	DedupTTL        time.Duration
	DedupSize       int

	Enricher *dns.Enricher
	Write    WriteFunc
	Metrics  *metrics.Metrics
}

// Pipeline is the certificate-to-database data flow. It is safe to call
// HandleIdentifiers from the CertSpotter monitor goroutines concurrently.
type Pipeline struct {
	cfg Config

	directQ chan model.Candidate
	dnsQ    chan string
	dbQ     chan Finding

	dedup    *dns.TTLCache[bool] // hostnames already queued for DNS
	seen     *dns.TTLCache[bool] // provider+service+provider_hostname already written
	producer sync.WaitGroup
	workers  sync.WaitGroup
}

func New(cfg Config) *Pipeline {
	if cfg.DirectQueueSize <= 0 {
		cfg.DirectQueueSize = 8192
	}
	if cfg.DNSQueueSize <= 0 {
		cfg.DNSQueueSize = 8192
	}
	if cfg.DBQueueSize <= 0 {
		cfg.DBQueueSize = 8192
	}
	if cfg.DirectWorkers <= 0 {
		cfg.DirectWorkers = 2
	}
	if cfg.DNSWorkers <= 0 {
		cfg.DNSWorkers = 16
	}
	if cfg.BatchSize <= 0 {
		cfg.BatchSize = 500
	}
	if cfg.FlushInterval <= 0 {
		cfg.FlushInterval = 2 * time.Second
	}
	if cfg.DedupTTL <= 0 {
		cfg.DedupTTL = time.Hour
	}
	if cfg.DedupSize <= 0 {
		cfg.DedupSize = 1 << 20
	}
	p := &Pipeline{
		cfg:     cfg,
		directQ: make(chan model.Candidate, cfg.DirectQueueSize),
		dnsQ:    make(chan string, cfg.DNSQueueSize),
		dbQ:     make(chan Finding, cfg.DBQueueSize),
		dedup:   dns.NewTTLCache[bool](cfg.DedupSize, 64),
		seen:    dns.NewTTLCache[bool](cfg.DedupSize, 64),
	}
	if p.cfg.Metrics != nil {
		p.cfg.Metrics.SetQueueDepth("direct", func() float64 { return float64(len(p.directQ)) })
		p.cfg.Metrics.SetQueueDepth("dns", func() float64 { return float64(len(p.dnsQ)) })
		p.cfg.Metrics.SetQueueDepth("db", func() float64 { return float64(len(p.dbQ)) })
	}
	return p
}

// Start launches the worker pools. Close must be called to drain.
func (p *Pipeline) Start(ctx context.Context) {
	for i := 0; i < p.cfg.DirectWorkers; i++ {
		p.producer.Add(1)
		go p.directWorker(ctx)
	}
	for i := 0; i < p.cfg.DNSWorkers; i++ {
		p.producer.Add(1)
		go p.dnsWorker(ctx)
	}
	p.workers.Add(1)
	go p.dbWorker(ctx)
}

// Close stops intake and drains all queues. It must only be called after the
// CertSpotter monitor has stopped calling HandleIdentifiers.
func (p *Pipeline) Close() {
	close(p.directQ)
	close(p.dnsQ)
	p.producer.Wait()
	close(p.dbQ)
	p.workers.Wait()
}

// HandleIdentifiers implements certspotter.IdentifierSink. It is deliberately
// allocation-light and never blocks: queues are bounded and overflow is dropped.
func (p *Pipeline) HandleIdentifiers(ctx context.Context, ids *certspotter.Identifiers) {
	if p.cfg.Metrics != nil {
		p.cfg.Metrics.CTCertificates.Add(1)
		p.cfg.Metrics.CTIdentifiers.Add(uint64(len(ids.DNSNames)))
	}
	for _, host := range ids.DNSNames {
		p.HandleHost(host)
	}
}

// HandleHost classifies a single hostname and enqueues it for direct handling
// or, when it does not match, for DNS/CNAME enrichment.
func (p *Pipeline) HandleHost(host string) {
	if m, ok := classify.Classify(host); ok {
		if !m.ResourceSpecific {
			// Generic service endpoints (bucket name in the path) are not
			// customer resources; do not create findings for them.
			return
		}
		if p.cfg.Metrics != nil {
			p.cfg.Metrics.DirectMatches.Add(1)
			p.cfg.Metrics.Candidates.Add(1)
		}
		if !p.markSeen(host) {
			return
		}
		cand := candidateFrom(host, host, m, model.SourceDirectCT)
		select {
		case p.directQ <- cand:
		default:
			p.drop()
		}
		return
	}
	if !classify.ShouldEnrich(host) {
		return
	}
	if _, dup := p.dedup.Get(host); dup {
		return
	}
	p.dedup.Set(host, true, p.cfg.DedupTTL)
	select {
	case p.dnsQ <- host:
	default:
		p.drop()
	}
}

func (p *Pipeline) drop() {
	if p.cfg.Metrics != nil {
		p.cfg.Metrics.QueueFull.Add(1)
		p.cfg.Metrics.QueueDrops.Add(1)
	}
}

func candidateFrom(host, original string, m model.Match, src model.Source) model.Candidate {
	return model.Candidate{
		Hostname:         host,
		Original:         original,
		Region:           m.Region,
		Resource:         m.Resource,
		Provider:         m.Provider,
		Service:          m.Service,
		Source:           src,
		PublicSemantic:   m.PublicSemantic,
		ResourceSpecific: m.ResourceSpecific,
	}
}

func appendIPs(f *Finding, ips []net.IP) {
	for _, ip := range ips {
		if v4 := ip.To4(); v4 != nil {
			f.IPv4 = append(f.IPv4, v4.String())
		} else {
			f.IPv6 = append(f.IPv6, ip.String())
		}
	}
}

func (p *Pipeline) enrich(ctx context.Context, c model.Candidate) Finding {
	f := Finding{Candidate: c}
	if p.cfg.Enricher == nil {
		return f
	}
	res, ok := p.cfg.Enricher.Resolve(ctx, c.Hostname)
	if !ok {
		return f
	}
	f.CNAMEChain = res.Chain
	appendIPs(&f, res.IPs)
	return f
}

func (p *Pipeline) directWorker(ctx context.Context) {
	defer p.producer.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case c, ok := <-p.directQ:
			if !ok {
				return
			}
			p.pushDB(ctx, p.enrich(ctx, c))
		}
	}
}

func (p *Pipeline) dnsWorker(ctx context.Context) {
	defer p.producer.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case host, ok := <-p.dnsQ:
			if !ok {
				return
			}
			if len(p.dbQ) >= cap(p.dbQ)/2 {
				// Load shedding: optional CNAME enrichment is dropped before
				// direct matches when the database writer is backed up.
				p.drop()
				continue
			}
			p.enrichHost(ctx, host)
		}
	}
}

func (p *Pipeline) enrichHost(ctx context.Context, host string) {
	res, ok := p.cfg.Enricher.Resolve(ctx, host)
	if !ok {
		return
	}
	for _, target := range res.Chain[1:] {
		m, matched := classify.Classify(target)
		if !matched || !m.ResourceSpecific {
			continue
		}
		if p.cfg.Metrics != nil {
			p.cfg.Metrics.CNAMEMatches.Add(1)
			p.cfg.Metrics.Candidates.Add(1)
		}
		if !p.markSeen(target) {
			continue
		}
		f := Finding{
			Candidate:  candidateFrom(target, host, m, model.SourceCNAME),
			CNAMEChain: res.Chain,
		}
		appendIPs(&f, res.IPs)
		p.pushDB(ctx, f)
	}
}

// markSeen dedups by hostname alone: a given hostname always classifies to the
// same provider/service, so no key concatenation (and no allocation) is needed.
func (p *Pipeline) markSeen(host string) bool {
	if _, dup := p.seen.Get(host); dup {
		return false
	}
	p.seen.Set(host, true, p.cfg.DedupTTL)
	return true
}

func (p *Pipeline) pushDB(ctx context.Context, f Finding) {
	select {
	case p.dbQ <- f:
	default:
		p.drop()
	}
}

func (p *Pipeline) dbWorker(ctx context.Context) {
	defer p.workers.Done()
	batch := make([]Finding, 0, p.cfg.BatchSize)
	ticker := time.NewTicker(p.cfg.FlushInterval)
	defer ticker.Stop()

	flush := func() {
		if len(batch) == 0 {
			return
		}
		if p.cfg.Write != nil {
			if err := p.cfg.Write(ctx, batch); err != nil {
				if p.cfg.Metrics != nil {
					p.cfg.Metrics.DBErrors.Add(1)
				}
			} else if p.cfg.Metrics != nil {
				p.cfg.Metrics.DBRows.Add(uint64(len(batch)))
			}
		}
		batch = batch[:0]
	}

	for {
		select {
		case <-ctx.Done():
			flush()
			return
		case f, ok := <-p.dbQ:
			if !ok {
				flush()
				return
			}
			batch = append(batch, f)
			if len(batch) >= p.cfg.BatchSize {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}