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