store.go

package main

import (
	"context"
	"database/sql"
	"fmt"
	"net/url"
	"os"
	"time"

	_ "modernc.org/sqlite"
)

type Session struct {
	ID           int64
	TS           time.Time
	Scheme       string // "HTTP" or "HTTPS"
	Method       string
	URL          string
	Host         string
	Status       int
	DurationMS   int64
	ReqHeaders   string
	ReqBody      []byte
	ReqBodySize  int64
	RespHeaders  string
	RespBody     []byte
	RespBodySize int64
	Err          string
}

type SessionRow struct {
	ID         int64
	TS         time.Time
	Scheme     string
	Method     string
	Host       string
	Path       string
	Status     int
	DurationMS int64
	Err        string
}

type Store struct {
	db   *sql.DB
	path string
}

const schema = `
CREATE TABLE IF NOT EXISTS sessions (
  id              INTEGER PRIMARY KEY AUTOINCREMENT,
  ts              INTEGER NOT NULL,
  scheme          TEXT NOT NULL,
  method          TEXT NOT NULL,
  url             TEXT NOT NULL,
  host            TEXT NOT NULL,
  status          INTEGER,
  duration_ms     INTEGER,
  req_headers     TEXT NOT NULL,
  req_body        BLOB,
  req_body_size   INTEGER NOT NULL DEFAULT 0,
  resp_headers    TEXT,
  resp_body       BLOB,
  resp_body_size  INTEGER NOT NULL DEFAULT 0,
  err             TEXT
);
CREATE INDEX IF NOT EXISTS idx_sessions_ts ON sessions(ts);

CREATE VIRTUAL TABLE IF NOT EXISTS sessions_fts USING fts5(
  url, host, req_headers, req_body, resp_headers, resp_body,
  content='sessions', content_rowid='id', tokenize='unicode61'
);

CREATE TRIGGER IF NOT EXISTS sessions_ai AFTER INSERT ON sessions BEGIN
  INSERT INTO sessions_fts(rowid, url, host, req_headers, req_body, resp_headers, resp_body)
  VALUES (new.id, new.url, new.host, new.req_headers, CAST(new.req_body AS TEXT), new.resp_headers, CAST(new.resp_body AS TEXT));
END;

CREATE TRIGGER IF NOT EXISTS sessions_ad AFTER DELETE ON sessions BEGIN
  INSERT INTO sessions_fts(sessions_fts, rowid, url, host, req_headers, req_body, resp_headers, resp_body)
  VALUES('delete', old.id, old.url, old.host, old.req_headers, CAST(old.req_body AS TEXT), old.resp_headers, CAST(old.resp_body AS TEXT));
END;
`

func OpenStore() (*Store, error) {
	dir, err := appDir()
	if err != nil {
		return nil, fmt.Errorf("locate app dir: %w", err)
	}
	// CreateTemp gives us a unique, race-free file in our own 0700 directory,
	// avoiding the predictable /tmp/netmiddler-<pid>.db pattern that allowed
	// symlink / hardlink clobbering.
	f, err := os.CreateTemp(dir, "session-*.db")
	if err != nil {
		return nil, fmt.Errorf("create temp db: %w", err)
	}
	path := f.Name()
	f.Close()
	// Tighten perms; CreateTemp gives 0600 on Unix but be explicit.
	if err := os.Chmod(path, 0o600); err != nil {
		os.Remove(path)
		return nil, fmt.Errorf("chmod db: %w", err)
	}

	dsn := fmt.Sprintf("file:%s?_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=busy_timeout(5000)", path)
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		os.Remove(path)
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	if _, err := db.ExecContext(context.Background(), schema); err != nil {
		db.Close()
		os.Remove(path)
		return nil, fmt.Errorf("apply schema: %w", err)
	}
	return &Store{db: db, path: path}, nil
}

func (s *Store) Close() error {
	if s == nil {
		return nil
	}
	err := s.db.Close()
	os.Remove(s.path)
	// WAL companions
	os.Remove(s.path + "-wal")
	os.Remove(s.path + "-shm")
	return err
}

func (s *Store) Save(ses *Session) (int64, error) {
	var status sql.NullInt64
	if ses.Status != 0 {
		status = sql.NullInt64{Int64: int64(ses.Status), Valid: true}
	}
	res, err := s.db.ExecContext(context.Background(), `
        INSERT INTO sessions
            (ts, scheme, method, url, host, status, duration_ms,
             req_headers, req_body, req_body_size, resp_headers, resp_body, resp_body_size, err)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		ses.TS.UnixNano(),
		ses.Scheme,
		ses.Method,
		ses.URL,
		ses.Host,
		status,
		ses.DurationMS,
		ses.ReqHeaders,
		ses.ReqBody,
		ses.ReqBodySize,
		ses.RespHeaders,
		ses.RespBody,
		ses.RespBodySize,
		nullableString(ses.Err),
	)
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

func nullableString(s string) any {
	if s == "" {
		return nil
	}
	return s
}

func (s *Store) List(limit int) ([]SessionRow, error) {
	rows, err := s.db.QueryContext(context.Background(), `
        SELECT id, ts, scheme, method, url, host, COALESCE(status, 0), COALESCE(duration_ms, 0), COALESCE(err, '')
        FROM sessions
        ORDER BY id DESC
        LIMIT ?`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanRows(rows)
}

func (s *Store) Search(query string, limit int) ([]SessionRow, error) {
	if query == "" {
		return s.List(limit)
	}
	rows, err := s.db.QueryContext(context.Background(), `
        SELECT s.id, s.ts, s.scheme, s.method, s.url, s.host,
               COALESCE(s.status, 0), COALESCE(s.duration_ms, 0), COALESCE(s.err, '')
        FROM sessions s
        JOIN sessions_fts f ON f.rowid = s.id
        WHERE sessions_fts MATCH ?
        ORDER BY s.id DESC
        LIMIT ?`, ftsEscape(query), limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanRows(rows)
}

func scanRows(rows *sql.Rows) ([]SessionRow, error) {
	var out []SessionRow
	for rows.Next() {
		var r SessionRow
		var tsNanos int64
		var rawURL string
		if err := rows.Scan(&r.ID, &tsNanos, &r.Scheme, &r.Method, &rawURL, &r.Host, &r.Status, &r.DurationMS, &r.Err); err != nil {
			return nil, err
		}
		r.TS = time.Unix(0, tsNanos)
		r.Path = pathOf(rawURL)
		out = append(out, r)
	}
	return out, rows.Err()
}

func pathOf(rawURL string) string {
	u, err := url.Parse(rawURL)
	if err != nil || u.Path == "" {
		return rawURL
	}
	if u.RawQuery != "" {
		return u.Path + "?" + u.RawQuery
	}
	return u.Path
}

// ftsEscape wraps the user query in double-quotes and escapes embedded quotes,
// so that arbitrary punctuation (slashes, dots, parens) is treated as literal
// text rather than FTS5 operators.
func ftsEscape(q string) string {
	out := make([]byte, 0, len(q)+2)
	out = append(out, '"')
	for i := 0; i < len(q); i++ {
		if q[i] == '"' {
			out = append(out, '"', '"')
		} else {
			out = append(out, q[i])
		}
	}
	out = append(out, '"')
	return string(out)
}

func (s *Store) Get(id int64) (*Session, error) {
	var (
		ses     Session
		tsNanos int64
		status  sql.NullInt64
		errStr  sql.NullString
	)
	err := s.db.QueryRowContext(context.Background(), `
        SELECT id, ts, scheme, method, url, host, status, COALESCE(duration_ms, 0),
               req_headers, req_body, req_body_size,
               COALESCE(resp_headers, ''), resp_body, resp_body_size, err
        FROM sessions WHERE id = ?`, id).Scan(
		&ses.ID, &tsNanos, &ses.Scheme, &ses.Method, &ses.URL, &ses.Host,
		&status, &ses.DurationMS,
		&ses.ReqHeaders, &ses.ReqBody, &ses.ReqBodySize,
		&ses.RespHeaders, &ses.RespBody, &ses.RespBodySize, &errStr,
	)
	if err != nil {
		return nil, err
	}
	ses.TS = time.Unix(0, tsNanos)
	if status.Valid {
		ses.Status = int(status.Int64)
	}
	if errStr.Valid {
		ses.Err = errStr.String
	}
	return &ses, nil
}

func (s *Store) Clear() error {
	ctx := context.Background()
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM sessions`); err != nil {
		tx.Rollback()
		return err
	}
	if _, err := tx.ExecContext(ctx, `INSERT INTO sessions_fts(sessions_fts) VALUES('rebuild')`); err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

func (s *Store) Count() (int, error) {
	var n int
	err := s.db.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM sessions`).Scan(&n)
	return n, err
}