path: root/internal/store/store.go

// Package store persists proxy history to SQLite in WAL mode. Request
// and response bytes are stored as-received where possible (see the
// Exact fields) rather than re-serialized from a parsed representation.
// An FTS5 index (history_fts) mirrors method/host/path and the raw
// request/response text for full-text search - see Search.
package store

import (
	"database/sql"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	_ "modernc.org/sqlite"

	"mitmux/internal/rules"
)

const schema = `
CREATE TABLE IF NOT EXISTS history (
	id             INTEGER PRIMARY KEY AUTOINCREMENT,
	started_at     INTEGER NOT NULL,
	duration_ms    INTEGER NOT NULL,
	method         TEXT NOT NULL,
	scheme         TEXT NOT NULL,
	host           TEXT NOT NULL,
	path           TEXT NOT NULL,
	status_code    INTEGER,
	request_raw    BLOB NOT NULL,
	response_raw   BLOB,
	request_exact  INTEGER NOT NULL,
	response_exact INTEGER NOT NULL,
	error          TEXT NOT NULL DEFAULT '',
	source         TEXT NOT NULL DEFAULT 'proxy',
	flagged        INTEGER NOT NULL DEFAULT 0,
	request_truncated  INTEGER NOT NULL DEFAULT 0,
	response_truncated INTEGER NOT NULL DEFAULT 0
);

CREATE VIRTUAL TABLE IF NOT EXISTS history_fts USING fts5(
	method, host, path, request_text, response_text,
	tokenize = 'unicode61 remove_diacritics 2'
);

CREATE TABLE IF NOT EXISTS rules (
	id       INTEGER PRIMARY KEY AUTOINCREMENT,
	enabled  INTEGER NOT NULL DEFAULT 1,
	name     TEXT NOT NULL DEFAULT '',
	scope    TEXT NOT NULL,
	part     TEXT NOT NULL,
	match    TEXT NOT NULL,
	replace  TEXT NOT NULL,
	is_regex INTEGER NOT NULL DEFAULT 0,
	position INTEGER NOT NULL DEFAULT 0
);
`

// Store is a handle to the history database. Safe for concurrent use.
type Store struct {
	db *sql.DB
}

// Open opens (creating if needed) the SQLite database at path in WAL mode.
func Open(path string) (*Store, error) {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("open db: %w", err)
	}
	// modernc.org/sqlite has no real connection pooling benefit here and
	// SQLite only supports one writer at a time; serializing access
	// through a single connection avoids SQLITE_BUSY entirely.
	db.SetMaxOpenConns(1)

	for _, pragma := range []string{
		"PRAGMA journal_mode = WAL",
		"PRAGMA synchronous = NORMAL",
		"PRAGMA foreign_keys = ON",
	} {
		if _, err := db.Exec(pragma); err != nil {
			db.Close()
			return nil, fmt.Errorf("%s: %w", pragma, err)
		}
	}
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("create schema: %w", err)
	}
	// Added after the initial schema; ignore the "duplicate column" error
	// on databases that already have it.
	db.Exec("ALTER TABLE history ADD COLUMN source TEXT NOT NULL DEFAULT 'proxy'")
	db.Exec("ALTER TABLE history ADD COLUMN flagged INTEGER NOT NULL DEFAULT 0")
	db.Exec("ALTER TABLE history ADD COLUMN request_truncated INTEGER NOT NULL DEFAULT 0")
	db.Exec("ALTER TABLE history ADD COLUMN response_truncated INTEGER NOT NULL DEFAULT 0")
	// Backfill history_fts for rows inserted before it existed. A no-op
	// once caught up, since every Insert keeps both tables in sync.
	if _, err := db.Exec(`
		INSERT INTO history_fts (rowid, method, host, path, request_text, response_text)
		SELECT id, method, host, path, CAST(request_raw AS TEXT), COALESCE(CAST(response_raw AS TEXT), '')
		FROM history WHERE id NOT IN (SELECT rowid FROM history_fts)
	`); err != nil {
		db.Close()
		return nil, fmt.Errorf("backfill search index: %w", err)
	}
	return &Store{db: db}, nil
}

// Close closes the underlying database.
func (s *Store) Close() error {
	return s.db.Close()
}

// Entry is one captured request/response pair.
type Entry struct {
	ID            int64
	StartedAt     time.Time
	Duration      time.Duration
	Method        string
	Scheme        string
	Host          string
	Path          string
	StatusCode    int // 0 if no response was received
	RequestRaw    []byte
	ResponseRaw   []byte // nil if no response was received
	RequestExact  bool   // true if RequestRaw is wire-exact, false if reconstructed (e.g. HTTP/2)
	ResponseExact bool
	// Truncated is true when the corresponding *Raw field would have
	// been an exact capture but hit maxCaptureBytes and had bytes
	// dropped off the end - distinct from a false *Exact, which also
	// covers HTTP/2's inherently-reconstructed (never wire-exact to
	// begin with) captures. Only meaningful when the matching *Exact
	// field is false; a capture can't be both exact and truncated.
	RequestTruncated  bool
	ResponseTruncated bool
	Error             string // network/transport error, if the request never got a response
	Source            string // "proxy" or "repeater"
	Flagged           bool
}

// Summary is the lightweight metadata used for the history list view -
// no request/response bodies.
type Summary struct {
	ID         int64
	StartedAt  time.Time
	Duration   time.Duration
	Method     string
	Scheme     string
	Host       string
	Path       string
	StatusCode int
	ReqSize    int
	RespSize   int
	Error      string
	Source     string
	Flagged    bool
}

// Insert stores e (and indexes it for search) and returns its assigned ID.
func (s *Store) Insert(e *Entry) (int64, error) {
	var statusCode any
	if e.StatusCode != 0 {
		statusCode = e.StatusCode
	}
	source := e.Source
	if source == "" {
		source = "proxy"
	}

	tx, err := s.db.Begin()
	if err != nil {
		return 0, fmt.Errorf("begin insert: %w", err)
	}
	defer tx.Rollback()

	res, err := tx.Exec(
		`INSERT INTO history
			(started_at, duration_ms, method, scheme, host, path, status_code,
			 request_raw, response_raw, request_exact, response_exact, error, source,
			 request_truncated, response_truncated)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		e.StartedAt.UnixMilli(), e.Duration.Milliseconds(), e.Method, e.Scheme, e.Host, e.Path,
		statusCode, e.RequestRaw, e.ResponseRaw, boolToInt(e.RequestExact), boolToInt(e.ResponseExact), e.Error, source,
		boolToInt(e.RequestTruncated), boolToInt(e.ResponseTruncated),
	)
	if err != nil {
		return 0, fmt.Errorf("insert history entry: %w", err)
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("insert history entry: %w", err)
	}

	if _, err := tx.Exec(
		`INSERT INTO history_fts (rowid, method, host, path, request_text, response_text)
		 VALUES (?, ?, ?, ?, ?, ?)`,
		id, e.Method, e.Host, e.Path, string(e.RequestRaw), string(e.ResponseRaw),
	); err != nil {
		return 0, fmt.Errorf("index history entry: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return 0, fmt.Errorf("commit insert: %w", err)
	}
	return id, nil
}

// Count returns the total number of history entries.
func (s *Store) Count() (int64, error) {
	var n int64
	if err := s.db.QueryRow(`SELECT count(*) FROM history`).Scan(&n); err != nil {
		return 0, fmt.Errorf("count history: %w", err)
	}
	return n, nil
}

// List returns up to limit history summaries older than beforeID (or the
// most recent if beforeID is 0), newest first.
func (s *Store) List(limit int, beforeID int64) ([]Summary, error) {
	if limit <= 0 || limit > 1000 {
		limit = 200
	}
	if beforeID <= 0 {
		beforeID = 1<<63 - 1
	}
	rows, err := s.db.Query(
		`SELECT id, started_at, duration_ms, method, scheme, host, path,
		        COALESCE(status_code, 0), length(request_raw), COALESCE(length(response_raw), 0), error, source, flagged
		 FROM history WHERE id < ? ORDER BY id DESC LIMIT ?`,
		beforeID, limit,
	)
	if err != nil {
		return nil, fmt.Errorf("list history: %w", err)
	}
	defer rows.Close()

	var out []Summary
	for rows.Next() {
		var sum Summary
		var startedAt, durationMs int64
		var flagged int
		if err := rows.Scan(&sum.ID, &startedAt, &durationMs, &sum.Method, &sum.Scheme, &sum.Host, &sum.Path,
			&sum.StatusCode, &sum.ReqSize, &sum.RespSize, &sum.Error, &sum.Source, &flagged); err != nil {
			return nil, fmt.Errorf("scan history row: %w", err)
		}
		sum.Flagged = flagged != 0
		sum.StartedAt = time.UnixMilli(startedAt)
		sum.Duration = time.Duration(durationMs) * time.Millisecond
		out = append(out, sum)
	}
	return out, rows.Err()
}

// Search returns up to limit history summaries older than beforeID (or
// the most recent if beforeID is 0) matching query, newest first among
// equally-ranked matches, best-relevance first otherwise. query is
// mostly plain text - "example.com", "x-forwarded-for", "192.168.1.1"
// all just work - plus a little FTS5 syntax for power users: AND/OR/NOT,
// and column filters like host:example.com. See prepareFTSQuery for why
// plain terms need help: FTS5's query grammar treats characters like
// '.', '-', '/', '@' as syntax, so an unquoted domain name is a parse
// error, not a search term, unless quoted first.
//
// Two more filters are recognized ahead of the FTS5 layer, since neither
// fits it - status_code isn't a text column FTS5 can index, and doesn't
// benefit from full-text matching (numeric comparison, not word search),
// and source is a plain low-cardinality column better matched exactly:
// status:404, status:>=400, status:!=200, status:4xx (also 2xx/3xx/5xx),
// and source:proxy / source:repeater / source:intruder. These combine
// with AND against any remaining free-text/FTS5 portion of the query.
func (s *Store) Search(query string, limit int, beforeID int64) ([]Summary, error) {
	if limit <= 0 || limit > 1000 {
		limit = 200
	}
	if beforeID <= 0 {
		beforeID = 1<<63 - 1
	}

	remaining, pred := extractStructured(query)
	where := []string{"h.id < ?"}
	args := []any{beforeID}
	if pred.statusSQL != "" {
		where = append(where, pred.statusSQL)
		args = append(args, pred.statusArgs...)
	}
	if pred.source != "" {
		where = append(where, "h.source = ?")
		args = append(args, pred.source)
	}
	if pred.flagged != nil {
		where = append(where, "h.flagged = ?")
		args = append(args, boolToInt(*pred.flagged))
	}

	var q string
	if remaining == "" {
		// No free-text component left - query history directly, no
		// FTS5 join or ranking needed.
		q = `SELECT h.id, h.started_at, h.duration_ms, h.method, h.scheme, h.host, h.path,
		            COALESCE(h.status_code, 0), length(h.request_raw), COALESCE(length(h.response_raw), 0),
		            h.error, h.source, h.flagged
		     FROM history h
		     WHERE ` + strings.Join(where, " AND ") + `
		     ORDER BY h.id DESC LIMIT ?`
	} else {
		// history_fts must appear unaliased for MATCH/bm25 to resolve
		// against it - this SQLite build doesn't support querying FTS5
		// through a table alias (confirmed directly against the sqlite3
		// CLI: aliasing it raises "no such column").
		where = append([]string{"history_fts MATCH ?"}, where...)
		args = append([]any{prepareFTSQuery(remaining)}, args...)
		q = `SELECT h.id, h.started_at, h.duration_ms, h.method, h.scheme, h.host, h.path,
		            COALESCE(h.status_code, 0), length(h.request_raw), COALESCE(length(h.response_raw), 0),
		            h.error, h.source, h.flagged
		     FROM history_fts
		     JOIN history h ON h.id = history_fts.rowid
		     WHERE ` + strings.Join(where, " AND ") + `
		     ORDER BY bm25(history_fts), h.id DESC LIMIT ?`
	}
	args = append(args, limit)

	rows, err := s.db.Query(q, args...)
	if err != nil {
		return nil, fmt.Errorf("search history: %w", err)
	}
	defer rows.Close()

	var out []Summary
	for rows.Next() {
		var sum Summary
		var startedAt, durationMs int64
		var flagged int
		if err := rows.Scan(&sum.ID, &startedAt, &durationMs, &sum.Method, &sum.Scheme, &sum.Host, &sum.Path,
			&sum.StatusCode, &sum.ReqSize, &sum.RespSize, &sum.Error, &sum.Source, &flagged); err != nil {
			return nil, fmt.Errorf("scan search row: %w", err)
		}
		sum.Flagged = flagged != 0
		sum.StartedAt = time.UnixMilli(startedAt)
		sum.Duration = time.Duration(durationMs) * time.Millisecond
		out = append(out, sum)
	}
	return out, rows.Err()
}

// structuredPredicate holds filters extracted from a search query that
// get applied as real SQL predicates instead of going through FTS5.
type structuredPredicate struct {
	statusSQL  string
	statusArgs []any
	source     string
	flagged    *bool
}

var (
	statusExactRe = regexp.MustCompile(`^(>=|<=|!=|>|<|=)?(\d{3})$`)
	statusRangeRe = regexp.MustCompile(`^([2-5])xx$`)
)

// extractStructured pulls status:/source: tokens out of query, returning
// what's left (for the FTS5 layer, if anything) and the predicates found.
func extractStructured(query string) (remaining string, pred structuredPredicate) {
	fields := strings.Fields(query)
	kept := fields[:0:0]
	for _, f := range fields {
		lower := strings.ToLower(f)
		switch {
		case strings.HasPrefix(lower, "status:"):
			val := strings.ToLower(strings.TrimPrefix(f, "status:"))
			if m := statusRangeRe.FindStringSubmatch(val); m != nil {
				lo, _ := strconv.Atoi(m[1] + "00")
				pred.statusSQL = "status_code >= ? AND status_code < ?"
				pred.statusArgs = []any{lo, lo + 100}
				continue
			}
			if m := statusExactRe.FindStringSubmatch(val); m != nil {
				op := m[1]
				if op == "" {
					op = "="
				}
				n, _ := strconv.Atoi(m[2])
				pred.statusSQL = "status_code " + op + " ?"
				pred.statusArgs = []any{n}
				continue
			}
		case strings.HasPrefix(lower, "source:"):
			pred.source = strings.ToLower(strings.TrimPrefix(f, "source:"))
			continue
		case strings.HasPrefix(lower, "flagged:"):
			val := strings.ToLower(strings.TrimPrefix(lower, "flagged:"))
			switch val {
			case "true", "1", "yes":
				b := true
				pred.flagged = &b
				continue
			case "false", "0", "no":
				b := false
				pred.flagged = &b
				continue
			}
		}
		kept = append(kept, f)
	}
	return strings.Join(kept, " "), pred
}

// Get returns the full entry (including raw bytes) for id.
func (s *Store) Get(id int64) (*Entry, error) {
	row := s.db.QueryRow(
		`SELECT id, started_at, duration_ms, method, scheme, host, path,
		        COALESCE(status_code, 0), request_raw, response_raw,
		        request_exact, response_exact, error, source, flagged,
		        request_truncated, response_truncated
		 FROM history WHERE id = ?`,
		id,
	)
	var e Entry
	var startedAt, durationMs int64
	var reqExact, respExact, flagged, reqTrunc, respTrunc int
	if err := row.Scan(&e.ID, &startedAt, &durationMs, &e.Method, &e.Scheme, &e.Host, &e.Path,
		&e.StatusCode, &e.RequestRaw, &e.ResponseRaw, &reqExact, &respExact, &e.Error, &e.Source, &flagged,
		&reqTrunc, &respTrunc); err != nil {
		return nil, fmt.Errorf("get history entry %d: %w", id, err)
	}
	e.StartedAt = time.UnixMilli(startedAt)
	e.Duration = time.Duration(durationMs) * time.Millisecond
	e.RequestExact = reqExact != 0
	e.ResponseExact = respExact != 0
	e.Flagged = flagged != 0
	e.RequestTruncated = reqTrunc != 0
	e.ResponseTruncated = respTrunc != 0
	return &e, nil
}

// SetFlagged toggles the flagged marker on a history entry - a simple
// "mark this, revisit later" bit, filterable via flagged:true/false in
// Search. Doesn't touch anything else about the entry.
func (s *Store) SetFlagged(id int64, flagged bool) error {
	if _, err := s.db.Exec(`UPDATE history SET flagged = ? WHERE id = ?`, boolToInt(flagged), id); err != nil {
		return fmt.Errorf("set flagged on entry %d: %w", id, err)
	}
	return nil
}

// prepareFTSQuery makes free text safe to hand to FTS5's MATCH, whose
// query grammar reserves a wide range of punctuation ('.', '-', '/',
// '@', '(', ')', and more - confirmed empirically, not just from docs)
// as syntax. A bareword containing any of it is a parse error, not a
// literal search term, which would make searching for the most common
// things in HTTP traffic (domains, paths, hyphenated headers, IPs) fail
// by default. So: quote every plain token as an FTS5 phrase, which is
// syntactically valid regardless of its contents, while still
// recognizing AND/OR/NOT and column:value filters for anyone using them
// deliberately.
//
// Known limitation: this splits on whitespace, so a query the user
// already wrote as a multi-word "quoted phrase" gets re-split and each
// word individually re-quoted (still a valid query, just no longer an
// exact-adjacency phrase match). Fine for the common case of typing
// plain, unquoted search terms, which is what this exists for.
func prepareFTSQuery(q string) string {
	fields := strings.Fields(q)
	for i, f := range fields {
		switch strings.ToUpper(f) {
		case "AND", "OR", "NOT":
			continue
		}
		if col, val, ok := strings.Cut(f, ":"); ok && col != "" && val != "" {
			fields[i] = col + `:"` + strings.ReplaceAll(val, `"`, `""`) + `"`
			continue
		}
		fields[i] = `"` + strings.ReplaceAll(f, `"`, `""`) + `"`
	}
	return strings.Join(fields, " ")
}

// ListRules returns every match-and-replace rule, ordered for application.
func (s *Store) ListRules() ([]rules.Rule, error) {
	rows, err := s.db.Query(
		`SELECT id, enabled, name, scope, part, match, replace, is_regex, position
		 FROM rules ORDER BY position, id`,
	)
	if err != nil {
		return nil, fmt.Errorf("list rules: %w", err)
	}
	defer rows.Close()

	var out []rules.Rule
	for rows.Next() {
		var r rules.Rule
		var enabled, isRegex int
		if err := rows.Scan(&r.ID, &enabled, &r.Name, &r.Scope, &r.Part, &r.Match, &r.Replace, &isRegex, &r.Position); err != nil {
			return nil, fmt.Errorf("scan rule row: %w", err)
		}
		r.Enabled = enabled != 0
		r.IsRegex = isRegex != 0
		out = append(out, r)
	}
	return out, rows.Err()
}

// EnabledRules returns enabled rules for scope ("request" or
// "response"), ordered for application.
func (s *Store) EnabledRules(scope string) ([]rules.Rule, error) {
	rows, err := s.db.Query(
		`SELECT id, enabled, name, scope, part, match, replace, is_regex, position
		 FROM rules WHERE enabled = 1 AND scope = ? ORDER BY position, id`,
		scope,
	)
	if err != nil {
		return nil, fmt.Errorf("enabled rules: %w", err)
	}
	defer rows.Close()

	var out []rules.Rule
	for rows.Next() {
		var r rules.Rule
		var enabled, isRegex int
		if err := rows.Scan(&r.ID, &enabled, &r.Name, &r.Scope, &r.Part, &r.Match, &r.Replace, &isRegex, &r.Position); err != nil {
			return nil, fmt.Errorf("scan rule row: %w", err)
		}
		r.Enabled = enabled != 0
		r.IsRegex = isRegex != 0
		out = append(out, r)
	}
	return out, rows.Err()
}

// AddRule stores r and returns its assigned ID.
func (s *Store) AddRule(r rules.Rule) (int64, error) {
	res, err := s.db.Exec(
		`INSERT INTO rules (enabled, name, scope, part, match, replace, is_regex, position)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
		boolToInt(r.Enabled), r.Name, r.Scope, r.Part, r.Match, r.Replace, boolToInt(r.IsRegex), r.Position,
	)
	if err != nil {
		return 0, fmt.Errorf("add rule: %w", err)
	}
	return res.LastInsertId()
}

// UpdateRule replaces the stored rule with the same ID as r.
func (s *Store) UpdateRule(r rules.Rule) error {
	_, err := s.db.Exec(
		`UPDATE rules SET enabled = ?, name = ?, scope = ?, part = ?, match = ?, replace = ?, is_regex = ?, position = ?
		 WHERE id = ?`,
		boolToInt(r.Enabled), r.Name, r.Scope, r.Part, r.Match, r.Replace, boolToInt(r.IsRegex), r.Position, r.ID,
	)
	if err != nil {
		return fmt.Errorf("update rule %d: %w", r.ID, err)
	}
	return nil
}

// SetRuleEnabled toggles a rule without touching its other fields.
func (s *Store) SetRuleEnabled(id int64, enabled bool) error {
	_, err := s.db.Exec(`UPDATE rules SET enabled = ? WHERE id = ?`, boolToInt(enabled), id)
	if err != nil {
		return fmt.Errorf("set rule %d enabled: %w", id, err)
	}
	return nil
}

// DeleteRule removes a rule.
func (s *Store) DeleteRule(id int64) error {
	if _, err := s.db.Exec(`DELETE FROM rules WHERE id = ?`, id); err != nil {
		return fmt.Errorf("delete rule %d: %w", id, err)
	}
	return nil
}

func boolToInt(b bool) int {
	if b {
		return 1
	}
	return 0
}