QOL & Security improvement

Signed-off-by: Hadi <112569860+anotherhadi@users.noreply.github.com>

internal/config/config.go

@@ -24,6 +24,7 @@ type Config struct {
 		ProjectDir    string `mapstructure:"project_dir"`
 		PluginsDir    string `mapstructure:"plugins_dir"`
 		UpstreamProxy string `mapstructure:"upstream_proxy"`
+		MaxBodySizeMB int    `mapstructure:"max_body_size_mb"`
 	} `mapstructure:"app"`
 
 	TUI struct {

internal/config/default_config.yaml

@@ -5,6 +5,7 @@ app:
   project_dir: ~/.local/share/spilltea
   plugins_dir: ~/.config/spilltea/plugins
   upstream_proxy: "" # e.g. http://corporate-proxy:8888 or http://user:pass@host:8888
+  max_body_size_mb: 50 # max response body size read into memory for large streamed responses (MB)
 
 intercept:
   default_intercept_enabled: true

internal/db/db.go

@@ -2,12 +2,14 @@ package db
 
 import (
 	"database/sql"
+	"sync"
 
 	_ "modernc.org/sqlite"
 )
 
 type DB struct {
-	conn *sql.DB
+	conn    *sql.DB
+	dedupMu sync.Mutex
 }
 
 func Open(path string) (*DB, error) {
@@ -33,7 +35,8 @@ func (d *DB) migrate() error {
 			path         TEXT NOT NULL,
 			status_code  INTEGER NOT NULL,
 			request_raw  TEXT NOT NULL,
-			response_raw TEXT NOT NULL
+			response_raw TEXT NOT NULL,
+			body_hash    TEXT NOT NULL DEFAULT ''
 		);
 CREATE TABLE IF NOT EXISTS replay_entries (
 			id           INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -65,6 +68,10 @@ CREATE TABLE IF NOT EXISTS replay_entries (
 			UNIQUE(plugin_name, dedup_key)
 		);
 	`)
+	if err != nil {
+		return err
+	}
+	_, err = d.conn.Exec(`CREATE INDEX IF NOT EXISTS idx_entries_dedup ON entries(method, host, path, body_hash)`)
 	return err
 }
 

internal/db/entries.go

@@ -1,6 +1,7 @@
 package db
 
 import (
+	"crypto/sha256"
 	"database/sql"
 	"fmt"
 	"strings"
@@ -18,40 +19,45 @@ type Entry struct {
 	ResponseRaw string
 }
 
-// HasDuplicate returns true if an entry with the same method, host, path and
-// request body already exists. Used to implement skip_duplicates filtering.
-func (d *DB) HasDuplicate(method, host, path, body string) (bool, error) {
-	rows, err := d.conn.Query(
-		`SELECT request_raw FROM entries WHERE method = ? AND host = ? AND path = ?`,
-		method, host, path,
-	)
-	if err != nil {
-		return false, err
-	}
-	defer rows.Close()
-	for rows.Next() {
-		var raw string
-		if err := rows.Scan(&raw); err != nil {
-			return false, err
-		}
-		parts := strings.SplitN(raw, "\n\n", 2)
-		entryBody := ""
-		if len(parts) == 2 {
-			entryBody = parts[1]
-		}
-		if entryBody == body {
-			return true, nil
-		}
-	}
-	return false, rows.Err()
+func bodyHash(body string) string {
+	sum := sha256.Sum256([]byte(body))
+	return fmt.Sprintf("%x", sum)
 }
 
-func (d *DB) InsertEntry(e Entry) (Entry, error) {
+// HasDuplicate returns true if an entry with the same method, host, path and
+// request body hash already exists.
+func (d *DB) HasDuplicate(method, host, path, body string) (bool, error) {
+	hash := bodyHash(body)
+	var exists int
+	err := d.conn.QueryRow(
+		`SELECT 1 FROM entries WHERE method = ? AND host = ? AND path = ? AND body_hash = ? LIMIT 1`,
+		method, host, path, hash,
+	).Scan(&exists)
+	if err == sql.ErrNoRows {
+		return false, nil
+	}
+	return err == nil, err
+}
+
+// InsertIfNotDuplicate atomically checks for a duplicate and inserts if none
+// exists. Returns (entry, isDuplicate, error).
+func (d *DB) InsertIfNotDuplicate(e Entry, body string) (Entry, bool, error) {
+	d.dedupMu.Lock()
+	defer d.dedupMu.Unlock()
+	dup, err := d.HasDuplicate(e.Method, e.Host, e.Path, body)
+	if err != nil || dup {
+		return e, dup, err
+	}
+	e, err = d.InsertEntry(e, body)
+	return e, false, err
+}
+
+func (d *DB) InsertEntry(e Entry, body string) (Entry, error) {
 	res, err := d.conn.Exec(
-		`INSERT INTO entries (timestamp, method, host, path, status_code, request_raw, response_raw)
-		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
+		`INSERT INTO entries (timestamp, method, host, path, status_code, request_raw, response_raw, body_hash)
+		 VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
 		e.Timestamp.UTC().Format(time.RFC3339),
-		e.Method, e.Host, e.Path, e.StatusCode, e.RequestRaw, e.ResponseRaw,
+		e.Method, e.Host, e.Path, e.StatusCode, e.RequestRaw, e.ResponseRaw, bodyHash(body),
 	)
 	if err != nil {
 		return e, err
@@ -102,16 +108,10 @@ func (d *DB) SearchEntries(term string) ([]Entry, error) {
 	return scanEntries(rows)
 }
 
-// QueryEntries executes a user-supplied query against the entries table.
-// If the query does not start with SELECT, it is treated as a WHERE expression
-// and wrapped automatically (e.g. "status_code = 404" becomes a full SELECT).
-func (d *DB) QueryEntries(rawSQL string) ([]Entry, error) {
-	q := strings.TrimSpace(rawSQL)
-	if !strings.HasPrefix(strings.ToUpper(q), "SELECT") {
-		q = "SELECT id, timestamp, method, host, path, status_code, request_raw, response_raw FROM entries WHERE " + q
-	} else if strings.ContainsAny(strings.ToUpper(q), "INSERTDELETEUPDATEDROP") {
-		return nil, fmt.Errorf("only SELECT queries are allowed")
-	}
+// QueryEntries runs a WHERE expression supplied by the user against the entries
+// table (e.g. "status_code = 404" or "host LIKE '%example.com%'").
+func (d *DB) QueryEntries(where string) ([]Entry, error) {
+	q := "SELECT id, timestamp, method, host, path, status_code, request_raw, response_raw FROM entries WHERE " + strings.TrimSpace(where)
 	rows, err := d.conn.Query(q)
 	if err != nil {
 		return nil, err

internal/intercept/broker.go

@@ -1,6 +1,7 @@
 package intercept
 
 import (
+	"log"
 	"regexp"
 	"sync"
 	"sync/atomic"
@@ -75,9 +76,12 @@ func (b *Broker) SetCaptureResponse(v bool) {
 func (b *Broker) SetAutoForwardRegex(patterns []string) {
 	compiled := make([]*regexp.Regexp, 0, len(patterns))
 	for _, p := range patterns {
-		if r, err := regexp.Compile(p); err == nil {
-			compiled = append(compiled, r)
+		r, err := regexp.Compile(p)
+		if err != nil {
+			log.Printf("intercept: invalid auto_forward_regex %q: %v", p, err)
+			continue
 		}
+		compiled = append(compiled, r)
 	}
 	b.autoFwdMu.Lock()
 	b.autoFwdRegexes = compiled
@@ -164,19 +168,14 @@ func (b *Broker) SaveEntry(f *proxy.Flow) {
 	if path == "" {
 		path = "/"
 	}
-	if config.Global.History.SkipDuplicates {
-		body := string(r.Body)
-		if dup, _ := d.HasDuplicate(r.Method, r.URL.Host, path, body); dup {
-			return
-		}
-	}
+	body := string(r.Body)
 	pending := db.Entry{
-		Timestamp:   time.Now(),
-		Method:      r.Method,
-		Host:        r.URL.Host,
-		Path:        path,
-		StatusCode:  status,
-		RequestRaw:  FormatRawRequest(f),
+		Timestamp:  time.Now(),
+		Method:     r.Method,
+		Host:       r.URL.Host,
+		Path:       path,
+		StatusCode: status,
+		RequestRaw: FormatRawRequest(f),
 		ResponseRaw: func() string {
 			if config.Global.History.KeepResponses {
 				return FormatRawResponse(f)
@@ -189,7 +188,19 @@ func (b *Broker) SaveEntry(f *proxy.Flow) {
 			return
 		}
 	}
-	entry, err := d.InsertEntry(pending)
+	var (
+		entry db.Entry
+		err   error
+	)
+	if config.Global.History.SkipDuplicates {
+		var dup bool
+		entry, dup, err = d.InsertIfNotDuplicate(pending, body)
+		if dup || err != nil {
+			return
+		}
+	} else {
+		entry, err = d.InsertEntry(pending, body)
+	}
 	if err == nil {
 		if cb := b.onNewEntry; cb != nil {
 			go cb(entry)

internal/plugins/lua.go

@@ -11,7 +11,22 @@ import (
 )
 
 func newLuaState(mgr *Manager, p *Plugin) *lua.LState {
-	L := lua.NewState()
+	L := lua.NewState(lua.Options{SkipOpenLibs: true})
+	for _, lib := range []struct {
+		name string
+		fn   lua.LGFunction
+	}{
+		{lua.LoadLibName, lua.OpenPackage},
+		{lua.BaseLibName, lua.OpenBase},
+		{lua.TabLibName, lua.OpenTable},
+		{lua.StringLibName, lua.OpenString},
+		{lua.MathLibName, lua.OpenMath},
+		{lua.CoroutineLibName, lua.OpenCoroutine},
+	} {
+		L.Push(L.NewFunction(lib.fn))
+		L.Push(lua.LString(lib.name))
+		L.Call(1, 0)
+	}
 	registerUtilities(L, mgr, p)
 	return L
 }

internal/plugins/manager.go

@@ -270,20 +270,22 @@ func (m *Manager) RunOnQuit() {
 	}
 }
 
-func (m *Manager) RunSyncOnRequest(f *goproxy.Flow) intercept.Decision {
+// runSyncDecisionForPlugins runs hookName synchronously for all enabled plugins
+// that registered it as sync, and returns the first non-Intercept decision.
+func (m *Manager) runSyncDecisionForPlugins(hookName string, argsFor func(*Plugin) []lua.LValue) intercept.Decision {
 	for _, p := range m.GetPlugins() {
 		if !p.Enabled {
 			continue
 		}
-		hc, ok := p.hooks["on_request"]
+		hc, ok := p.hooks[hookName]
 		if !ok || !hc.Sync {
 			continue
 		}
 		p.mu.Lock()
-		result, err := callHook(p, "on_request", pushRequest(p.L, f))
+		result, err := callHook(p, hookName, argsFor(p)...)
 		p.mu.Unlock()
 		if err != nil {
-			log.Printf("plugin %s on_request: %v", p.Name, err)
+			log.Printf("plugin %s %s: %v", p.Name, hookName, err)
 			continue
 		}
 		switch result {
@@ -296,68 +298,49 @@ func (m *Manager) RunSyncOnRequest(f *goproxy.Flow) intercept.Decision {
 	return intercept.Intercept
 }
 
+// runAsyncForPlugins fires hookName asynchronously for all enabled plugins
+// that registered it as async.
+func (m *Manager) runAsyncForPlugins(hookName string, argsFor func(*Plugin) []lua.LValue) {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		hc, ok := p.hooks[hookName]
+		if !ok || hc.Sync {
+			continue
+		}
+		go func(p *Plugin) {
+			p.mu.Lock()
+			if _, err := callHook(p, hookName, argsFor(p)...); err != nil {
+				log.Printf("plugin %s %s: %v", p.Name, hookName, err)
+			}
+			p.mu.Unlock()
+		}(p)
+	}
+}
+
+func (m *Manager) RunSyncOnRequest(f *goproxy.Flow) intercept.Decision {
+	return m.runSyncDecisionForPlugins("on_request", func(p *Plugin) []lua.LValue {
+		return []lua.LValue{pushRequest(p.L, f)}
+	})
+}
+
 func (m *Manager) RunAsyncOnRequest(f *goproxy.Flow) {
-	for _, p := range m.GetPlugins() {
-		if !p.Enabled {
-			continue
-		}
-		hc, ok := p.hooks["on_request"]
-		if !ok || hc.Sync {
-			continue
-		}
-		go func(p *Plugin) {
-			p.mu.Lock()
-			if _, err := callHook(p, "on_request", pushRequest(p.L, f)); err != nil {
-				log.Printf("plugin %s on_request: %v", p.Name, err)
-			}
-			p.mu.Unlock()
-		}(p)
-	}
+	m.runAsyncForPlugins("on_request", func(p *Plugin) []lua.LValue {
+		return []lua.LValue{pushRequest(p.L, f)}
+	})
 }
 
 func (m *Manager) RunSyncOnResponse(f *goproxy.Flow) intercept.Decision {
-	for _, p := range m.GetPlugins() {
-		if !p.Enabled {
-			continue
-		}
-		hc, ok := p.hooks["on_response"]
-		if !ok || !hc.Sync {
-			continue
-		}
-		p.mu.Lock()
-		result, err := callHook(p, "on_response", pushRequest(p.L, f), pushResponse(p.L, f))
-		p.mu.Unlock()
-		if err != nil {
-			log.Printf("plugin %s on_response: %v", p.Name, err)
-			continue
-		}
-		switch result {
-		case "drop":
-			return intercept.Drop
-		case "forward":
-			return intercept.Forward
-		}
-	}
-	return intercept.Intercept
+	return m.runSyncDecisionForPlugins("on_response", func(p *Plugin) []lua.LValue {
+		return []lua.LValue{pushRequest(p.L, f), pushResponse(p.L, f)}
+	})
 }
 
 func (m *Manager) RunAsyncOnResponse(f *goproxy.Flow) {
-	for _, p := range m.GetPlugins() {
-		if !p.Enabled {
-			continue
-		}
-		hc, ok := p.hooks["on_response"]
-		if !ok || hc.Sync {
-			continue
-		}
-		go func(p *Plugin) {
-			p.mu.Lock()
-			if _, err := callHook(p, "on_response", pushRequest(p.L, f), pushResponse(p.L, f)); err != nil {
-				log.Printf("plugin %s on_response: %v", p.Name, err)
-			}
-			p.mu.Unlock()
-		}(p)
-	}
+	m.runAsyncForPlugins("on_response", func(p *Plugin) []lua.LValue {
+		return []lua.LValue{pushRequest(p.L, f), pushResponse(p.L, f)}
+	})
 }
 
 // RunSyncOnHistoryEntry is called before DB insert; returns false to skip saving.
@@ -385,20 +368,7 @@ func (m *Manager) RunSyncOnHistoryEntry(e db.Entry) bool {
 }
 
 func (m *Manager) RunAsyncOnHistoryEntry(e db.Entry) {
-	for _, p := range m.GetPlugins() {
-		if !p.Enabled {
-			continue
-		}
-		hc, ok := p.hooks["on_history_entry"]
-		if !ok || hc.Sync {
-			continue
-		}
-		go func(p *Plugin) {
-			p.mu.Lock()
-			if _, err := callHook(p, "on_history_entry", pushEntry(p.L, e)); err != nil {
-				log.Printf("plugin %s on_history_entry: %v", p.Name, err)
-			}
-			p.mu.Unlock()
-		}(p)
-	}
+	m.runAsyncForPlugins("on_history_entry", func(p *Plugin) []lua.LValue {
+		return []lua.LValue{pushEntry(p.L, e)}
+	})
 }

internal/proxy/proxy.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"os"
 
@@ -63,7 +64,11 @@ func (a *interceptAddon) Request(f *goproxy.Flow) {
 func (a *interceptAddon) Response(f *goproxy.Flow) {
 	if f.Response != nil {
 		if len(f.Response.Body) == 0 && f.Response.BodyReader != nil {
-			body, _ := io.ReadAll(f.Response.BodyReader)
+			limit := int64(config.Global.App.MaxBodySizeMB) * 1024 * 1024
+			body, err := io.ReadAll(io.LimitReader(f.Response.BodyReader, limit))
+			if err != nil {
+				log.Printf("proxy: reading response body: %v", err)
+			}
 			f.Response.Body = body
 			f.Response.BodyReader = nil
 		}

internal/ui/components/copyas/formats.go

@@ -237,7 +237,7 @@ func toHAR(pr parsedRequest) string {
 		} `json:"timings"`
 	}
 	type harLog struct {
-		Version string     `json:"version"`
+		Version string `json:"version"`
 		Creator struct {
 			Name    string `json:"name"`
 			Version string `json:"version"`