go-otel/internal/storage/csp.go

package storage

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
)

// CSPReport is one Content-Security-Policy violation report, normalized from
// either the legacy report-uri format ({"csp-report": {...}}) or the newer
// Reporting API (application/reports+json).
type CSPReport struct {
	// ID is a stable content hash, computed at ingest, usable as a UI key.
	ID                 string    `json:"id"`
	Timestamp          time.Time `json:"timestamp"`
	Site               string    `json:"site"` // host of document_url, e.g. app.k6n.net
	DocumentURL        string    `json:"documentUrl"`
	Referrer           string    `json:"referrer,omitempty"`
	BlockedURL         string    `json:"blockedUrl,omitempty"`
	ViolatedDirective  string    `json:"violatedDirective,omitempty"`
	EffectiveDirective string    `json:"effectiveDirective,omitempty"`
	OriginalPolicy     string    `json:"originalPolicy,omitempty"`
	Disposition        string    `json:"disposition,omitempty"` // enforce | report
	StatusCode         int       `json:"statusCode,omitempty"`
	Sample             string    `json:"sample,omitempty"`
	// ReportType is the Reporting API type ("csp-violation"); empty for
	// legacy report-uri payloads.
	ReportType string `json:"reportType,omitempty"`
	UserAgent  string `json:"userAgent,omitempty"`
	SourceIP   string `json:"sourceIP,omitempty"`
}

// CSPReportsQuery filters the stored reports for the UI list.
type CSPReportsQuery struct {
	// Site filters on the document host ("app.k6n.net"); empty = all.
	Site string
	// Directive filters on the effective directive ("script-src"); empty = all.
	Directive string
	// Search is a substring match across urls, referrer and sample.
	Search string
	// Since (zero = all time) drops reports older than the given instant.
	Since time.Time
	Limit int
}

// CSPCount is one named bucket of the summary (directive, host, site).
type CSPCount struct {
	Name  string `json:"name"`
	Count uint64 `json:"count"`
}

// CSPSummary aggregates reports over one window for the view's overview cards.
type CSPSummary struct {
	Total        uint64     `json:"total"`
	Sites        uint64     `json:"sites"`
	ByDirective  []CSPCount `json:"byDirective"`
	ByBlockedURL []CSPCount `json:"byBlockedHost"`
	BySite       []CSPCount `json:"bySite"`
}

// CSPStore persists CSP violation reports and answers the view's queries.
// ClickHouseStore implements it; CSPFileStore is the fallback when no
// ClickHouse is configured.
type CSPStore interface {
	AppendCSPReport(ctx context.Context, r CSPReport) error
	CSPReports(ctx context.Context, q CSPReportsQuery) ([]CSPReport, error)
	CSPSummary(ctx context.Context, hours int) (CSPSummary, error)
}

// EnsureCSPSchema creates the dedicated csp_reports table: one row per
// browser report, typed columns instead of a JSON blob so the view can
// aggregate server-side. Idempotent.
func EnsureCSPSchema(ctx context.Context, conn driver.Conn, db string) error {
	return conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.csp_reports (
		id                  String,
		timestamp           DateTime64(9, 'UTC'),
		site                LowCardinality(String),
		document_url        String,
		referrer            String,
		blocked_url         String,
		violated_directive  LowCardinality(String),
		effective_directive LowCardinality(String),
		original_policy     String,
		disposition         LowCardinality(String),
		status_code         UInt16,
		sample              String,
		report_type         LowCardinality(String),
		user_agent          String,
		source_ip           String
	) ENGINE = MergeTree
	PARTITION BY toYYYYMM(timestamp)
	ORDER BY (site, effective_directive, timestamp)
	TTL toDateTime(timestamp) + INTERVAL 90 DAY`, db))
}

func (s *ClickHouseStore) AppendCSPReport(ctx context.Context, r CSPReport) error {
	return s.conn.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.csp_reports
		(id, timestamp, site, document_url, referrer, blocked_url, violated_directive,
		 effective_directive, original_policy, disposition, status_code, sample,
		 report_type, user_agent, source_ip) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`, s.db),
		r.ID, r.Timestamp.UTC(), r.Site, r.DocumentURL, r.Referrer, r.BlockedURL,
		r.ViolatedDirective, r.EffectiveDirective, r.OriginalPolicy, r.Disposition,
		r.StatusCode, r.Sample, r.ReportType, r.UserAgent, r.SourceIP)
}

func (s *ClickHouseStore) CSPReports(ctx context.Context, q CSPReportsQuery) ([]CSPReport, error) {
	if q.Limit <= 0 || q.Limit > 1000 {
		q.Limit = 200
	}
	where, args := "1", []any{}
	if q.Site != "" {
		where += " AND site = ?"
		args = append(args, q.Site)
	}
	if q.Directive != "" {
		where += " AND effective_directive = ?"
		args = append(args, q.Directive)
	}
	if q.Search != "" {
		where += " AND (position(document_url, ?) > 0 OR position(blocked_url, ?) > 0 OR position(referrer, ?) > 0 OR position(sample, ?) > 0)"
		for i := 0; i < 4; i++ {
			args = append(args, q.Search)
		}
	}
	if !q.Since.IsZero() {
		where += " AND timestamp >= ?"
		args = append(args, q.Since.UTC())
	}
	args = append(args, q.Limit)
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT id, timestamp, site, document_url, referrer, blocked_url,
		violated_directive, effective_directive, original_policy, disposition, status_code,
		sample, report_type, user_agent, source_ip
		FROM %s.csp_reports WHERE %s ORDER BY timestamp DESC LIMIT ?`, s.db, where), args...)
	if err != nil {
		return nil, fmt.Errorf("clickhouse csp reports: %w", err)
	}
	defer rows.Close()
	out := []CSPReport{}
	for rows.Next() {
		var r CSPReport
		var code uint16
		if err := rows.Scan(&r.ID, &r.Timestamp, &r.Site, &r.DocumentURL, &r.Referrer,
			&r.BlockedURL, &r.ViolatedDirective, &r.EffectiveDirective, &r.OriginalPolicy,
			&r.Disposition, &code, &r.Sample, &r.ReportType, &r.UserAgent, &r.SourceIP); err != nil {
			return nil, err
		}
		r.StatusCode = int(code)
		out = append(out, r)
	}
	return out, rows.Err()
}

func (s *ClickHouseStore) CSPSummary(ctx context.Context, hours int) (CSPSummary, error) {
	if hours <= 0 {
		hours = 24
	}
	since := time.Now().UTC().Add(-time.Duration(hours) * time.Hour)
	var sum CSPSummary
	if err := s.conn.QueryRow(ctx, fmt.Sprintf(
		`SELECT count(), uniqExact(site) FROM %s.csp_reports WHERE timestamp >= ?`, s.db),
		since).Scan(&sum.Total, &sum.Sites); err != nil {
		return sum, fmt.Errorf("clickhouse csp summary: %w", err)
	}
	counts := func(expr string) ([]CSPCount, error) {
		rows, err := s.conn.Query(ctx, fmt.Sprintf(
			`SELECT %s AS name, count() AS c FROM %s.csp_reports
			 WHERE timestamp >= ? GROUP BY name ORDER BY c DESC LIMIT 20`, expr, s.db), since)
		if err != nil {
			return nil, err
		}
		defer rows.Close()
		var out []CSPCount
		for rows.Next() {
			var v CSPCount
			if err := rows.Scan(&v.Name, &v.Count); err != nil {
				return nil, err
			}
			out = append(out, v)
		}
		return out, rows.Err()
	}
	var err error
	if sum.ByDirective, err = counts("effective_directive"); err != nil {
		return sum, err
	}
	if sum.ByBlockedURL, err = counts("domain(blocked_url)"); err != nil {
		return sum, err
	}
	if sum.BySite, err = counts("site"); err != nil {
		return sum, err
	}
	return sum, nil
}

// URLLowerHost extracts the lowercase host (port stripped) from a URL-ish
// string without failing on scheme-relative or malformed input; shared by
// ingest normalization and the file-store summary.
func URLLowerHost(raw string) string {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return ""
	}
	if i := strings.IndexAny(raw, "?#"); i != -1 {
		raw = raw[:i]
	}
	raw = strings.TrimPrefix(strings.TrimPrefix(raw, "https://"), "http://")
	raw = strings.TrimPrefix(raw, "//")
	if i := strings.IndexByte(raw, '/'); i != -1 {
		raw = raw[:i]
	}
	if i := strings.LastIndexByte(raw, '@'); i != -1 {
		raw = raw[i+1:]
	}
	if i := strings.LastIndexByte(raw, ':'); i != -1 && strings.LastIndexByte(raw, ']') < i {
		raw = raw[:i]
	}
	return strings.ToLower(raw)
}

// CSPFileStore is the no-ClickHouse fallback: newline-delimited JSON next to
// the telemetry log, read back with in-memory filtering.
type CSPFileStore struct {
	filename string
	mu       sync.Mutex
}

func NewCSPFileStore(filename string) (*CSPFileStore, error) {
	if dir := filepath.Dir(filename); dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return nil, err
		}
	}
	return &CSPFileStore{filename: filename}, nil
}

func (fs *CSPFileStore) AppendCSPReport(_ context.Context, r CSPReport) error {
	fs.mu.Lock()
	defer fs.mu.Unlock()
	f, err := os.OpenFile(fs.filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer f.Close()
	return json.NewEncoder(f).Encode(r)
}

// readAll loads every stored report, oldest first.
func (fs *CSPFileStore) readAll() ([]CSPReport, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()
	f, err := os.Open(fs.filename)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	defer f.Close()
	var out []CSPReport
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		var r CSPReport
		if err := json.Unmarshal([]byte(line), &r); err == nil {
			out = append(out, r)
		}
	}
	return out, sc.Err()
}

func (fs *CSPFileStore) CSPReports(_ context.Context, q CSPReportsQuery) ([]CSPReport, error) {
	if q.Limit <= 0 || q.Limit > 1000 {
		q.Limit = 200
	}
	all, err := fs.readAll()
	if err != nil {
		return nil, err
	}
	match := func(hay, needle string) bool { return needle == "" || strings.Contains(hay, needle) }
	out := []CSPReport{}
	for _, r := range all {
		if q.Site != "" && r.Site != q.Site {
			continue
		}
		if q.Directive != "" && r.EffectiveDirective != q.Directive {
			continue
		}
		if !q.Since.IsZero() && r.Timestamp.Before(q.Since) {
			continue
		}
		if q.Search != "" &&
			!match(r.DocumentURL, q.Search) && !match(r.BlockedURL, q.Search) &&
			!match(r.Referrer, q.Search) && !match(r.Sample, q.Search) {
			continue
		}
		out = append(out, r)
	}
	// Newest first, bounded.
	sort.Slice(out, func(i, j int) bool { return out[i].Timestamp.After(out[j].Timestamp) })
	if len(out) > q.Limit {
		out = out[:q.Limit]
	}
	return out, nil
}

func (fs *CSPFileStore) CSPSummary(_ context.Context, hours int) (CSPSummary, error) {
	if hours <= 0 {
		hours = 24
	}
	since := time.Now().Add(-time.Duration(hours) * time.Hour)
	all, err := fs.readAll()
	if err != nil {
		return CSPSummary{}, err
	}
	var sum CSPSummary
	directive, blocked, site := map[string]uint64{}, map[string]uint64{}, map[string]uint64{}
	sites := map[string]bool{}
	add := func(m map[string]uint64, name string) {
		if name != "" {
			m[name]++
		}
	}
	for _, r := range all {
		if r.Timestamp.Before(since) {
			continue
		}
		sum.Total++
		sites[r.Site] = true
		add(directive, r.EffectiveDirective)
		add(blocked, URLLowerHost(r.BlockedURL))
		add(site, r.Site)
	}
	sum.Sites = uint64(len(sites))
	sum.ByDirective = topCounts(directive)
	sum.ByBlockedURL = topCounts(blocked)
	sum.BySite = topCounts(site)
	return sum, nil
}

// topCounts sorts a counter map, heaviest first, capped like the ClickHouse
// summary so both stores return comparable payloads.
func topCounts(m map[string]uint64) []CSPCount {
	out := make([]CSPCount, 0, len(m))
	for name, count := range m {
		out = append(out, CSPCount{Name: name, Count: count})
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].Count > out[j-1].Count; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	if len(out) > 20 {
		out = out[:20]
	}
	return out
}