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"
	"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'
);

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'")
	// 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
	Error         string // network/transport error, if the request never got a response
	Source        string // "proxy" or "repeater"
}

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

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

// 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
		 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
		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); err != nil {
			return nil, fmt.Errorf("scan history row: %w", err)
		}
		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.
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
	}
	query = prepareFTSQuery(query)
	// 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").
	rows, err := s.db.Query(
		`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
		 FROM history_fts
		 JOIN history h ON h.id = history_fts.rowid
		 WHERE history_fts MATCH ? AND h.id < ?
		 ORDER BY bm25(history_fts), h.id DESC LIMIT ?`,
		query, beforeID, limit,
	)
	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
		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); err != nil {
			return nil, fmt.Errorf("scan search row: %w", err)
		}
		sum.StartedAt = time.UnixMilli(startedAt)
		sum.Duration = time.Duration(durationMs) * time.Millisecond
		out = append(out, sum)
	}
	return out, rows.Err()
}

// 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
		 FROM history WHERE id = ?`,
		id,
	)
	var e Entry
	var startedAt, durationMs int64
	var reqExact, respExact 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); 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
	return &e, 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
}