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 }

// PoolStats reports pgx pool utilization for logging/metrics.
func (s *Store) PoolStats() (total, idle, acquired int32) {
	st := s.pool.Stat()
	return st.TotalConns(), st.IdleConns(), st.AcquiredConns()
}

// 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
}

// IndexEntry is one indexed resource item for the UI.
type IndexEntry struct {
	ID           int64      `json:"id"`
	AssetID      int64      `json:"asset_id"`
	Provider     string     `json:"provider"`
	Service      string     `json:"service"`
	ProviderHost string     `json:"provider_hostname"`
	Kind         string     `json:"kind"`
	Name         string     `json:"name"`
	Size         int64      `json:"size"`
	Updated      *time.Time `json:"updated"`
	ContentType  string     `json:"content_type"`
	IndexedAt    time.Time  `json:"indexed_at"`
}

// ReplaceEntries replaces the stored index for a readable asset. Returns the
// number of entries stored; a no-op if the asset row is unknown or entries is
// empty.
func (s *Store) ReplaceEntries(ctx context.Context, provider, service, host string, entries []model.Entry) (int, error) {
	if len(entries) == 0 {
		return 0, nil
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback(ctx)

	var assetID int64
	err = tx.QueryRow(ctx,
		`SELECT id FROM cloud_assets WHERE provider = $1 AND service = $2 AND provider_hostname = $3`,
		provider, service, host).Scan(&assetID)
	if err == pgx.ErrNoRows {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	if _, err := tx.Exec(ctx, `DELETE FROM asset_entries WHERE asset_id = $1`, assetID); err != nil {
		return 0, err
	}
	b := &pgx.Batch{}
	for i := range entries {
		e := &entries[i]
		var updated any
		if !e.Updated.IsZero() {
			updated = e.Updated
		}
		b.Queue(`
			INSERT INTO asset_entries (asset_id, kind, name, size, updated, content_type, etag)
			VALUES ($1, $2, $3, $4, $5, $6, $7)
			ON CONFLICT (asset_id, kind, name) DO UPDATE
			SET size = EXCLUDED.size, updated = EXCLUDED.updated,
			    content_type = EXCLUDED.content_type, etag = EXCLUDED.etag, indexed_at = now()`,
			assetID, e.Kind, e.Name, e.Size, updated, e.ContentType, e.ETag)
	}
	br := tx.SendBatch(ctx, b)
	for range entries {
		if _, err := br.Exec(); err != nil {
			br.Close()
			return 0, err
		}
	}
	if err := br.Close(); err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return len(entries), nil
}

// SearchEntries searches the index across all readable assets, newest first.
func (s *Store) SearchEntries(ctx context.Context, query string, limit, offset int) ([]IndexEntry, int, error) {
	if limit <= 0 || limit > 500 {
		limit = 100
	}
	rows, err := s.pool.Query(ctx, `
		SELECT e.id, e.asset_id, a.provider, a.service, a.provider_hostname,
		       e.kind, e.name, e.size, e.updated, e.content_type, e.indexed_at
		FROM asset_entries e JOIN cloud_assets a ON a.id = e.asset_id
		WHERE ($1 = '' OR e.name ILIKE '%' || $1 || '%' OR a.provider_hostname ILIKE '%' || $1 || '%')
		ORDER BY e.indexed_at DESC, e.id DESC
		LIMIT $2 OFFSET $3`, query, limit, offset)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	out := []IndexEntry{}
	for rows.Next() {
		var e IndexEntry
		if err := rows.Scan(&e.ID, &e.AssetID, &e.Provider, &e.Service, &e.ProviderHost,
			&e.Kind, &e.Name, &e.Size, &e.Updated, &e.ContentType, &e.IndexedAt); err != nil {
			return nil, 0, err
		}
		out = append(out, e)
	}
	if err := rows.Err(); err != nil {
		return nil, 0, err
	}
	var total int
	if err := s.pool.QueryRow(ctx, `
		SELECT count(*) FROM asset_entries e JOIN cloud_assets a ON a.id = e.asset_id
		WHERE ($1 = '' OR e.name ILIKE '%' || $1 || '%' OR a.provider_hostname ILIKE '%' || $1 || '%')`,
		query).Scan(&total); err != nil {
		return nil, 0, err
	}
	return out, total, nil
}

// CountAssets returns the number of assets matching the filter (for pagination).
func (s *Store) CountAssets(ctx context.Context, f AssetFilter) (int, error) {
	var n int
	err := s.pool.QueryRow(ctx, `
		SELECT count(*) 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 || '%')`,
		f.Provider, f.Service, f.Exposure, f.Query).Scan(&n)
	return n, err
}

// 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
	Sort     string
	Limit    int
	Offset   int
}

// orderClause maps a UI sort key to a fixed ORDER BY, preventing injection.
func orderClause(sort string) string {
	switch sort {
	case "first_seen":
		return "first_seen DESC"
	case "provider":
		return "provider ASC, service ASC, provider_hostname ASC"
	case "hostname":
		return "provider_hostname ASC"
	case "exposure":
		return "CASE exposure_state WHEN 'public_read' THEN 0 WHEN 'reachable' THEN 1 WHEN 'auth_required' THEN 2 WHEN 'not_found' THEN 3 WHEN 'private' THEN 4 ELSE 5 END, last_seen DESC"
	default:
		return "last_seen DESC"
	}
}

// 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 `+orderClause(f.Sort)+`
		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"`
	Reachable  int64   `json:"reachable"`
	Entries    int64   `json:"entries"`
	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'),
		       count(*) FILTER (WHERE exposure_state = 'reachable')
		FROM cloud_assets`).Scan(&st.Total, &st.Readable, &st.Reachable); err != nil {
		return st, err
	}
	if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM asset_entries`).Scan(&st.Entries); 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
}