bucketeer/internal/store/store.go

// Package store owns the PostgreSQL schema, the batched findings writer and the
// restart-safe CT checkpoint state provider.
package store

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

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

// Store wraps a pgx pool with the operations ctcloud needs.
type Store struct {
	pool    *pgxpool.Pool
	metrics *metrics.Metrics
}

// Open connects to PostgreSQL. The pool is bounded; the caller configures the
// maximum number of connections.
func Open(ctx context.Context, databaseURL string, maxConns int) (*Store, error) {
	cfg, err := pgxpool.ParseConfig(databaseURL)
	if err != nil {
		return nil, fmt.Errorf("invalid DATABASE_URL: %w", err)
	}
	if maxConns > 0 {
		cfg.MaxConns = int32(maxConns)
	}
	cfg.MaxConnLifetime = 30 * time.Minute
	cfg.MinConns = 1
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, err
	}
	return &Store{pool: pool}, nil
}

func (s *Store) SetMetrics(m *metrics.Metrics) { s.metrics = m }

func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }

func (s *Store) Close() { s.pool.Close() }

func (s *Store) Pool() *pgxpool.Pool { return s.pool }

// Migrate applies the embedded schema idempotently.
func (s *Store) Migrate(ctx context.Context) error {
	for _, name := range migrations.SQL() {
		sqlBytes, err := migrations.FS.ReadFile(name)
		if err != nil {
			return err
		}
		if _, err := s.pool.Exec(ctx, string(sqlBytes)); err != nil {
			return fmt.Errorf("applying %s: %w", name, err)
		}
	}
	return nil
}

const upsertFindingSQL = `
WITH asset AS (
    INSERT INTO cloud_assets
        (provider, service, hostname, provider_hostname, resource_identifier,
         region, resource_specific, public_semantic, first_seen, last_seen)
    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now(), now())
    ON CONFLICT (provider, service, provider_hostname)
    DO UPDATE SET
        last_seen           = now(),
        hostname            = EXCLUDED.hostname,
        resource_identifier = CASE WHEN EXCLUDED.resource_identifier = '' THEN cloud_assets.resource_identifier ELSE EXCLUDED.resource_identifier END,
        region              = CASE WHEN EXCLUDED.region = '' THEN cloud_assets.region ELSE EXCLUDED.region END
    RETURNING id
)
INSERT INTO observations
    (asset_id, source_hostname, source_type, cname_chain, ipv4, ipv6)
SELECT asset.id, $9, $10, $11, $12, $13 FROM asset`

// WriteFindings writes one batch of findings in a single round trip. Asset rows
// are upserted and an observation is appended for each finding.
func (s *Store) WriteFindings(ctx context.Context, batch []pipeline.Finding) error {
	if len(batch) == 0 {
		return nil
	}
	var lastErr error
	// ponytail: three quick attempts; the pipeline keeps running and the
	// bounded DB queue sheds load if PostgreSQL stays down.
	for attempt := 0; attempt < 3; attempt++ {
		if err := s.writeBatch(ctx, batch); err != nil {
			lastErr = err
			if ctx.Err() != nil {
				return err
			}
			backoff := time.Duration(1<<attempt) * 100 * time.Millisecond
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(backoff):
			}
			continue
		}
		return nil
	}
	return lastErr
}

func (s *Store) writeBatch(ctx context.Context, batch []pipeline.Finding) error {
	b := &pgx.Batch{}
	for i := range batch {
		f := &batch[i]
		b.Queue(upsertFindingSQL,
			f.Provider.String(), f.Service.String(), f.Original, f.Hostname,
			f.Resource, f.Region, f.ResourceSpecific, f.PublicSemantic.String(),
			f.Original, f.Source.String(),
			nonNil(f.CNAMEChain), nonNil(f.IPv4), nonNil(f.IPv6),
		)
	}
	br := s.pool.SendBatch(ctx, b)
	defer br.Close()
	for range batch {
		if _, err := br.Exec(); err != nil {
			return err
		}
	}
	return nil
}

func nonNil(s []string) []string {
	if s == nil {
		return []string{}
	}
	return s
}

// UpdateExposure records the outcome of an optional reachability probe.
func (s *Store) UpdateExposure(ctx context.Context, provider, service, providerHostname string, state model.ExposureState, detail string) error {
	_, err := s.pool.Exec(ctx, `
		UPDATE cloud_assets
		SET exposure_state = $4, exposure_detail = $5, probed_at = now()
		WHERE provider = $1 AND service = $2 AND provider_hostname = $3`,
		provider, service, providerHostname, state.String(), detail)
	return err
}

// Asset is a row of the cloud_assets table for the read-only UI.
type Asset struct {
	ID             int64      `json:"id"`
	Provider       string     `json:"provider"`
	Service        string     `json:"service"`
	Hostname       string     `json:"hostname"`
	ProviderHost   string     `json:"provider_hostname"`
	Resource       string     `json:"resource_identifier"`
	Region         string     `json:"region"`
	PublicSemantic string     `json:"public_semantic"`
	Exposure       string     `json:"exposure_state"`
	ExposureDetail string     `json:"exposure_detail"`
	FirstSeen      time.Time  `json:"first_seen"`
	LastSeen       time.Time  `json:"last_seen"`
	ProbedAt       *time.Time `json:"probed_at"`
}

// AssetFilter selects assets for the UI. Empty fields are ignored.
type AssetFilter struct {
	Provider string
	Service  string
	Exposure string
	Query    string
	Limit    int
	Offset   int
}

// ListAssets returns assets matching the filter, newest first.
func (s *Store) ListAssets(ctx context.Context, f AssetFilter) ([]Asset, error) {
	if f.Limit <= 0 || f.Limit > 500 {
		f.Limit = 100
	}
	rows, err := s.pool.Query(ctx, `
		SELECT id, provider, service, hostname, provider_hostname, resource_identifier,
		       region, public_semantic, exposure_state, exposure_detail, first_seen, last_seen, probed_at
		FROM cloud_assets
		WHERE ($1 = '' OR provider = $1)
		  AND ($2 = '' OR service = $2)
		  AND ($3 = '' OR exposure_state = $3)
		  AND ($4 = '' OR provider_hostname ILIKE '%' || $4 || '%' OR hostname ILIKE '%' || $4 || '%' OR resource_identifier ILIKE '%' || $4 || '%')
		ORDER BY last_seen DESC
		LIMIT $5 OFFSET $6`,
		f.Provider, f.Service, f.Exposure, f.Query, f.Limit, f.Offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []Asset{}
	for rows.Next() {
		var a Asset
		if err := rows.Scan(&a.ID, &a.Provider, &a.Service, &a.Hostname, &a.ProviderHost,
			&a.Resource, &a.Region, &a.PublicSemantic, &a.Exposure, &a.ExposureDetail,
			&a.FirstSeen, &a.LastSeen, &a.ProbedAt); err != nil {
			return nil, err
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

// Count is a (label, count) pair for the UI summary.
type Count struct {
	Label string `json:"label"`
	Count int64  `json:"count"`
}

// Stats summarizes the assets table.
type Stats struct {
	Total      int64   `json:"total"`
	Readable   int64   `json:"readable"`
	ByProvider []Count `json:"by_provider"`
	ByService  []Count `json:"by_service"`
	ByExposure []Count `json:"by_exposure"`
}

func (s *Store) Stats(ctx context.Context) (Stats, error) {
	var st Stats
	if err := s.pool.QueryRow(ctx, `SELECT count(*), count(*) FILTER (WHERE exposure_state = 'public_read') FROM cloud_assets`).Scan(&st.Total, &st.Readable); err != nil {
		return st, err
	}
	group := func(column string) ([]Count, error) {
		rows, err := s.pool.Query(ctx, `SELECT `+column+`, count(*) FROM cloud_assets GROUP BY 1 ORDER BY 2 DESC`)
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var counts []Count
		for rows.Next() {
			var c Count
			if err := rows.Scan(&c.Label, &c.Count); err != nil {
				return nil, err
			}
			counts = append(counts, c)
		}
		return counts, rows.Err()
	}
	var err error
	if st.ByProvider, err = group("provider"); err != nil {
		return st, err
	}
	if st.ByService, err = group("service"); err != nil {
		return st, err
	}
	if st.ByExposure, err = group("exposure_state"); err != nil {
		return st, err
	}
	return st, nil
}