From 0a477fd288d064d265395c88dbd0f261b723804e Mon Sep 17 00:00:00 2001 From: spiperac Date: Mon, 15 Jun 2026 06:42:13 +0200 Subject: Initial release --- internal/cert/ca.go | 298 ++++++ internal/cert/ca_test.go | 345 +++++++ internal/cert/install.go | 149 +++ internal/events/bus.go | 305 ++++++ internal/events/bus_test.go | 746 ++++++++++++++ internal/events/events.go | 122 +++ internal/http/compression.go | 142 +++ internal/http/compression_test.go | 152 +++ internal/http/headers.go | 54 + internal/http/headers_test.go | 117 +++ internal/http/rawrequest.go | 327 ++++++ internal/http/rawrequest_test.go | 254 +++++ internal/logger/logger.go | 49 + internal/plugin/api.go | 313 ++++++ internal/plugin/engine.go | 347 +++++++ internal/plugin/engine_test.go | 1003 ++++++++++++++++++ internal/plugin/plugin.go | 197 ++++ internal/plugin/plugin_test.go | 391 +++++++ internal/proxy/forward.go | 29 + internal/proxy/forward_test.go | 197 ++++ internal/proxy/mitm.go | 214 ++++ internal/proxy/mitm2.go | 211 ++++ internal/proxy/mitm2_test.go | 280 +++++ internal/proxy/mitm_test.go | 281 +++++ internal/proxy/proxy.go | 200 ++++ internal/proxy/proxy_test.go | 218 ++++ internal/proxy/recover.go | 13 + internal/proxy/websocket.go | 225 ++++ internal/proxy/websocket_test.go | 764 ++++++++++++++ internal/store/annotation.go | 58 ++ internal/store/annotation_test.go | 92 ++ internal/store/har.go | 239 +++++ internal/store/history.go | 401 ++++++++ internal/store/history_test.go | 190 ++++ internal/store/intercept.go | 233 +++++ internal/store/intruder.go | 72 ++ internal/store/loot.go | 103 ++ internal/store/loot_test.go | 155 +++ internal/store/plugin.go | 156 +++ internal/store/plugin_test.go | 368 +++++++ internal/store/project.go | 245 +++++ internal/store/repeater.go | 103 ++ internal/store/repeater_test.go | 265 +++++ internal/store/scope.go | 116 +++ internal/store/scope_test.go | 135 +++ internal/store/search_test.go | 105 ++ internal/store/store_test.go | 22 + internal/store/websocket.go | 198 ++++ internal/store/websocket_test.go | 200 ++++ internal/ui/app.go | 239 +++++ internal/ui/browser.go | 283 +++++ internal/ui/decoder.go | 254 +++++ internal/ui/format.go | 120 +++ internal/ui/har.go | 88 ++ internal/ui/helpers.go | 91 ++ internal/ui/history.go | 1303 +++++++++++++++++++++++ internal/ui/history_test.go | 124 +++ internal/ui/inspector.go | 125 +++ internal/ui/intercept.go | 216 ++++ internal/ui/intruder.go | 742 ++++++++++++++ internal/ui/keybinds.go | 92 ++ internal/ui/launch.go | 219 ++++ internal/ui/loot.go | 394 +++++++ internal/ui/plugins.go | 179 ++++ internal/ui/repeater.go | 224 ++++ internal/ui/repeater_http.go | 310 ++++++ internal/ui/repeater_test.go | 378 +++++++ internal/ui/repeater_ws.go | 418 ++++++++ internal/ui/scope.go | 126 +++ internal/ui/settings.go | 401 ++++++++ internal/ui/theme.go | 373 +++++++ internal/ui/websocket.go | 379 +++++++ internal/ui/widgets/datatable.go | 876 ++++++++++++++++ internal/ui/widgets/datatable_test.go | 199 ++++ internal/ui/widgets/textview.go | 836 +++++++++++++++ internal/ui/widgets/textview_test.go | 189 ++++ internal/ui/widgets/textviewentry.go | 1589 +++++++++++++++++++++++++++++ internal/ui/widgets/textviewentry_test.go | 381 +++++++ internal/ui/widgets/tokeniser.go | 362 +++++++ internal/ui/widgets/tokeniser_test.go | 213 ++++ 80 files changed, 23122 insertions(+) create mode 100644 internal/cert/ca.go create mode 100644 internal/cert/ca_test.go create mode 100644 internal/cert/install.go create mode 100644 internal/events/bus.go create mode 100644 internal/events/bus_test.go create mode 100644 internal/events/events.go create mode 100644 internal/http/compression.go create mode 100644 internal/http/compression_test.go create mode 100644 internal/http/headers.go create mode 100644 internal/http/headers_test.go create mode 100644 internal/http/rawrequest.go create mode 100644 internal/http/rawrequest_test.go create mode 100644 internal/logger/logger.go create mode 100644 internal/plugin/api.go create mode 100644 internal/plugin/engine.go create mode 100644 internal/plugin/engine_test.go create mode 100644 internal/plugin/plugin.go create mode 100644 internal/plugin/plugin_test.go create mode 100644 internal/proxy/forward.go create mode 100644 internal/proxy/forward_test.go create mode 100644 internal/proxy/mitm.go create mode 100644 internal/proxy/mitm2.go create mode 100644 internal/proxy/mitm2_test.go create mode 100644 internal/proxy/mitm_test.go create mode 100644 internal/proxy/proxy.go create mode 100644 internal/proxy/proxy_test.go create mode 100644 internal/proxy/recover.go create mode 100644 internal/proxy/websocket.go create mode 100644 internal/proxy/websocket_test.go create mode 100644 internal/store/annotation.go create mode 100644 internal/store/annotation_test.go create mode 100644 internal/store/har.go create mode 100644 internal/store/history.go create mode 100644 internal/store/history_test.go create mode 100644 internal/store/intercept.go create mode 100644 internal/store/intruder.go create mode 100644 internal/store/loot.go create mode 100644 internal/store/loot_test.go create mode 100644 internal/store/plugin.go create mode 100644 internal/store/plugin_test.go create mode 100644 internal/store/project.go create mode 100644 internal/store/repeater.go create mode 100644 internal/store/repeater_test.go create mode 100644 internal/store/scope.go create mode 100644 internal/store/scope_test.go create mode 100644 internal/store/search_test.go create mode 100644 internal/store/store_test.go create mode 100644 internal/store/websocket.go create mode 100644 internal/store/websocket_test.go create mode 100644 internal/ui/app.go create mode 100644 internal/ui/browser.go create mode 100644 internal/ui/decoder.go create mode 100644 internal/ui/format.go create mode 100644 internal/ui/har.go create mode 100644 internal/ui/helpers.go create mode 100644 internal/ui/history.go create mode 100644 internal/ui/history_test.go create mode 100644 internal/ui/inspector.go create mode 100644 internal/ui/intercept.go create mode 100644 internal/ui/intruder.go create mode 100644 internal/ui/keybinds.go create mode 100644 internal/ui/launch.go create mode 100644 internal/ui/loot.go create mode 100644 internal/ui/plugins.go create mode 100644 internal/ui/repeater.go create mode 100644 internal/ui/repeater_http.go create mode 100644 internal/ui/repeater_test.go create mode 100644 internal/ui/repeater_ws.go create mode 100644 internal/ui/scope.go create mode 100644 internal/ui/settings.go create mode 100644 internal/ui/theme.go create mode 100644 internal/ui/websocket.go create mode 100644 internal/ui/widgets/datatable.go create mode 100644 internal/ui/widgets/datatable_test.go create mode 100644 internal/ui/widgets/textview.go create mode 100644 internal/ui/widgets/textview_test.go create mode 100644 internal/ui/widgets/textviewentry.go create mode 100644 internal/ui/widgets/textviewentry_test.go create mode 100644 internal/ui/widgets/tokeniser.go create mode 100644 internal/ui/widgets/tokeniser_test.go (limited to 'internal') diff --git a/internal/cert/ca.go b/internal/cert/ca.go new file mode 100644 index 0000000..76579a4 --- /dev/null +++ b/internal/cert/ca.go @@ -0,0 +1,298 @@ +// Package cert handles all certificate concerns for Shiv: +// - Root CA generation and persistent storage +// - Per-hostname TLS certificate generation and signing +// - In-memory cert cache (sync.Map, rebuilt each launch) +// - OS trust store installation (see install.go) +package cert + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "errors" + "fmt" + "math/big" + "net" + "os" + "path/filepath" + "sync" + "time" +) + +const ( + caKeySize = 4096 + hostKeySize = 2048 + caValidity = 10 * 365 * 24 * time.Hour + hostValidity = 365 * 24 * time.Hour +) + +// CA holds the loaded root certificate authority keypair and an +// in-memory cache of already-generated per-host certificates. +type CA struct { + cert *x509.Certificate + key *rsa.PrivateKey + cache sync.Map // map[string]*tls.Certificate + fresh bool +} + +// Fresh returns true if the CA was generated this run (not loaded from disk). +// Use this to decide whether to attempt OS trust store installation. +func (ca *CA) Fresh() bool { + return ca.fresh +} + +// Dir returns the directory where Shiv stores the CA files. +// Uses os.UserConfigDir()/shiv — created if it does not exist. +func Dir() (string, error) { + configDir, err := os.UserConfigDir() + if err != nil { + return "", fmt.Errorf("cert: cannot find user config dir: %w", err) + } + dir := filepath.Join(configDir, "shiv") + if err := os.MkdirAll(dir, 0700); err != nil { + return "", fmt.Errorf("cert: cannot create config dir %s: %w", dir, err) + } + return dir, nil +} + +// Load loads the CA from disk, generating it if it does not exist. +// The CA keypair is stored in dir/ca.key and dir/ca.crt. +func Load() (*CA, error) { + dir, err := Dir() + if err != nil { + return nil, err + } + keyPath := filepath.Join(dir, "ca.key") + crtPath := filepath.Join(dir, "ca.crt") + + // If both files exist, load them. + if fileExists(keyPath) && fileExists(crtPath) { + return loadFromDisk(keyPath, crtPath) + } + + // Otherwise generate a fresh CA and persist it. + return generate(keyPath, crtPath) +} + +// LoadFromDir is like Load but uses an explicit directory. Useful in tests. +func LoadFromDir(dir string) (*CA, error) { + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, fmt.Errorf("cert: cannot create dir %s: %w", dir, err) + } + keyPath := filepath.Join(dir, "ca.key") + crtPath := filepath.Join(dir, "ca.crt") + + if fileExists(keyPath) && fileExists(crtPath) { + return loadFromDisk(keyPath, crtPath) + } + return generate(keyPath, crtPath) +} + +// CertPEMPath returns the path to the CA certificate PEM file. +func CertPEMPath() (string, error) { + dir, err := Dir() + if err != nil { + return "", err + } + return filepath.Join(dir, "ca.crt"), nil +} + +// CertPEM returns the PEM-encoded CA certificate bytes. +func (ca *CA) CertPEM() ([]byte, error) { + return encodeCertPEM(ca.cert.Raw), nil +} + +// TLSCertForHost returns a *tls.Certificate for the given hostname, +// generating and caching one if it does not exist yet. +func (ca *CA) TLSCertForHost(host string) (*tls.Certificate, error) { + // Strip port if present. + hostname := host + if h, _, err := net.SplitHostPort(host); err == nil { + hostname = h + } + + if cachedCert, ok := ca.cache.Load(hostname); ok { + return cachedCert.(*tls.Certificate), nil + } + + tlsCert, err := ca.generateHostCert(hostname) + if err != nil { + return nil, err + } + + // Store-if-absent to avoid races generating the same cert twice. + actual, _ := ca.cache.LoadOrStore(hostname, tlsCert) + return actual.(*tls.Certificate), nil +} + +// X509Cert returns the underlying *x509.Certificate for the CA. +// Used by tests and trust store installation. +func (ca *CA) X509Cert() *x509.Certificate { + return ca.cert +} + +// --------------------------------------------------------------------- +// internal helpers +// --------------------------------------------------------------------- + +func generate(keyPath, crtPath string) (*CA, error) { + caKey, err := rsa.GenerateKey(rand.Reader, caKeySize) + if err != nil { + return nil, fmt.Errorf("cert: CA key generation failed: %w", err) + } + + serial, err := randomSerial() + if err != nil { + return nil, err + } + + tmpl := &x509.Certificate{ + SerialNumber: serial, + Subject: pkix.Name{ + CommonName: "Shiv Local CA", + Organization: []string{"Shiv"}, + }, + NotBefore: time.Now().Add(-1 * time.Minute), // small back-date for clock skew + NotAfter: time.Now().Add(caValidity), + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign, + BasicConstraintsValid: true, + IsCA: true, + MaxPathLen: 1, + } + + derBytes, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &caKey.PublicKey, caKey) + if err != nil { + return nil, fmt.Errorf("cert: CA cert creation failed: %w", err) + } + + parsedCert, err := x509.ParseCertificate(derBytes) + if err != nil { + return nil, fmt.Errorf("cert: CA cert parse failed: %w", err) + } + + // Persist key. + if err := writePEM(keyPath, "RSA PRIVATE KEY", x509.MarshalPKCS1PrivateKey(caKey), 0600); err != nil { + return nil, err + } + // Persist cert. + if err := writePEM(crtPath, "CERTIFICATE", derBytes, 0644); err != nil { + return nil, err + } + + return &CA{cert: parsedCert, key: caKey, fresh: true}, nil +} + +func loadFromDisk(keyPath, crtPath string) (*CA, error) { + keyPEM, err := os.ReadFile(keyPath) + if err != nil { + return nil, fmt.Errorf("cert: read CA key: %w", err) + } + crtPEM, err := os.ReadFile(crtPath) + if err != nil { + return nil, fmt.Errorf("cert: read CA cert: %w", err) + } + + keyBlock, _ := pem.Decode(keyPEM) + if keyBlock == nil { + return nil, errors.New("cert: CA key PEM decode failed") + } + caKey, err := x509.ParsePKCS1PrivateKey(keyBlock.Bytes) + if err != nil { + return nil, fmt.Errorf("cert: CA key parse: %w", err) + } + + certBlock, _ := pem.Decode(crtPEM) + if certBlock == nil { + return nil, errors.New("cert: CA cert PEM decode failed") + } + parsedCert, err := x509.ParseCertificate(certBlock.Bytes) + if err != nil { + return nil, fmt.Errorf("cert: CA cert parse: %w", err) + } + + if time.Now().After(parsedCert.NotAfter) { + return generate(keyPath, crtPath) + } + + return &CA{cert: parsedCert, key: caKey}, nil +} + +func (ca *CA) generateHostCert(hostname string) (*tls.Certificate, error) { + key, err := rsa.GenerateKey(rand.Reader, hostKeySize) + if err != nil { + return nil, fmt.Errorf("cert: host key generation failed for %s: %w", hostname, err) + } + + serial, err := randomSerial() + if err != nil { + return nil, err + } + + certTemplate := &x509.Certificate{ + SerialNumber: serial, + Subject: pkix.Name{ + CommonName: hostname, + }, + NotBefore: time.Now().Add(-1 * time.Minute), + NotAfter: time.Now().Add(hostValidity), + KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment, + ExtKeyUsage: []x509.ExtKeyUsage{ + x509.ExtKeyUsageServerAuth, + }, + } + + // Support both IP addresses and DNS names. + if ip := net.ParseIP(hostname); ip != nil { + certTemplate.IPAddresses = []net.IP{ip} + } else { + certTemplate.DNSNames = []string{hostname} + } + + derBytes, err := x509.CreateCertificate(rand.Reader, certTemplate, ca.cert, &key.PublicKey, ca.key) + if err != nil { + return nil, fmt.Errorf("cert: host cert creation failed for %s: %w", hostname, err) + } + + tlsCert := &tls.Certificate{ + Certificate: [][]byte{derBytes, ca.cert.Raw}, + PrivateKey: key, + } + + // Parse and attach for field access (e.g. in tests). + leaf, err := x509.ParseCertificate(derBytes) + if err != nil { + return nil, fmt.Errorf("cert: host cert parse failed: %w", err) + } + tlsCert.Leaf = leaf + + return tlsCert, nil +} + +func randomSerial() (*big.Int, error) { + serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128)) + if err != nil { + return nil, fmt.Errorf("cert: serial generation failed: %w", err) + } + return serial, nil +} + +func writePEM(path, pemType string, derBytes []byte, mode os.FileMode) error { + pemFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode) + if err != nil { + return fmt.Errorf("cert: open %s for write: %w", path, err) + } + defer pemFile.Close() + return pem.Encode(pemFile, &pem.Block{Type: pemType, Bytes: derBytes}) +} + +func encodeCertPEM(derBytes []byte) []byte { + return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}) +} + +func fileExists(path string) bool { + _, err := os.Stat(path) + return err == nil +} diff --git a/internal/cert/ca_test.go b/internal/cert/ca_test.go new file mode 100644 index 0000000..8922c32 --- /dev/null +++ b/internal/cert/ca_test.go @@ -0,0 +1,345 @@ +package cert_test + +import ( + "crypto/tls" + "crypto/x509" + "os" + "path/filepath" + "slices" + "sync" + "testing" + "time" + + "github.com/shiv/internal/cert" +) + +// newTestCA creates a fresh CA in a temp dir and returns it along with +// the dir path. The caller owns cleanup via t.Cleanup. +func newTestCA(t *testing.T) (*cert.CA, string) { + t.Helper() + dir := t.TempDir() + ca, err := cert.LoadFromDir(dir) + if err != nil { + t.Fatalf("LoadFromDir: %v", err) + } + return ca, dir +} + +// ----------------------------------------------------------------------- +// CA generation +// ----------------------------------------------------------------------- + +func TestCAGeneration(t *testing.T) { + ca, dir := newTestCA(t) + + // Files must exist on disk. + for _, name := range []string{"ca.key", "ca.crt"} { + path := filepath.Join(dir, name) + if _, err := os.Stat(path); err != nil { + t.Errorf("expected file %s to exist: %v", name, err) + } + } + + x := ca.X509Cert() + if x == nil { + t.Fatal("X509Cert() returned nil") + } + if !x.IsCA { + t.Error("CA cert must have IsCA = true") + } + if x.Subject.CommonName != "Shiv Local CA" { + t.Errorf("unexpected CN: %q", x.Subject.CommonName) + } + if time.Until(x.NotAfter) < 9*365*24*time.Hour { + t.Errorf("CA cert validity too short: expires %v", x.NotAfter) + } +} + +// ----------------------------------------------------------------------- +// CA persistence — load from existing files +// ----------------------------------------------------------------------- + +func TestCAPersistence(t *testing.T) { + _, dir := newTestCA(t) + + // Load again from the same directory — must succeed without regenerating. + ca2, err := cert.LoadFromDir(dir) + if err != nil { + t.Fatalf("second LoadFromDir: %v", err) + } + + // Serial numbers must match — same cert was loaded, not regenerated. + _, dir2 := newTestCA(t) + ca3, err := cert.LoadFromDir(dir2) + if err != nil { + t.Fatalf("third LoadFromDir (fresh dir): %v", err) + } + + x2 := ca2.X509Cert() + x3 := ca3.X509Cert() + + // Same dir → same serial. Different dir → different serial (fresh CA). + // Load ca again from dir for serial comparison. + ca1reload, _ := cert.LoadFromDir(dir) + x1reload := ca1reload.X509Cert() + + if x1reload.SerialNumber.Cmp(x2.SerialNumber) != 0 { + t.Errorf("reloaded CA has different serial: got %v want %v", + x2.SerialNumber, x1reload.SerialNumber) + } + if x2.SerialNumber.Cmp(x3.SerialNumber) == 0 { + t.Error("independent CAs should have different serials") + } +} + +// ----------------------------------------------------------------------- +// Per-host cert — basic properties +// ----------------------------------------------------------------------- + +func TestHostCertProperties(t *testing.T) { + ca, _ := newTestCA(t) + + hostnames := []string{"example.com", "sub.example.com", "localhost"} + for _, h := range hostnames { + t.Run(h, func(t *testing.T) { + tlsCert, err := ca.TLSCertForHost(h) + if err != nil { + t.Fatalf("TLSCertForHost(%q): %v", h, err) + } + + leaf := tlsCert.Leaf + if leaf == nil { + t.Fatal("Leaf is nil") + } + if leaf.Subject.CommonName != h { + t.Errorf("CN = %q, want %q", leaf.Subject.CommonName, h) + } + + found := slices.Contains(leaf.DNSNames, h) + if !found { + t.Errorf("hostname %q not in DNSNames %v", h, leaf.DNSNames) + } + + if time.Until(leaf.NotAfter) < 364*24*time.Hour { + t.Errorf("host cert validity too short: expires %v", leaf.NotAfter) + } + }) + } +} + +// ----------------------------------------------------------------------- +// Per-host cert — IP address +// ----------------------------------------------------------------------- + +func TestHostCertIP(t *testing.T) { + ca, _ := newTestCA(t) + + tlsCert, err := ca.TLSCertForHost("127.0.0.1") + if err != nil { + t.Fatalf("TLSCertForHost(127.0.0.1): %v", err) + } + leaf := tlsCert.Leaf + if len(leaf.IPAddresses) == 0 { + t.Fatal("expected IP in IPAddresses, got none") + } + if leaf.IPAddresses[0].String() != "127.0.0.1" { + t.Errorf("got IP %v, want 127.0.0.1", leaf.IPAddresses[0]) + } + if len(leaf.DNSNames) != 0 { + t.Errorf("IP cert should have no DNSNames, got %v", leaf.DNSNames) + } +} + +// ----------------------------------------------------------------------- +// Chain validation — the acceptance criterion from the design doc +// ----------------------------------------------------------------------- + +func TestCertChainValidation(t *testing.T) { + ca, _ := newTestCA(t) + + tlsCert, err := ca.TLSCertForHost("example.com") + if err != nil { + t.Fatalf("TLSCertForHost: %v", err) + } + + // Build a cert pool containing only our CA — simulates what a browser + // with the CA installed would do. + pool := x509.NewCertPool() + pool.AddCert(ca.X509Cert()) + + opts := x509.VerifyOptions{ + DNSName: "example.com", + Roots: pool, + } + + leaf := tlsCert.Leaf + if _, err := leaf.Verify(opts); err != nil { + t.Errorf("chain verification failed: %v", err) + } +} + +// ----------------------------------------------------------------------- +// TLS handshake — end-to-end using net/tls in-process +// ----------------------------------------------------------------------- + +func TestTLSHandshake(t *testing.T) { + ca, _ := newTestCA(t) + + serverTLSCert, err := ca.TLSCertForHost("testhost.local") + if err != nil { + t.Fatalf("TLSCertForHost: %v", err) + } + serverCfg := &tls.Config{ + Certificates: []tls.Certificate{*serverTLSCert}, + } + + pool := x509.NewCertPool() + pool.AddCert(ca.X509Cert()) + clientCfg := &tls.Config{ + RootCAs: pool, + ServerName: "testhost.local", + } + + // tls.Dial inside newTLSPipe already completes the handshake. + // If it returns without error, the chain validated successfully. + serverConn, clientConn := newTLSPipe(t, serverCfg, clientCfg) + defer serverConn.Close() + defer clientConn.Close() +} + +// newTLSPipe creates an in-process TLS client/server pair over TCP loopback. +func newTLSPipe(t *testing.T, serverCfg, clientCfg *tls.Config) (*tls.Conn, *tls.Conn) { + t.Helper() + + ch := make(chan *tls.Conn, 1) + ln, err := tls.Listen("tcp", "127.0.0.1:0", serverCfg) + if err != nil { + t.Fatalf("tls.Listen: %v", err) + } + t.Cleanup(func() { ln.Close() }) + + go func() { + conn, err := ln.Accept() + if err != nil { + ch <- nil + return + } + tlsConn := conn.(*tls.Conn) + tlsConn.Handshake() // complete server side + ch <- tlsConn + }() + + client, err := tls.Dial("tcp", ln.Addr().String(), clientCfg) + if err != nil { + t.Fatalf("tls.Dial: %v", err) + } + + server := <-ch + if server == nil { + t.Fatal("server accept failed") + } + + return server, client +} + +// ----------------------------------------------------------------------- +// Cache — same *tls.Certificate pointer returned for same host +// ----------------------------------------------------------------------- + +func TestHostCertCache(t *testing.T) { + ca, _ := newTestCA(t) + + first, err := ca.TLSCertForHost("cache-test.com") + if err != nil { + t.Fatalf("first call: %v", err) + } + second, err := ca.TLSCertForHost("cache-test.com") + if err != nil { + t.Fatalf("second call: %v", err) + } + + if first != second { + t.Error("expected same *tls.Certificate pointer from cache") + } +} + +// ----------------------------------------------------------------------- +// Cache — concurrent access does not race or duplicate +// ----------------------------------------------------------------------- + +func TestHostCertCacheConcurrent(t *testing.T) { + ca, _ := newTestCA(t) + + const goroutines = 20 + results := make([]*tls.Certificate, goroutines) + var wg sync.WaitGroup + + for i := range goroutines { + i := i + wg.Go(func() { + c, err := ca.TLSCertForHost("concurrent.example.com") + if err != nil { + t.Errorf("goroutine %d: %v", i, err) + return + } + results[i] = c + }) + } + wg.Wait() + + for i, c := range results { + if c == nil { + t.Errorf("goroutine %d returned nil", i) + continue + } + if c != results[0] { + t.Errorf("goroutine %d returned different pointer", i) + } + } +} + +// ----------------------------------------------------------------------- +// Host + port — port should be stripped before cert lookup +// ----------------------------------------------------------------------- + +func TestHostCertStripsPort(t *testing.T) { + ca, _ := newTestCA(t) + + withPort, err := ca.TLSCertForHost("strip-test.com:443") + if err != nil { + t.Fatalf("with port: %v", err) + } + withoutPort, err := ca.TLSCertForHost("strip-test.com") + if err != nil { + t.Fatalf("without port: %v", err) + } + + if withPort != withoutPort { + t.Error("port stripping broken: different certs returned for same host") + } + + if withPort.Leaf.Subject.CommonName != "strip-test.com" { + t.Errorf("CN should be bare hostname, got %q", withPort.Leaf.Subject.CommonName) + } +} + +// ----------------------------------------------------------------------- +// CertPEM — returned bytes are valid PEM +// ----------------------------------------------------------------------- + +func TestCertPEM(t *testing.T) { + ca, _ := newTestCA(t) + + pem, err := ca.CertPEM() + if err != nil { + t.Fatalf("CertPEM: %v", err) + } + if len(pem) == 0 { + t.Fatal("CertPEM returned empty bytes") + } + + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(pem) { + t.Error("AppendCertsFromPEM failed — invalid PEM") + } +} diff --git a/internal/cert/install.go b/internal/cert/install.go new file mode 100644 index 0000000..8ea25b9 --- /dev/null +++ b/internal/cert/install.go @@ -0,0 +1,149 @@ +package cert + +import ( + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" +) + +// InstallCA attempts to install the Shiv CA into the user-level trust store. +// Never requires sudo. Returns a human-readable status string, or an error +// with manual instructions if automatic installation is not possible. +func (ca *CA) InstallCA() (string, error) { + certPath, err := CertPEMPath() + if err != nil { + return "", err + } + + switch runtime.GOOS { + case "linux": + return installLinux(certPath) + case "darwin": + return installDarwin(certPath) + case "windows": + return installWindows(certPath) + default: + return "", manualInstructions(certPath, fmt.Sprintf("unsupported OS: %s", runtime.GOOS)) + } +} + +func installLinux(certPath string) (string, error) { + if _, err := exec.LookPath("certutil"); err != nil { + return "", manualInstructions(certPath, "certutil not found (install libnss3-tools on Debian/Ubuntu or nss on Arch)") + } + + var installed []string + var errs []string + + // Chrome/Chromium user NSS database. + home, err := os.UserHomeDir() + if err == nil { + nssdb := filepath.Join(home, ".pki", "nssdb") + if err := os.MkdirAll(nssdb, 0700); err == nil { + cmd := exec.Command("certutil", "-A", "-n", "Shiv CA", "-t", "CT,,", "-i", certPath, "-d", "sql:"+nssdb) + if cmdOut, err := cmd.CombinedOutput(); err != nil { + errs = append(errs, fmt.Sprintf("Chrome NSS: %v\n%s", err, cmdOut)) + } else { + installed = append(installed, "Chrome/Chromium") + } + } + } + + // Firefox profiles. + profiles, _ := firefoxProfiles() + for _, profile := range profiles { + cmd := exec.Command("certutil", "-A", "-n", "Shiv CA", "-t", "CT,,", "-i", certPath, "-d", profile) + if out, err := cmd.CombinedOutput(); err != nil { + errs = append(errs, fmt.Sprintf("Firefox profile %s: %v\n%s", profile, err, out)) + } else { + installed = append(installed, "Firefox ("+filepath.Base(profile)+")") + } + } + + if len(installed) == 0 { + msg := "certutil ran but no stores were updated" + if len(errs) > 0 { + msg = strings.Join(errs, "\n") + } + return "", manualInstructions(certPath, msg) + } + + result := "CA certificate installed for: " + strings.Join(installed, ", ") + "." + if len(errs) > 0 { + result += "\nSome stores failed:\n" + strings.Join(errs, "\n") + } + return result, nil +} + +func installDarwin(certPath string) (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", manualInstructions(certPath, fmt.Sprintf("cannot find home directory: %v", err)) + } + + loginKeychain := filepath.Join(home, "Library", "Keychains", "login.keychain-db") + cmd := exec.Command("security", "add-trusted-cert", + "-r", "trustRoot", + "-k", loginKeychain, + certPath, + ) + if out, err := cmd.CombinedOutput(); err != nil { + return "", manualInstructions(certPath, fmt.Sprintf("security add-trusted-cert failed: %v\n%s", err, out)) + } + return "CA certificate installed to macOS login keychain.", nil +} + +func installWindows(certPath string) (string, error) { + // certutil is built into Windows, no extra tools needed. + // -user flag installs to current user store, no admin required. + cmd := exec.Command("certutil", "-addstore", "-user", "-f", "ROOT", certPath) + if out, err := cmd.CombinedOutput(); err != nil { + return "", manualInstructions(certPath, fmt.Sprintf("certutil failed: %v\n%s", err, out)) + } + return "CA certificate installed to Windows user root store.", nil +} + +func firefoxProfiles() ([]string, error) { + home, err := os.UserHomeDir() + if err != nil { + return nil, err + } + + var profiles []string + + profileDirs := []string{ + filepath.Join(home, ".mozilla", "firefox"), + filepath.Join(home, ".librewolf"), + filepath.Join(home, "snap", "firefox", "common", ".mozilla", "firefox"), + filepath.Join(home, ".var", "app", "org.mozilla.firefox", ".mozilla", "firefox"), + filepath.Join(home, ".var", "app", "io.gitlab.librewolf-community", ".librewolf"), + } + + for _, dir := range profileDirs { + entries, err := os.ReadDir(dir) + if err != nil { + continue + } + for _, dirEntry := range entries { + if dirEntry.IsDir() { + profiles = append(profiles, filepath.Join(dir, dirEntry.Name())) + } + } + } + + return profiles, nil +} + +func manualInstructions(certPath, reason string) error { + return errors.New(reason + "\n\n" + + "To install manually, import this file into your browser:\n" + + " " + certPath + "\n\n" + + "Browser instructions:\n" + + " Chrome/Chromium: Settings → Privacy → Security → Manage certificates → Authorities → Import\n" + + " Firefox: Settings → Privacy & Security → Certificates → View Certificates → Authorities → Import\n" + + " NixOS: security.pki.certificateFiles = [ \"" + certPath + "\" ]") +} diff --git a/internal/events/bus.go b/internal/events/bus.go new file mode 100644 index 0000000..4856c53 --- /dev/null +++ b/internal/events/bus.go @@ -0,0 +1,305 @@ +package events + +import "sync" + +// ── Handler interfaces ──────────────────────────────────────────────────────── + +// RequestMiddleware is called synchronously for every HTTP request before +// forwarding. Returning Drop=true causes the proxy to send a 403 and stop. +// The implementation may modify Request and Body (e.g. intercept gate). +type RequestMiddleware interface { + HandleRequest(RequestEvent) RequestResult +} + +// ResponseObserver is called synchronously after every completed HTTP +// transaction. Ordering is preserved: observers see responses in the exact +// order the proxy received them. +type ResponseObserver interface { + ObserveResponse(ResponseEvent) +} + +// WebSocketConnectionObserver is called synchronously when a WebSocket +// connection is established. It returns a uint64 connection ID that the +// bus passes back to WebSocketFrameObservers as ConnectionID. Return 0 +// to indicate the connection should not be tracked. +type WebSocketConnectionObserver interface { + ObserveWebSocketConnection(WebSocketConnectionEvent) uint64 +} + +// WebSocketFrameObserver is called synchronously for every WebSocket frame. +// Ordering is preserved within a connection. The returned WebSocketFrameResult +// carries the (possibly modified) payload to forward. +type WebSocketFrameObserver interface { + ObserveWebSocketFrame(WebSocketFrameEvent) WebSocketFrameResult +} + +// PluginLogObserver is called synchronously when a plugin emits a log line. +// The store implements this to buffer lines in memory and notify the UI. +type PluginLogObserver interface { + ObservePluginLog(PluginLogEvent) +} + +// PluginEnabledObserver is called synchronously when the UI toggles a plugin. +// Both the engine (to skip hooks) and the store (to persist state) implement this. +type PluginEnabledObserver interface { + ObservePluginEnabled(SetPluginEnabledEvent) +} + +// LoadPluginObserver is called synchronously when the user imports a new +// plugin. The engine observes it to copy, load, and register the plugin. +type LoadPluginObserver interface { + ObserveLoadPlugin(LoadPluginEvent) +} + +// ProxyCommandObserver is implemented by the proxy to react to UI-initiated +// control events. Restart starts or restarts the proxy on the given address. +// Stop shuts the proxy down. +type ProxyCommandObserver interface { + ObserveProxyRestart(ProxyRestartEvent) + ObserveProxyStop(ProxyStopEvent) +} + +// ── Func adapters ───────────────────────────────────────────────────────────── +// Allow main.go to register plain closures without declaring named types. + +type RequestMiddlewareFunc func(RequestEvent) RequestResult + +func (f RequestMiddlewareFunc) HandleRequest(e RequestEvent) RequestResult { + return f(e) +} + +type ResponseObserverFunc func(ResponseEvent) + +func (f ResponseObserverFunc) ObserveResponse(e ResponseEvent) { + f(e) +} + +type WebSocketConnectionObserverFunc func(WebSocketConnectionEvent) uint64 + +func (f WebSocketConnectionObserverFunc) ObserveWebSocketConnection(e WebSocketConnectionEvent) uint64 { + return f(e) +} + +type WebSocketFrameObserverFunc func(WebSocketFrameEvent) WebSocketFrameResult + +func (f WebSocketFrameObserverFunc) ObserveWebSocketFrame(e WebSocketFrameEvent) WebSocketFrameResult { + return f(e) +} + +type PluginLogObserverFunc func(PluginLogEvent) + +func (f PluginLogObserverFunc) ObservePluginLog(e PluginLogEvent) { + f(e) +} + +type PluginEnabledObserverFunc func(SetPluginEnabledEvent) + +func (f PluginEnabledObserverFunc) ObservePluginEnabled(e SetPluginEnabledEvent) { + f(e) +} + +type LoadPluginObserverFunc func(LoadPluginEvent) + +func (f LoadPluginObserverFunc) ObserveLoadPlugin(e LoadPluginEvent) { + f(e) +} + +// ── Bus ─────────────────────────────────────────────────────────────────────── + +// Bus holds registered handlers and dispatches events to them. +// All methods are safe for concurrent use. All emit methods run +// synchronously to preserve the ordering guarantee that existed before +// the refactor: store.Log emits to the Updates channel which the UI reads +// in arrival order. Making observers async would break this. +type Bus struct { + mu sync.RWMutex + requestMiddlewares []RequestMiddleware + responseObservers []ResponseObserver + wsConnObservers []WebSocketConnectionObserver + wsFrameObservers []WebSocketFrameObserver + pluginLogObservers []PluginLogObserver + pluginEnabledObservers []PluginEnabledObserver + pluginLoadObservers []LoadPluginObserver + proxyCommandObservers []ProxyCommandObserver +} + +// NewBus returns a ready-to-use Bus with no handlers registered. +func NewBus() *Bus { + return &Bus{} +} + +// Register adds a handler to the bus. The handler is matched against all +// handler interfaces — a single type may implement multiple interfaces +// and will be registered for each one it satisfies. Panics if h implements +// none of the interfaces (programming error). +func (b *Bus) Register(h any) { + b.mu.Lock() + defer b.mu.Unlock() + + matched := false + if v, ok := h.(RequestMiddleware); ok { + b.requestMiddlewares = append(b.requestMiddlewares, v) + matched = true + } + if v, ok := h.(ResponseObserver); ok { + b.responseObservers = append(b.responseObservers, v) + matched = true + } + if v, ok := h.(WebSocketConnectionObserver); ok { + b.wsConnObservers = append(b.wsConnObservers, v) + matched = true + } + if v, ok := h.(WebSocketFrameObserver); ok { + b.wsFrameObservers = append(b.wsFrameObservers, v) + matched = true + } + if v, ok := h.(PluginLogObserver); ok { + b.pluginLogObservers = append(b.pluginLogObservers, v) + matched = true + } + if v, ok := h.(PluginEnabledObserver); ok { + b.pluginEnabledObservers = append(b.pluginEnabledObservers, v) + matched = true + } + if v, ok := h.(LoadPluginObserver); ok { + b.pluginLoadObservers = append(b.pluginLoadObservers, v) + matched = true + } + if v, ok := h.(ProxyCommandObserver); ok { + b.proxyCommandObservers = append(b.proxyCommandObservers, v) + matched = true + } + if !matched { + panic("events: Register called with a type that implements no handler interface") + } +} + +// EmitRequest runs all RequestMiddlewares in registration order. +// The first middleware to return Drop=true short-circuits — subsequent +// middlewares are not called. If no middleware drops the request, the +// final Request and Body are taken from the last middleware result. +// If no middlewares are registered the original event values are returned +// as a pass-through result. +func (b *Bus) EmitRequest(e RequestEvent) RequestResult { + b.mu.RLock() + mws := b.requestMiddlewares + b.mu.RUnlock() + + result := RequestResult{Request: e.Request, Body: e.Body} + for _, mw := range mws { + result = mw.HandleRequest(RequestEvent{Request: result.Request, Body: result.Body}) + if result.Drop { + return result + } + } + return result +} + +// EmitResponse calls all ResponseObservers in registration order. +// Runs synchronously; returns when all observers have returned. +func (b *Bus) EmitResponse(e ResponseEvent) { + b.mu.RLock() + obs := b.responseObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObserveResponse(e) + } +} + +// EmitWebSocketConnection calls all WebSocketConnectionObservers in +// registration order. The first non-zero uint64 returned is used as the +// connection ID for subsequent frame events. Runs synchronously. +func (b *Bus) EmitWebSocketConnection(e WebSocketConnectionEvent) uint64 { + b.mu.RLock() + obs := b.wsConnObservers + b.mu.RUnlock() + + var connID uint64 + for _, o := range obs { + if id := o.ObserveWebSocketConnection(e); id != 0 && connID == 0 { + connID = id + } + } + return connID +} + +// EmitWebSocketFrame calls all WebSocketFrameObservers in registration order. +// Runs synchronously to preserve per-connection frame ordering. Each observer +// may return a modified payload — the last non-nil payload wins. The final +// WebSocketFrameResult carries the payload to forward. +func (b *Bus) EmitWebSocketFrame(e WebSocketFrameEvent) WebSocketFrameResult { + b.mu.RLock() + obs := b.wsFrameObservers + b.mu.RUnlock() + + result := WebSocketFrameResult{Payload: e.Payload} + for _, o := range obs { + r := o.ObserveWebSocketFrame(e) + if r.Payload != nil { + result.Payload = r.Payload + e.Payload = r.Payload + } + } + return result +} + +// EmitPluginLog calls all PluginLogObservers in registration order. +// Runs synchronously. +func (b *Bus) EmitPluginLog(e PluginLogEvent) { + b.mu.RLock() + obs := b.pluginLogObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObservePluginLog(e) + } +} + +// EmitSetPluginEnabled calls all PluginEnabledObservers in registration order. +// Runs synchronously. +func (b *Bus) EmitSetPluginEnabled(e SetPluginEnabledEvent) { + b.mu.RLock() + obs := b.pluginEnabledObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObservePluginEnabled(e) + } +} + +// EmitLoadPlugin calls all LoadPluginObservers in registration order. +// Runs synchronously. +func (b *Bus) EmitLoadPlugin(e LoadPluginEvent) { + b.mu.RLock() + obs := b.pluginLoadObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObserveLoadPlugin(e) + } +} + +// EmitProxyRestart calls all ProxyCommandObservers to restart on the given address. +// Runs synchronously. +func (b *Bus) EmitProxyRestart(e ProxyRestartEvent) { + b.mu.RLock() + obs := b.proxyCommandObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObserveProxyRestart(e) + } +} + +// EmitProxyStop calls all ProxyCommandObservers to stop the proxy. +// Runs synchronously. +func (b *Bus) EmitProxyStop(e ProxyStopEvent) { + b.mu.RLock() + obs := b.proxyCommandObservers + b.mu.RUnlock() + + for _, o := range obs { + o.ObserveProxyStop(e) + } +} diff --git a/internal/events/bus_test.go b/internal/events/bus_test.go new file mode 100644 index 0000000..633d950 --- /dev/null +++ b/internal/events/bus_test.go @@ -0,0 +1,746 @@ +package events_test + +import ( + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +func newRequest(t *testing.T, method, url string) *http.Request { + t.Helper() + return httptest.NewRequest(method, url, nil) +} + +func newRequestEvent(t *testing.T) events.RequestEvent { + t.Helper() + return events.RequestEvent{ + Request: newRequest(t, http.MethodGet, "http://example.com/path"), + Body: []byte("request body"), + } +} + +func newResponseEvent() events.ResponseEvent { + return events.ResponseEvent{ + Timestamp: time.Now(), + Host: "example.com", + Proto: "HTTP/1.1", + Method: "GET", + URL: "http://example.com/path", + StatusCode: 200, + RespBody: []byte("response body"), + DurationMs: 42, + TLS: false, + } +} + +func newWSConnectionEvent() events.WebSocketConnectionEvent { + return events.WebSocketConnectionEvent{ + Host: "example.com:443", + URL: "wss://example.com/ws", + TLS: true, + Timestamp: time.Now(), + } +} + +func newWSFrameEvent(connID uint64) events.WebSocketFrameEvent { + return events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: []byte("hello"), + } +} + +// ── Register ────────────────────────────────────────────────────────────────── + +func TestRegister_PanicsOnUnknownType(t *testing.T) { + b := events.NewBus() + assert.Panics(t, func() { + b.Register(struct{}{}) + }) +} + +func TestRegister_AcceptsRequestMiddleware(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + }) +} + +func TestRegister_AcceptsResponseObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) {})) + }) +} + +func TestRegister_AcceptsWebSocketConnectionObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 1 })) + }) +} + +func TestRegister_AcceptsWebSocketFrameObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: e.Payload} + })) + }) +} + +func TestRegister_AcceptsPluginLogObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) {})) + }) +} + +func TestRegister_AcceptsPluginEnabledObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) {})) + }) +} + +func TestRegister_AcceptsLoadPluginObserver(t *testing.T) { + b := events.NewBus() + require.NotPanics(t, func() { + b.Register(events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) {})) + }) +} + +func TestRegister_MultiInterfaceTypeRegisteredForAll(t *testing.T) { + b := events.NewBus() + + var requestCalled, responseCalled, wsConnCalled, wsFrameCalled bool + var pluginLogCalled, pluginEnabledCalled, loadPluginCalled bool + + all := &allHandler{ + onRequest: func() { requestCalled = true }, + onResponse: func() { responseCalled = true }, + onWSConn: func() { wsConnCalled = true }, + onWSFrame: func() { wsFrameCalled = true }, + onPluginLog: func() { pluginLogCalled = true }, + onPluginEnabled: func() { pluginEnabledCalled = true }, + onLoadPlugin: func() { loadPluginCalled = true }, + } + + b.Register(all) + + b.EmitRequest(newRequestEvent(t)) + b.EmitResponse(newResponseEvent()) + b.EmitWebSocketConnection(newWSConnectionEvent()) + b.EmitWebSocketFrame(newWSFrameEvent(1)) + b.EmitPluginLog(events.PluginLogEvent{Name: "test.lua", Message: "hello"}) + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + b.EmitLoadPlugin(events.LoadPluginEvent{SourcePath: "/tmp/test.lua"}) + + assert.True(t, requestCalled) + assert.True(t, responseCalled) + assert.True(t, wsConnCalled) + assert.True(t, wsFrameCalled) + assert.True(t, pluginLogCalled) + assert.True(t, pluginEnabledCalled) + assert.True(t, loadPluginCalled) +} + +// ── EmitRequest ─────────────────────────────────────────────────────────────── + +func TestEmitRequest_NoHandlers_PassThrough(t *testing.T) { + b := events.NewBus() + ev := newRequestEvent(t) + result := b.EmitRequest(ev) + assert.False(t, result.Drop) + assert.Equal(t, ev.Request, result.Request) + assert.Equal(t, ev.Body, result.Body) +} + +func TestEmitRequest_SingleMiddleware_PassThrough(t *testing.T) { + b := events.NewBus() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + ev := newRequestEvent(t) + result := b.EmitRequest(ev) + assert.False(t, result.Drop) + assert.Equal(t, ev.Request, result.Request) +} + +func TestEmitRequest_MiddlewareDrop_ReturnsDrop(t *testing.T) { + b := events.NewBus() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Drop: true, Request: e.Request, Body: e.Body} + })) + result := b.EmitRequest(newRequestEvent(t)) + assert.True(t, result.Drop) +} + +func TestEmitRequest_FirstDropShortCircuits(t *testing.T) { + b := events.NewBus() + secondCalled := false + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Drop: true, Request: e.Request, Body: e.Body} + })) + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + secondCalled = true + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + b.EmitRequest(newRequestEvent(t)) + assert.False(t, secondCalled) +} + +func TestEmitRequest_MiddlewareModifiesBody(t *testing.T) { + b := events.NewBus() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: []byte("modified body")} + })) + result := b.EmitRequest(newRequestEvent(t)) + assert.False(t, result.Drop) + assert.Equal(t, []byte("modified body"), result.Body) +} + +func TestEmitRequest_MiddlewareModifiesRequest(t *testing.T) { + b := events.NewBus() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + e.Request.Header.Set("X-Injected", "yes") + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + result := b.EmitRequest(newRequestEvent(t)) + assert.Equal(t, "yes", result.Request.Header.Get("X-Injected")) +} + +func TestEmitRequest_ChainedMiddlewares_ModificationsAccumulate(t *testing.T) { + b := events.NewBus() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: []byte("step1")} + })) + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: append(e.Body, []byte("-step2")...)} + })) + result := b.EmitRequest(newRequestEvent(t)) + assert.Equal(t, []byte("step1-step2"), result.Body) +} + +func TestEmitRequest_MiddlewareCalledInRegistrationOrder(t *testing.T) { + b := events.NewBus() + var order []int + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + order = append(order, 1) + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + order = append(order, 2) + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + order = append(order, 3) + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + b.EmitRequest(newRequestEvent(t)) + assert.Equal(t, []int{1, 2, 3}, order) +} + +// ── EmitResponse ────────────────────────────────────────────────────────────── + +func TestEmitResponse_NoHandlers_NoPanic(t *testing.T) { + b := events.NewBus() + assert.NotPanics(t, func() { b.EmitResponse(newResponseEvent()) }) +} + +func TestEmitResponse_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.ResponseEvent + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { received = e })) + ev := newResponseEvent() + b.EmitResponse(ev) + assert.Equal(t, ev.Host, received.Host) + assert.Equal(t, ev.StatusCode, received.StatusCode) + assert.Equal(t, ev.DurationMs, received.DurationMs) +} + +func TestEmitResponse_MultipleObserversAllCalled(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { calls++ })) + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { calls++ })) + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { calls++ })) + b.EmitResponse(newResponseEvent()) + assert.Equal(t, 3, calls) +} + +func TestEmitResponse_ObserversCalledInRegistrationOrder(t *testing.T) { + b := events.NewBus() + var order []int + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { order = append(order, 1) })) + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { order = append(order, 2) })) + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { order = append(order, 3) })) + b.EmitResponse(newResponseEvent()) + assert.Equal(t, []int{1, 2, 3}, order) +} + +func TestEmitResponse_IsSynchronous(t *testing.T) { + b := events.NewBus() + done := false + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { done = true })) + b.EmitResponse(newResponseEvent()) + assert.True(t, done) +} + +// ── EmitWebSocketConnection ─────────────────────────────────────────────────── + +func TestEmitWebSocketConnection_NoHandlers_ReturnsZero(t *testing.T) { + b := events.NewBus() + assert.Equal(t, uint64(0), b.EmitWebSocketConnection(newWSConnectionEvent())) +} + +func TestEmitWebSocketConnection_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.WebSocketConnectionEvent + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { + received = e + return 1 + })) + ev := newWSConnectionEvent() + b.EmitWebSocketConnection(ev) + assert.Equal(t, ev.Host, received.Host) + assert.Equal(t, ev.URL, received.URL) + assert.Equal(t, ev.TLS, received.TLS) +} + +func TestEmitWebSocketConnection_ReturnsFirstNonZeroID(t *testing.T) { + b := events.NewBus() + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 0 })) + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 42 })) + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 99 })) + assert.Equal(t, uint64(42), b.EmitWebSocketConnection(newWSConnectionEvent())) +} + +func TestEmitWebSocketConnection_AllObserversCalledEvenAfterNonZero(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { calls++; return 1 })) + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { calls++; return 2 })) + b.EmitWebSocketConnection(newWSConnectionEvent()) + assert.Equal(t, 2, calls) +} + +func TestEmitWebSocketConnection_ZeroReturnNotUsedAsID(t *testing.T) { + b := events.NewBus() + b.Register(events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 0 })) + assert.Equal(t, uint64(0), b.EmitWebSocketConnection(newWSConnectionEvent())) +} + +// ── EmitWebSocketFrame ──────────────────────────────────────────────────────── + +func TestEmitWebSocketFrame_NoHandlers_ReturnsOriginalPayload(t *testing.T) { + b := events.NewBus() + ev := newWSFrameEvent(1) + result := b.EmitWebSocketFrame(ev) + assert.Equal(t, ev.Payload, result.Payload) +} + +func TestEmitWebSocketFrame_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.WebSocketFrameEvent + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + received = e + return events.WebSocketFrameResult{Payload: e.Payload} + })) + ev := newWSFrameEvent(7) + b.EmitWebSocketFrame(ev) + assert.Equal(t, uint64(7), received.ConnectionID) + assert.Equal(t, events.WebSocketClient, received.Direction) + assert.Equal(t, events.WebSocketText, received.Opcode) + assert.Equal(t, []byte("hello"), received.Payload) +} + +func TestEmitWebSocketFrame_ObserverModifiesPayload(t *testing.T) { + b := events.NewBus() + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: []byte("modified")} + })) + result := b.EmitWebSocketFrame(newWSFrameEvent(1)) + assert.Equal(t, []byte("modified"), result.Payload) +} + +func TestEmitWebSocketFrame_ChainedObservers_ModificationsAccumulate(t *testing.T) { + b := events.NewBus() + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: []byte("step1")} + })) + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: append(e.Payload, []byte("-step2")...)} + })) + result := b.EmitWebSocketFrame(newWSFrameEvent(1)) + assert.Equal(t, []byte("step1-step2"), result.Payload) +} + +func TestEmitWebSocketFrame_NilPayloadResult_KeepsPrevious(t *testing.T) { + b := events.NewBus() + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: []byte("keep-me")} + })) + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + return events.WebSocketFrameResult{Payload: nil} + })) + result := b.EmitWebSocketFrame(newWSFrameEvent(1)) + assert.Equal(t, []byte("keep-me"), result.Payload) +} + +func TestEmitWebSocketFrame_MultipleObserversAllCalled(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + calls++ + return events.WebSocketFrameResult{Payload: e.Payload} + })) + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + calls++ + return events.WebSocketFrameResult{Payload: e.Payload} + })) + b.EmitWebSocketFrame(newWSFrameEvent(1)) + assert.Equal(t, 2, calls) +} + +func TestEmitWebSocketFrame_OrderingPreserved(t *testing.T) { + b := events.NewBus() + var received []uint64 + b.Register(events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + received = append(received, e.ConnectionID) + return events.WebSocketFrameResult{Payload: e.Payload} + })) + b.EmitWebSocketFrame(events.WebSocketFrameEvent{ConnectionID: 1, Payload: []byte("a")}) + b.EmitWebSocketFrame(events.WebSocketFrameEvent{ConnectionID: 2, Payload: []byte("b")}) + b.EmitWebSocketFrame(events.WebSocketFrameEvent{ConnectionID: 3, Payload: []byte("c")}) + assert.Equal(t, []uint64{1, 2, 3}, received) +} + +// ── EmitPluginLog ───────────────────────────────────────────────────────────── + +func TestEmitPluginLog_NoHandlers_NoPanic(t *testing.T) { + b := events.NewBus() + assert.NotPanics(t, func() { + b.EmitPluginLog(events.PluginLogEvent{Name: "test.lua", Message: "hello"}) + }) +} + +func TestEmitPluginLog_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.PluginLogEvent + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { received = e })) + b.EmitPluginLog(events.PluginLogEvent{Name: "myplugin.lua", Message: "log line"}) + assert.Equal(t, "myplugin.lua", received.Name) + assert.Equal(t, "log line", received.Message) +} + +func TestEmitPluginLog_MultipleObserversAllCalled(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { calls++ })) + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { calls++ })) + b.EmitPluginLog(events.PluginLogEvent{Name: "test.lua", Message: "msg"}) + assert.Equal(t, 2, calls) +} + +func TestEmitPluginLog_IsSynchronous(t *testing.T) { + b := events.NewBus() + done := false + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { done = true })) + b.EmitPluginLog(events.PluginLogEvent{Name: "test.lua", Message: "msg"}) + assert.True(t, done) +} + +// ── EmitSetPluginEnabled ────────────────────────────────────────────────────── + +func TestEmitSetPluginEnabled_NoHandlers_NoPanic(t *testing.T) { + b := events.NewBus() + assert.NotPanics(t, func() { + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + }) +} + +func TestEmitSetPluginEnabled_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.SetPluginEnabledEvent + b.Register(events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) { received = e })) + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "myplugin.lua", Enabled: false}) + assert.Equal(t, "myplugin.lua", received.Name) + assert.False(t, received.Enabled) +} + +func TestEmitSetPluginEnabled_MultipleObserversAllCalled(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) { calls++ })) + b.Register(events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) { calls++ })) + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + assert.Equal(t, 2, calls) +} + +func TestEmitSetPluginEnabled_EnabledFlagPropagated(t *testing.T) { + b := events.NewBus() + var gotEnabled bool + b.Register(events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) { gotEnabled = e.Enabled })) + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + assert.True(t, gotEnabled) + b.EmitSetPluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + assert.False(t, gotEnabled) +} + +// ── EmitLoadPlugin ──────────────────────────────────────────────────────────── + +func TestEmitLoadPlugin_NoHandlers_NoPanic(t *testing.T) { + b := events.NewBus() + assert.NotPanics(t, func() { + b.EmitLoadPlugin(events.LoadPluginEvent{SourcePath: "/tmp/test.lua"}) + }) +} + +func TestEmitLoadPlugin_ObserverReceivesEvent(t *testing.T) { + b := events.NewBus() + var received events.LoadPluginEvent + b.Register(events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) { received = e })) + b.EmitLoadPlugin(events.LoadPluginEvent{SourcePath: "/home/user/myplugin.lua"}) + assert.Equal(t, "/home/user/myplugin.lua", received.SourcePath) +} + +func TestEmitLoadPlugin_MultipleObserversAllCalled(t *testing.T) { + b := events.NewBus() + var calls int + b.Register(events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) { calls++ })) + b.Register(events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) { calls++ })) + b.EmitLoadPlugin(events.LoadPluginEvent{SourcePath: "/tmp/test.lua"}) + assert.Equal(t, 2, calls) +} + +func TestEmitLoadPlugin_IsSynchronous(t *testing.T) { + b := events.NewBus() + done := false + b.Register(events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) { done = true })) + b.EmitLoadPlugin(events.LoadPluginEvent{SourcePath: "/tmp/test.lua"}) + assert.True(t, done) +} + +// ── Func adapters ───────────────────────────────────────────────────────────── + +func TestRequestMiddlewareFunc_ImplementsInterface(t *testing.T) { + var _ events.RequestMiddleware = events.RequestMiddlewareFunc(nil) +} + +func TestResponseObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.ResponseObserver = events.ResponseObserverFunc(nil) +} + +func TestWebSocketConnectionObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.WebSocketConnectionObserver = events.WebSocketConnectionObserverFunc(nil) +} + +func TestWebSocketFrameObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.WebSocketFrameObserver = events.WebSocketFrameObserverFunc(nil) +} + +func TestPluginLogObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.PluginLogObserver = events.PluginLogObserverFunc(nil) +} + +func TestPluginEnabledObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.PluginEnabledObserver = events.PluginEnabledObserverFunc(nil) +} + +func TestLoadPluginObserverFunc_ImplementsInterface(t *testing.T) { + var _ events.LoadPluginObserver = events.LoadPluginObserverFunc(nil) +} + +func TestRequestMiddlewareFunc_DelegatesCall(t *testing.T) { + called := false + f := events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + called = true + return events.RequestResult{Request: e.Request, Body: e.Body} + }) + f.HandleRequest(events.RequestEvent{Request: newRequest(t, http.MethodGet, "http://example.com")}) + assert.True(t, called) +} + +func TestResponseObserverFunc_DelegatesCall(t *testing.T) { + called := false + f := events.ResponseObserverFunc(func(e events.ResponseEvent) { called = true }) + f.ObserveResponse(newResponseEvent()) + assert.True(t, called) +} + +func TestWebSocketConnectionObserverFunc_DelegatesCall(t *testing.T) { + f := events.WebSocketConnectionObserverFunc(func(e events.WebSocketConnectionEvent) uint64 { return 55 }) + assert.Equal(t, uint64(55), f.ObserveWebSocketConnection(newWSConnectionEvent())) +} + +func TestWebSocketFrameObserverFunc_DelegatesCall(t *testing.T) { + called := false + f := events.WebSocketFrameObserverFunc(func(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + called = true + return events.WebSocketFrameResult{Payload: e.Payload} + }) + f.ObserveWebSocketFrame(newWSFrameEvent(1)) + assert.True(t, called) +} + +func TestPluginLogObserverFunc_DelegatesCall(t *testing.T) { + called := false + f := events.PluginLogObserverFunc(func(e events.PluginLogEvent) { called = true }) + f.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "hello"}) + assert.True(t, called) +} + +func TestPluginEnabledObserverFunc_DelegatesCall(t *testing.T) { + var got events.SetPluginEnabledEvent + f := events.PluginEnabledObserverFunc(func(e events.SetPluginEnabledEvent) { got = e }) + f.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + assert.Equal(t, "test.lua", got.Name) + assert.True(t, got.Enabled) +} + +func TestLoadPluginObserverFunc_DelegatesCall(t *testing.T) { + var got events.LoadPluginEvent + f := events.LoadPluginObserverFunc(func(e events.LoadPluginEvent) { got = e }) + f.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: "/tmp/test.lua"}) + assert.Equal(t, "/tmp/test.lua", got.SourcePath) +} + +// ── Concurrency ─────────────────────────────────────────────────────────────── + +func TestBus_ConcurrentEmitRequest_Safe(t *testing.T) { + b := events.NewBus() + var count atomic.Int64 + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + count.Add(1) + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for range goroutines { + go func() { + defer wg.Done() + b.EmitRequest(events.RequestEvent{Request: httptest.NewRequest(http.MethodGet, "http://example.com", nil)}) + }() + } + wg.Wait() + assert.Equal(t, int64(goroutines), count.Load()) +} + +func TestBus_ConcurrentEmitResponse_Safe(t *testing.T) { + b := events.NewBus() + var count atomic.Int64 + b.Register(events.ResponseObserverFunc(func(e events.ResponseEvent) { count.Add(1) })) + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for range goroutines { + go func() { + defer wg.Done() + b.EmitResponse(newResponseEvent()) + }() + } + wg.Wait() + assert.Equal(t, int64(goroutines), count.Load()) +} + +func TestBus_ConcurrentEmitPluginLog_Safe(t *testing.T) { + b := events.NewBus() + var count atomic.Int64 + b.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { count.Add(1) })) + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for range goroutines { + go func() { + defer wg.Done() + b.EmitPluginLog(events.PluginLogEvent{Name: "test.lua", Message: "msg"}) + }() + } + wg.Wait() + assert.Equal(t, int64(goroutines), count.Load()) +} + +func TestBus_ConcurrentRegisterAndEmit_Safe(t *testing.T) { + b := events.NewBus() + var wg sync.WaitGroup + for range 20 { + wg.Add(1) + go func() { + defer wg.Done() + b.EmitRequest(events.RequestEvent{Request: httptest.NewRequest(http.MethodGet, "http://example.com", nil)}) + }() + } + for range 20 { + wg.Add(1) + go func() { + defer wg.Done() + b.Register(events.RequestMiddlewareFunc(func(e events.RequestEvent) events.RequestResult { + return events.RequestResult{Request: e.Request, Body: e.Body} + })) + }() + } + wg.Wait() +} + +// ── WebSocket direction and opcode constants ────────────────────────────────── + +func TestWebSocketDirection_Values(t *testing.T) { + assert.Equal(t, events.WebSocketDirection(0), events.WebSocketClient) + assert.Equal(t, events.WebSocketDirection(1), events.WebSocketServer) +} + +func TestWebSocketOpcode_Values(t *testing.T) { + assert.Equal(t, events.WebSocketOpcode(1), events.WebSocketText) + assert.Equal(t, events.WebSocketOpcode(2), events.WebSocketBinary) + assert.Equal(t, events.WebSocketOpcode(8), events.WebSocketClose) + assert.Equal(t, events.WebSocketOpcode(9), events.WebSocketPing) + assert.Equal(t, events.WebSocketOpcode(10), events.WebSocketPong) +} + +// ── allHandler — test double implementing all seven interfaces ──────────────── + +type allHandler struct { + onRequest func() + onResponse func() + onWSConn func() + onWSFrame func() + onPluginLog func() + onPluginEnabled func() + onLoadPlugin func() +} + +func (h *allHandler) HandleRequest(e events.RequestEvent) events.RequestResult { + h.onRequest() + return events.RequestResult{Request: e.Request, Body: e.Body} +} + +func (h *allHandler) ObserveResponse(e events.ResponseEvent) { h.onResponse() } + +func (h *allHandler) ObserveWebSocketConnection(e events.WebSocketConnectionEvent) uint64 { + h.onWSConn() + return 1 +} + +func (h *allHandler) ObserveWebSocketFrame(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + h.onWSFrame() + return events.WebSocketFrameResult{Payload: e.Payload} +} + +func (h *allHandler) ObservePluginLog(e events.PluginLogEvent) { h.onPluginLog() } +func (h *allHandler) ObservePluginEnabled(e events.SetPluginEnabledEvent) { h.onPluginEnabled() } +func (h *allHandler) ObserveLoadPlugin(e events.LoadPluginEvent) { h.onLoadPlugin() } diff --git a/internal/events/events.go b/internal/events/events.go new file mode 100644 index 0000000..5258a21 --- /dev/null +++ b/internal/events/events.go @@ -0,0 +1,122 @@ +package events + +import ( + "net/http" + "time" +) + +// RequestEvent carries a fully-read HTTP request and its body, ready to be +// forwarded. The proxy emits this before forwarding. Consumers may inspect +// or modify the request. The result controls whether forwarding proceeds. +type RequestEvent struct { + Request *http.Request + Body []byte +} + +// ResponseEvent carries a completed HTTP transaction. The proxy emits this +// after the upstream response is fully received and the body is decompressed. +// All fields are populated; no further store knowledge is required. +type ResponseEvent struct { + Timestamp time.Time + Host string + Proto string // "HTTP/1.1" or "HTTP/2" + Method string + URL string + ReqHeaders http.Header + ReqBody []byte + StatusCode int + RespHeaders http.Header + RespBody []byte // decompressed; nil if binary + DurationMs int64 + TLS bool +} + +// RequestResult is returned by RequestMiddleware. If Drop is true the proxy +// sends a 403 and does not forward. Request and Body are the (possibly +// modified) request to forward when Drop is false. +type RequestResult struct { + Drop bool + Request *http.Request + Body []byte +} + +// WebSocketDirection indicates which endpoint sent a frame. +type WebSocketDirection int + +const ( + WebSocketClient WebSocketDirection = 0 // browser → server + WebSocketServer WebSocketDirection = 1 // server → browser +) + +// WebSocketOpcode mirrors gorilla/websocket message type constants. +type WebSocketOpcode int + +const ( + WebSocketText WebSocketOpcode = 1 + WebSocketBinary WebSocketOpcode = 2 + WebSocketPing WebSocketOpcode = 9 + WebSocketPong WebSocketOpcode = 10 + WebSocketClose WebSocketOpcode = 8 +) + +// WebSocketConnectionEvent is emitted once when a WebSocket upgrade handshake +// completes successfully. The proxy emits this after the upstream dial and +// browser upgrade both succeed. +type WebSocketConnectionEvent struct { + Host string + URL string + TLS bool + Timestamp time.Time +} + +// WebSocketFrameResult is returned by WebSocketFrameObserver. Payload is the +// (possibly modified) payload to forward. If Payload is nil the original is used. +type WebSocketFrameResult struct { + Payload []byte +} + +// WebSocketFrameEvent is emitted for every proxied WebSocket frame in either +// direction. ConnectionID is the value returned by the +// WebSocketConnectionObserver that handled the corresponding +// WebSocketConnectionEvent. +type WebSocketFrameEvent struct { + ConnectionID uint64 + Timestamp time.Time + Direction WebSocketDirection + Opcode WebSocketOpcode + Payload []byte +} + +// PluginLogEvent is emitted by the plugin engine when a Lua plugin calls +// log(). The store observes it, buffers the line in memory, and notifies +// the UI via the PluginEntries channel. +type PluginLogEvent struct { + Name string // plugin filename (e.g. "myplugin.lua") + Message string +} + +// SetPluginEnabledEvent is emitted by the UI when the user toggles a plugin. +// The engine observes it to skip or include that plugin's hooks. +// The store observes it to persist the new state to the database. +type SetPluginEnabledEvent struct { + Name string // plugin filename + Enabled bool +} + +// LoadPluginEvent is emitted by the UI when the user selects a .lua file to +// import. The engine observes it, copies the file into the plugins directory, +// loads it, and registers it for subsequent hook calls. +type LoadPluginEvent struct { + SourcePath string // absolute path to the selected .lua file +} + +// ProxyRestartEvent is emitted by the UI when the user saves proxy settings +// or toggles the proxy on. Addr is the new listen address (host:port). +// The proxy observes this and restarts on the new address. +type ProxyRestartEvent struct { + Addr string +} + +// ProxyStopEvent is emitted by the UI when the user toggles the proxy off. +// The proxy observes this and stops listening. +type ProxyStopEvent struct{} diff --git a/internal/http/compression.go b/internal/http/compression.go new file mode 100644 index 0000000..aba8ef1 --- /dev/null +++ b/internal/http/compression.go @@ -0,0 +1,142 @@ +package http + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "compress/zlib" + "io" + "net/http" + "strings" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" +) + +const MaxDisplayBytes = 64 * 1024 // 64 KB + +// Decompress decompresses body according to the Content-Encoding header. +// Returns the original body unchanged if encoding is unknown or decompression fails. +// +// Deflate note: servers claiming "deflate" Content-Encoding almost universally +// send zlib-wrapped data (RFC 1950), not raw DEFLATE (RFC 1951). We try zlib +// first and fall back to raw deflate so both are handled correctly. +func Decompress(header http.Header, body []byte) []byte { + encoding := strings.ToLower(strings.TrimSpace(header.Get("Content-Encoding"))) + + switch encoding { + case "gzip": + gr, err := gzip.NewReader(bytes.NewReader(body)) + if err != nil { + return body + } + defer gr.Close() + out, err := io.ReadAll(gr) + if err != nil { + return body + } + return out + + case "deflate": + // Try zlib (RFC 1950) first — the de-facto standard despite the name. + if out, err := decompressZlib(body); err == nil { + return out + } + // Fall back to raw DEFLATE (RFC 1951). + out, err := io.ReadAll(flate.NewReader(bytes.NewReader(body))) + if err != nil { + return body + } + return out + + case "br": + out, err := io.ReadAll(brotli.NewReader(bytes.NewReader(body))) + if err != nil { + return body + } + return out + + case "zstd": + dec, err := zstd.NewReader(bytes.NewReader(body)) + if err != nil { + return body + } + defer dec.Close() + out, err := io.ReadAll(dec) + if err != nil { + return body + } + return out + + case "", "identity": + // No encoding or explicit identity — return as-is. + return body + } + + // Unknown encoding — return original bytes unchanged. + return body +} + +// decompressZlib attempts to decompress body as zlib (RFC 1950). +func decompressZlib(body []byte) ([]byte, error) { + r, err := zlib.NewReader(bytes.NewReader(body)) + if err != nil { + return nil, err + } + defer r.Close() + return io.ReadAll(r) +} + +// IsBinary returns true if the response body should not be stored or displayed as text. +// +// If Content-Encoding is a known compression we support, the body will be +// decompressed to text before this is called, so those are not binary. +// Unknown or future Content-Encoding values are treated as binary since we +// cannot safely decode them. +func IsBinary(header http.Header) bool { + encoding := strings.ToLower(strings.TrimSpace(header.Get("Content-Encoding"))) + switch encoding { + case "", "identity", "gzip", "deflate", "br", "zstd": + // Known encodings we can decompress — fall through to Content-Type check. + default: + // Unknown encoding — we can't decode it, treat as binary. + return true + } + + ct := strings.ToLower(header.Get("Content-Type")) + if ct == "" { + // No Content-Type — assume text so we don't silently discard bodies. + return false + } + + textPrefixes := []string{ + "text/", + "application/json", + "application/xml", + "application/javascript", + "application/x-www-form-urlencoded", + "application/xhtml", + "application/ld+json", + "application/graphql", + "application/x-yaml", + "application/yaml", + } + for _, prefix := range textPrefixes { + if strings.Contains(ct, prefix) { + return false + } + } + return true +} + +// TruncateBody returns body truncated to MaxDisplayBytes with a notice appended. +func TruncateBody(body []byte) []byte { + if len(body) <= MaxDisplayBytes { + return body + } + notice := []byte("\n... truncated") + out := make([]byte, MaxDisplayBytes+len(notice)) + copy(out, body[:MaxDisplayBytes]) + copy(out[MaxDisplayBytes:], notice) + return out +} diff --git a/internal/http/compression_test.go b/internal/http/compression_test.go new file mode 100644 index 0000000..b675341 --- /dev/null +++ b/internal/http/compression_test.go @@ -0,0 +1,152 @@ +package http + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "net/http" + "strings" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── Decompress ──────────────────────────────────────────────────────────────── + +func TestDecompress_NoEncoding(t *testing.T) { + body := []byte("plain") + assert.Equal(t, body, Decompress(http.Header{}, body)) +} + +func TestDecompress_UnknownEncoding(t *testing.T) { + body := []byte("mystery") + assert.Equal(t, body, Decompress(http.Header{"Content-Encoding": []string{"identity"}}, body)) +} + +func TestDecompress_Gzip(t *testing.T) { + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write([]byte("gzip content")) + require.NoError(t, err) + w.Close() + assert.Equal(t, "gzip content", string(Decompress(http.Header{"Content-Encoding": []string{"gzip"}}, buf.Bytes()))) +} + +func TestDecompress_Gzip_Corrupt(t *testing.T) { + body := []byte("not gzip") + assert.Equal(t, body, Decompress(http.Header{"Content-Encoding": []string{"gzip"}}, body)) +} + +func TestDecompress_Deflate(t *testing.T) { + var buf bytes.Buffer + w, err := flate.NewWriter(&buf, flate.DefaultCompression) + require.NoError(t, err) + w.Write([]byte("deflate content")) + w.Close() + assert.Equal(t, "deflate content", string(Decompress(http.Header{"Content-Encoding": []string{"deflate"}}, buf.Bytes()))) +} + +func TestDecompress_Deflate_Corrupt(t *testing.T) { + body := []byte("not deflate") + assert.Equal(t, body, Decompress(http.Header{"Content-Encoding": []string{"deflate"}}, body)) +} + +func TestDecompress_Brotli(t *testing.T) { + var buf bytes.Buffer + w := brotli.NewWriter(&buf) + w.Write([]byte("brotli content")) + w.Close() + assert.Equal(t, "brotli content", string(Decompress(http.Header{"Content-Encoding": []string{"br"}}, buf.Bytes()))) +} + +func TestDecompress_Brotli_Corrupt(t *testing.T) { + body := []byte("not brotli") + assert.Equal(t, body, Decompress(http.Header{"Content-Encoding": []string{"br"}}, body)) +} + +func TestDecompress_Zstd(t *testing.T) { + enc, err := zstd.NewWriter(nil) + require.NoError(t, err) + compressed := enc.EncodeAll([]byte("zstd content"), nil) + assert.Equal(t, "zstd content", string(Decompress(http.Header{"Content-Encoding": []string{"zstd"}}, compressed))) +} + +func TestDecompress_Zstd_Corrupt(t *testing.T) { + body := []byte("not zstd") + assert.Equal(t, body, Decompress(http.Header{"Content-Encoding": []string{"zstd"}}, body)) +} + +func TestDecompress_CaseInsensitive(t *testing.T) { + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + w.Write([]byte("case")) + w.Close() + assert.Equal(t, "case", string(Decompress(http.Header{"Content-Encoding": []string{"GZIP"}}, buf.Bytes()))) +} + +func TestDecompress_EmptyBody(t *testing.T) { + assert.Equal(t, []byte{}, Decompress(http.Header{"Content-Encoding": []string{"gzip"}}, []byte{})) +} + +func TestDecompress_LargeBody(t *testing.T) { + large := []byte(strings.Repeat("a", 100000)) + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + w.Write(large) + w.Close() + out := Decompress(http.Header{"Content-Encoding": []string{"gzip"}}, buf.Bytes()) + assert.Equal(t, large, out) +} + +// ── IsBinary ────────────────────────────────────────────────────────────────── + +func TestIsBinary_KnownEncodingsNotBinary(t *testing.T) { + // Known encodings we can decompress fall through to the Content-Type check. + // image/png is still binary even with a known encoding — the encoding tells + // us we can decompress it, but the content type determines binary/text. + for _, enc := range []string{"gzip", "deflate", "br", "zstd"} { + h := http.Header{ + "Content-Encoding": []string{enc}, + "Content-Type": []string{"text/plain"}, + } + assert.False(t, IsBinary(h), "encoding %s with text/plain should not be binary", enc) + } + // image/png with a known encoding is still binary. + for _, enc := range []string{"gzip", "deflate", "br", "zstd"} { + h := http.Header{ + "Content-Encoding": []string{enc}, + "Content-Type": []string{"image/png"}, + } + assert.True(t, IsBinary(h), "encoding %s with image/png should be binary", enc) + } +} + +func TestIsBinary_BinaryContentTypes(t *testing.T) { + for _, ct := range []string{"image/png", "image/jpeg", "application/octet-stream", "video/mp4", "audio/mpeg"} { + h := http.Header{"Content-Type": []string{ct}} + assert.True(t, IsBinary(h), "%s should be binary", ct) + } +} + +func TestIsBinary_TextContentTypes(t *testing.T) { + for _, ct := range []string{ + "text/plain", "text/html", "text/css", + "application/json", "application/xml", + "application/javascript", "application/x-www-form-urlencoded", + "application/xhtml+xml", "application/ld+json", + } { + h := http.Header{"Content-Type": []string{ct}} + assert.False(t, IsBinary(h), "%s should not be binary", ct) + } +} + +func TestIsBinary_EmptyContentType(t *testing.T) { + assert.False(t, IsBinary(http.Header{"Content-Type": []string{""}})) +} + +func TestIsBinary_NoHeaders(t *testing.T) { + assert.False(t, IsBinary(http.Header{})) +} diff --git a/internal/http/headers.go b/internal/http/headers.go new file mode 100644 index 0000000..d06d584 --- /dev/null +++ b/internal/http/headers.go @@ -0,0 +1,54 @@ +package http + +import ( + "net" + "net/http" +) + +// NormalizeHost strips the port from a Host header value when it is the +// default for the scheme (443 for HTTPS, 80 for HTTP), so that we never +// send "Host: example.com:443" to an upstream server. +func NormalizeHost(host string, tls bool) string { + h, port, err := net.SplitHostPort(host) + if err != nil { + return host + } + if (tls && port == "443") || (!tls && port == "80") { + return h + } + return host +} + +// StripRequestCacheHeaders removes conditional-request headers so the +// upstream always returns a fresh response. +func StripRequestCacheHeaders(h http.Header) { + h.Del("If-None-Match") + h.Del("If-Modified-Since") + h.Del("If-Range") + h.Del("If-Match") + h.Del("If-Unmodified-Since") +} + +// StripResponseCacheHeaders removes caching directives from a response and +// sets no-store so the browser never serves a stale cached copy. +func StripResponseCacheHeaders(h http.Header) { + h.Del("Cache-Control") + h.Del("Expires") + h.Del("ETag") + h.Del("Last-Modified") + h.Del("Pragma") + h.Set("Cache-Control", "no-store, no-cache, must-revalidate") +} + +// HopByHop is the set of headers that must not be forwarded to the upstream. +var HopByHop = map[string]bool{ + "Connection": true, + "Proxy-Connection": true, + "Keep-Alive": true, + "Proxy-Authenticate": true, + "Proxy-Authorization": true, + "Te": true, + "Trailers": true, + "Transfer-Encoding": true, + "Upgrade": true, +} diff --git a/internal/http/headers_test.go b/internal/http/headers_test.go new file mode 100644 index 0000000..2378da6 --- /dev/null +++ b/internal/http/headers_test.go @@ -0,0 +1,117 @@ +package http + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" +) + +// ── NormalizeHost ───────────────────────────────────────────────────────────── + +func TestNormalizeHost_TLS_443_Stripped(t *testing.T) { + assert.Equal(t, "example.com", NormalizeHost("example.com:443", true)) +} + +func TestNormalizeHost_Plain_80_Stripped(t *testing.T) { + assert.Equal(t, "example.com", NormalizeHost("example.com:80", false)) +} + +func TestNormalizeHost_TLS_NonDefault_Kept(t *testing.T) { + assert.Equal(t, "example.com:8443", NormalizeHost("example.com:8443", true)) +} + +func TestNormalizeHost_Plain_NonDefault_Kept(t *testing.T) { + assert.Equal(t, "example.com:8080", NormalizeHost("example.com:8080", false)) +} + +func TestNormalizeHost_Plain_443_Kept(t *testing.T) { + // port 443 on plain (non-TLS) should not be stripped + assert.Equal(t, "example.com:443", NormalizeHost("example.com:443", false)) +} + +func TestNormalizeHost_TLS_80_Kept(t *testing.T) { + // port 80 on TLS should not be stripped + assert.Equal(t, "example.com:80", NormalizeHost("example.com:80", true)) +} + +func TestNormalizeHost_NoPort_ReturnedAsIs(t *testing.T) { + assert.Equal(t, "example.com", NormalizeHost("example.com", true)) +} + +func TestNormalizeHost_IP_WithPort443(t *testing.T) { + assert.Equal(t, "127.0.0.1", NormalizeHost("127.0.0.1:443", true)) +} + +// ── StripRequestCacheHeaders ────────────────────────────────────────────────── + +func TestStripRequestCacheHeaders_AllStripped(t *testing.T) { + h := http.Header{ + "If-None-Match": []string{`"etag"`}, + "If-Modified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "If-Range": []string{`"etag"`}, + "If-Match": []string{`"etag"`}, + "If-Unmodified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Authorization": []string{"Bearer token"}, + } + StripRequestCacheHeaders(h) + + assert.Empty(t, h.Get("If-None-Match")) + assert.Empty(t, h.Get("If-Modified-Since")) + assert.Empty(t, h.Get("If-Range")) + assert.Empty(t, h.Get("If-Match")) + assert.Empty(t, h.Get("If-Unmodified-Since")) + assert.Equal(t, "Bearer token", h.Get("Authorization")) +} + +func TestStripRequestCacheHeaders_AlreadyAbsent(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/json"}} + assert.NotPanics(t, func() { StripRequestCacheHeaders(h) }) + assert.Equal(t, "application/json", h.Get("Content-Type")) +} + +// ── StripResponseCacheHeaders ───────────────────────────────────────────────── + +func TestStripResponseCacheHeaders_AllStripped(t *testing.T) { + h := http.Header{ + "Cache-Control": []string{"max-age=3600"}, + "Expires": []string{"Mon, 01 Jan 2025 00:00:00 GMT"}, + "ETag": []string{`"etag"`}, + "Last-Modified": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Pragma": []string{"no-cache"}, + "Content-Type": []string{"text/html"}, + } + StripResponseCacheHeaders(h) + + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) + assert.Empty(t, h.Get("Expires")) + assert.Empty(t, h.Get("ETag")) + assert.Empty(t, h.Get("Last-Modified")) + assert.Empty(t, h.Get("Pragma")) + assert.Equal(t, "text/html", h.Get("Content-Type")) +} + +func TestStripResponseCacheHeaders_SetEvenWhenAbsent(t *testing.T) { + h := http.Header{} + StripResponseCacheHeaders(h) + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) +} + +// ── HopByHop ────────────────────────────────────────────────────────────────── + +func TestHopByHop_ContainsExpectedHeaders(t *testing.T) { + expected := []string{ + "Connection", "Proxy-Connection", "Keep-Alive", + "Proxy-Authenticate", "Proxy-Authorization", + "Te", "Trailers", "Transfer-Encoding", "Upgrade", + } + for _, h := range expected { + assert.True(t, HopByHop[h], "expected %s in HopByHop", h) + } +} + +func TestHopByHop_DoesNotContainRegularHeaders(t *testing.T) { + assert.False(t, HopByHop["Content-Type"]) + assert.False(t, HopByHop["Authorization"]) + assert.False(t, HopByHop["Host"]) +} diff --git a/internal/http/rawrequest.go b/internal/http/rawrequest.go new file mode 100644 index 0000000..090787a --- /dev/null +++ b/internal/http/rawrequest.go @@ -0,0 +1,327 @@ +package http + +import ( + "bufio" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "strconv" + "strings" + "time" +) + +// RawRequestOptions controls how SendRaw dials and sends. +type RawRequestOptions struct { + Host string + Port int + TLS bool + RawReq string + CookieJar map[string]string // optional; merged into Cookie header + TimeoutMs int // 0 means use defaults (10s dial, 30s total) +} + +// RawResponse is returned by SendRaw. +type RawResponse struct { + Raw string // full response as text (status line + headers + body) + StatusCode int + Headers http.Header + Body []byte + Cookies []*http.Cookie +} + +// SendRaw dials host:port, sends rawReq over a plain or TLS connection, and +// returns the decompressed, decoded response. It is intentionally low-level +// so the caller (Repeater, Intruder) controls the exact bytes on the wire. +// +// The request line is always rewritten to HTTP/1.1 — H2 frames cannot be +// sent as raw text, and a raw socket is always HTTP/1.1. +func SendRaw(opts RawRequestOptions) (*RawResponse, error) { + addr := net.JoinHostPort(opts.Host, strconv.Itoa(opts.Port)) + + timeout := 30 * time.Second + if opts.TimeoutMs > 0 { + timeout = time.Duration(opts.TimeoutMs) * time.Millisecond + } + + var conn net.Conn + var err error + if opts.TLS { + dialer := &net.Dialer{Timeout: timeout} + tlsCfg := &tls.Config{ + ServerName: opts.Host, + InsecureSkipVerify: true, //nolint:gosec + } + conn, err = tls.DialWithDialer(dialer, "tcp", addr, tlsCfg) + } else { + conn, err = net.DialTimeout("tcp", addr, timeout) + } + if err != nil { + return nil, fmt.Errorf("dial %s: %w", addr, err) + } + defer conn.Close() + + if err := conn.SetDeadline(time.Now().Add(timeout)); err != nil { + return nil, fmt.Errorf("set deadline: %w", err) + } + + raw := NormalizeLineEndings(opts.RawReq) + raw = rewriteRequestLine(raw) + raw = rewriteHeaders(raw, opts.CookieJar) + + if _, err := fmt.Fprint(conn, raw); err != nil { + return nil, fmt.Errorf("write request: %w", err) + } + + httpResp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + defer httpResp.Body.Close() + + body, err := io.ReadAll(io.LimitReader(httpResp.Body, 2<<20)) + if err != nil { + return nil, fmt.Errorf("read response body: %w", err) + } + body = Decompress(httpResp.Header, body) + body = TruncateBody(body) + rawResp := buildRawResponse(httpResp, body) + + return &RawResponse{ + Raw: rawResp, + StatusCode: httpResp.StatusCode, + Headers: httpResp.Header, + Body: body, + Cookies: httpResp.Cookies(), + }, nil +} + +// rewriteRequestLine rewrites the first line of a raw HTTP request to use +// HTTP/1.1, regardless of what was captured. HTTP/2 cannot be sent as raw +// text over a socket. +func rewriteRequestLine(raw string) string { + idx := strings.Index(raw, "\r\n") + if idx < 0 { + return raw + } + firstLine := raw[:idx] + parts := strings.Fields(firstLine) + if len(parts) < 2 { + return raw + } + // Rewrite: METHOD PATH HTTP/1.1 + newFirstLine := parts[0] + " " + parts[1] + " HTTP/1.1" + return newFirstLine + raw[idx:] +} + +// ParseHostFromRaw extracts host, port, and TLS flag from the Host header +// in a raw HTTP request string. The caller should pass the known TLS state +// from the captured transaction when available, as the Host header alone +// cannot distinguish HTTP from HTTPS. +func ParseHostFromRaw(raw string) (host string, port int, useTLS bool) { + // Check the scheme in the request line first (e.g. absolute-form URLs). + firstLine := strings.SplitN(raw, "\n", 2)[0] + parts := strings.Fields(firstLine) + if len(parts) >= 2 { + url := parts[1] + if strings.HasPrefix(url, "https://") { + useTLS = true + } else if strings.HasPrefix(url, "http://") { + useTLS = false + } + } + + for line := range strings.SplitSeq(raw, "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(strings.ToLower(line), "host:") { + continue + } + hostVal := strings.TrimSpace(line[5:]) + if h, portStr, err := net.SplitHostPort(hostVal); err == nil { + host = h + port, _ = strconv.Atoi(portStr) + // Port in Host header is the definitive signal. + useTLS = port == 443 + } else { + host = hostVal + if net.ParseIP(hostVal) != nil { + // Bare IP without port — can't tell, default to plaintext. + port = 80 + useTLS = false + } else { + // Domain without port — use whatever we detected from the + // request line scheme, or fall back to HTTPS. + if port == 0 { + if useTLS { + port = 443 + } else { + port = 80 + } + } + // If we got here without a scheme signal, default to HTTPS + // since most modern interception is TLS. + if port == 443 { + useTLS = true + } + } + } + return + } + return "", 0, false +} + +// ParseRawHeaders parses the headers section of a raw HTTP response string +// into an http.Header map. +func ParseRawHeaders(raw string) http.Header { + headers := http.Header{} + lines := strings.Split(raw, "\r\n") + if len(lines) == 0 { + return headers + } + // Skip the status/request line. + for _, line := range lines[1:] { + if line == "" { + break + } + parts := strings.SplitN(line, ": ", 2) + if len(parts) == 2 { + headers.Add(parts[0], parts[1]) + } + } + return headers +} + +// SplitRaw splits a raw HTTP message into its header section and body at the +// first blank line. Handles both CRLF and LF delimiters. The returned header +// section does not include the blank-line separator. +func SplitRaw(raw string) (header string, body []byte) { + if idx := strings.Index(raw, "\r\n\r\n"); idx >= 0 { + return raw[:idx], []byte(raw[idx+4:]) + } + if idx := strings.Index(raw, "\n\n"); idx >= 0 { + return raw[:idx], []byte(raw[idx+2:]) + } + return raw, nil +} + +// NormalizeLineEndings ensures CRLF throughout the header section only. +// It does NOT touch the body to avoid corrupting binary or multi-byte content. +func NormalizeLineEndings(raw string) string { + header, body := SplitRaw(raw) + header = strings.ReplaceAll(header, "\r\n", "\n") + header = strings.ReplaceAll(header, "\n", "\r\n") + return header + "\r\n\r\n" + string(body) +} + +// rewriteHeaders rewrites the header section of a raw request: +// - Injects Connection: close (required for single-use raw TCP connections) +// - Strips Accept-Encoding (we handle decompression ourselves) +// - Strips Content-Length (recalculated below) +// - Strips Transfer-Encoding (not valid in HTTP/1.1 raw send context) +// - Merges Cookie header from jar (replaces existing if jar is non-empty) +// - Strips default port from Host when it matches the scheme +// - Appends correct Content-Length based on byte length of body +func rewriteHeaders(raw string, jar map[string]string) string { + // Split on the canonical blank line separator (normalizeLineEndings runs first). + parts := strings.SplitN(raw, "\r\n\r\n", 2) + headerSection := parts[0] + body := "" + if len(parts) == 2 { + body = parts[1] + // Trim trailing CRLF from body only if it was added by normalisation, + // not if the user intentionally added trailing newlines. + // We use the raw byte length for Content-Length regardless. + } + + lines := strings.Split(headerSection, "\r\n") + out := make([]string, 0, len(lines)+3) + connectionSeen := false + + for _, line := range lines { + lower := strings.ToLower(line) + switch { + case strings.HasPrefix(lower, "cookie:"): + // Drop original Cookie header only if jar has entries to replace it. + if len(jar) == 0 { + out = append(out, line) + } + case strings.HasPrefix(lower, "accept-encoding:"): + // Drop — we decompress ourselves so we want uncompressed bytes. + continue + case strings.HasPrefix(lower, "content-length:"): + // Drop — we recalculate below from actual body byte length. + continue + case strings.HasPrefix(lower, "transfer-encoding:"): + // Not valid to send in this raw H1 context. + continue + case strings.HasPrefix(lower, "connection:"): + // Replace whatever was there with Connection: close. + out = append(out, "Connection: close") + connectionSeen = true + case strings.HasPrefix(lower, "host:"): + hostVal := strings.TrimSpace(line[5:]) + // Strip default ports from Host header. + if h, port, err := net.SplitHostPort(hostVal); err == nil { + if port == "443" || port == "80" { + line = "Host: " + h + } + } + out = append(out, line) + default: + out = append(out, line) + } + } + + // Ensure Connection: close is present even if the original had no Connection header. + if !connectionSeen { + out = append(out, "Connection: close") + } + + // Append merged Cookie header from jar. + if len(jar) > 0 { + cookies := make([]string, 0, len(jar)) + for k, v := range jar { + cookies = append(cookies, k+"="+v) + } + out = append(out, "Cookie: "+strings.Join(cookies, "; ")) + } + + // Append Content-Length based on actual byte length of the body. + if len(body) > 0 { + out = append(out, fmt.Sprintf("Content-Length: %d", len(body))) + } + + return strings.Join(out, "\r\n") + "\r\n\r\n" + body +} + +func buildRawResponse(resp *http.Response, body []byte) string { + var b strings.Builder + fmt.Fprintf(&b, "HTTP/1.1 %d %s\r\n", resp.StatusCode, http.StatusText(resp.StatusCode)) + for k, vals := range resp.Header { + for _, v := range vals { + fmt.Fprintf(&b, "%s: %s\r\n", k, v) + } + } + b.WriteString("\r\n") + b.Write(body) + return b.String() +} + +// ExtractURL extracts the path from the first line of a raw HTTP request. +func ExtractURL(rawRequest string) string { + parts := strings.Fields(strings.SplitN(rawRequest, "\n", 2)[0]) + if len(parts) >= 2 { + return parts[1] + } + return "/" +} + +// ExtractMethod extracts the HTTP method from the first line of a raw HTTP request. +func ExtractMethod(rawRequest string) string { + parts := strings.Fields(strings.SplitN(rawRequest, "\n", 2)[0]) + if len(parts) >= 1 { + return parts[0] + } + return "GET" +} diff --git a/internal/http/rawrequest_test.go b/internal/http/rawrequest_test.go new file mode 100644 index 0000000..5dd411c --- /dev/null +++ b/internal/http/rawrequest_test.go @@ -0,0 +1,254 @@ +package http + +import ( + "fmt" + "net" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── ParseHostFromRaw ────────────────────────────────────────────────────────── + +func TestParseHostFromRaw_HostOnly(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n") + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_HostWithPort443(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nHost: example.com:443\r\n\r\n") + assert.Equal(t, "example.com", host) + assert.Equal(t, 443, port) + assert.True(t, useTLS) +} + +func TestParseHostFromRaw_HostWithPort80(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nHost: example.com:80\r\n\r\n") + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_CustomPort(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nHost: example.com:8080\r\n\r\n") + assert.Equal(t, "example.com", host) + assert.Equal(t, 8080, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_IP(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nHost: 1.2.3.4\r\n\r\n") + assert.Equal(t, "1.2.3.4", host) + assert.Equal(t, 80, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_NoHostHeader(t *testing.T) { + host, port, useTLS := ParseHostFromRaw("GET / HTTP/1.1\r\nContent-Type: text/plain\r\n\r\n") + assert.Empty(t, host) + assert.Equal(t, 0, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_CaseInsensitive(t *testing.T) { + host, _, _ := ParseHostFromRaw("GET / HTTP/1.1\r\nHOST: example.com\r\n\r\n") + assert.Equal(t, "example.com", host) +} + +func TestParseHostFromRaw_LFOnly(t *testing.T) { + host, port, _ := ParseHostFromRaw("GET / HTTP/1.1\nHost: example.com\n\n") + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) +} + +// ── ParseRawHeaders ─────────────────────────────────────────────────────────── + +func TestParseRawHeaders_Single(t *testing.T) { + raw := "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n" + h := ParseRawHeaders(raw) + assert.Equal(t, "text/html", h.Get("Content-Type")) +} + +func TestParseRawHeaders_Multiple(t *testing.T) { + raw := "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nX-Custom: value\r\n\r\n" + h := ParseRawHeaders(raw) + assert.Equal(t, "text/html", h.Get("Content-Type")) + assert.Equal(t, "value", h.Get("X-Custom")) +} + +func TestParseRawHeaders_StopsAtBlankLine(t *testing.T) { + raw := "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\nbody here" + h := ParseRawHeaders(raw) + assert.Equal(t, "text/html", h.Get("Content-Type")) + assert.Empty(t, h.Get("body here")) +} + +func TestParseRawHeaders_Empty(t *testing.T) { + h := ParseRawHeaders("HTTP/1.1 200 OK\r\n\r\n") + assert.Equal(t, http.Header{}, h) +} + +// ── normalizeLineEndings ────────────────────────────────────────────────────── + +func TestNormalizeLineEndings_LFtoCRLF(t *testing.T) { + out := NormalizeLineEndings("GET / HTTP/1.1\nHost: example.com\n\n") + assert.Contains(t, out, "\r\n") + assert.NotContains(t, out, "\n\n") +} + +func TestNormalizeLineEndings_AlreadyCRLF(t *testing.T) { + in := "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n" + out := NormalizeLineEndings(in) + assert.Equal(t, in, out) +} + +func TestNormalizeLineEndings_FoldsContinuation(t *testing.T) { + // Continuation line folding was removed to avoid corrupting body content. + // normalizeLineEndings only touches the header section and converts LF to CRLF. + // Continuation lines are preserved as-is. + in := "GET / HTTP/1.1\r\nHost: example.com\r\n continued\r\n\r\n" + out := NormalizeLineEndings(in) + assert.Equal(t, in, out) +} + +// ── rewriteHeaders ──────────────────────────────────────────────────────────── + +func TestRewriteHeaders_StripsAcceptEncoding(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com\r\nAccept-Encoding: gzip\r\n\r\n" + out := rewriteHeaders(raw, nil) + assert.NotContains(t, out, "Accept-Encoding") +} + +func TestRewriteHeaders_StripsContentLength(t *testing.T) { + raw := "POST / HTTP/1.1\r\nHost: example.com\r\nContent-Length: 999\r\n\r\nbody" + out := rewriteHeaders(raw, nil) + assert.NotContains(t, out, "Content-Length: 999") + assert.Contains(t, out, fmt.Sprintf("Content-Length: %d", len("body"))) +} + +func TestRewriteHeaders_PreservesCookieWhenNoJar(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com\r\nCookie: session=abc\r\n\r\n" + out := rewriteHeaders(raw, nil) + assert.Contains(t, out, "Cookie: session=abc") +} + +func TestRewriteHeaders_ReplacesCookieFromJar(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com\r\nCookie: session=old\r\n\r\n" + out := rewriteHeaders(raw, map[string]string{"session": "new"}) + assert.NotContains(t, out, "session=old") + assert.Contains(t, out, "session=new") +} + +func TestRewriteHeaders_StripsDefaultPort443(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:443\r\n\r\n" + out := rewriteHeaders(raw, nil) + assert.Contains(t, out, "Host: example.com\r\n") + assert.NotContains(t, out, "Host: example.com:443") +} + +func TestRewriteHeaders_StripsDefaultPort80(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:80\r\n\r\n" + out := rewriteHeaders(raw, nil) + assert.Contains(t, out, "Host: example.com\r\n") + assert.NotContains(t, out, "Host: example.com:80") +} + +func TestRewriteHeaders_KeepsNonDefaultPort(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:8080\r\n\r\n" + out := rewriteHeaders(raw, nil) + assert.Contains(t, out, "Host: example.com:8080") +} + +// ── SendRaw (mock TCP server) ───────────────────────────────────────────────── + +func startMockServer(t *testing.T, response string) (string, int) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + buf := make([]byte, 4096) + conn.Read(buf) + fmt.Fprint(conn, response) + }() + t.Cleanup(func() { ln.Close() }) + addr := ln.Addr().(*net.TCPAddr) + return "127.0.0.1", addr.Port +} + +func TestSendRaw_200OK(t *testing.T) { + host, port := startMockServer(t, "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello") + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host, port) + result, err := SendRaw(RawRequestOptions{Host: host, Port: port, RawReq: raw}) + require.NoError(t, err) + assert.Equal(t, 200, result.StatusCode) + assert.Equal(t, "hello", string(result.Body)) +} + +func TestSendRaw_404(t *testing.T) { + host, port := startMockServer(t, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n") + raw := fmt.Sprintf("GET /missing HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host, port) + result, err := SendRaw(RawRequestOptions{Host: host, Port: port, RawReq: raw}) + require.NoError(t, err) + assert.Equal(t, 404, result.StatusCode) +} + +func TestSendRaw_ConnectionRefused(t *testing.T) { + _, err := SendRaw(RawRequestOptions{ + Host: "127.0.0.1", Port: 1, + RawReq: "GET / HTTP/1.1\r\nHost: 127.0.0.1:1\r\n\r\n", + }) + assert.Error(t, err) +} + +func TestSendRaw_CookieJarApplied(t *testing.T) { + received := make(chan string, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { ln.Close() }) + go func() { + conn, _ := ln.Accept() + defer conn.Close() + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + received <- string(buf[:n]) + fmt.Fprint(conn, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + }() + addr := ln.Addr().(*net.TCPAddr) + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: 127.0.0.1:%d\r\n\r\n", addr.Port) + SendRaw(RawRequestOptions{ + Host: "127.0.0.1", Port: addr.Port, + RawReq: raw, + CookieJar: map[string]string{"session": "abc123"}, + }) + req := <-received + assert.Contains(t, req, "Cookie: session=abc123") +} + +func TestSendRaw_AcceptEncodingStripped(t *testing.T) { + received := make(chan string, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { ln.Close() }) + go func() { + conn, _ := ln.Accept() + defer conn.Close() + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + received <- string(buf[:n]) + fmt.Fprint(conn, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + }() + addr := ln.Addr().(*net.TCPAddr) + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: 127.0.0.1:%d\r\nAccept-Encoding: gzip\r\n\r\n", addr.Port) + SendRaw(RawRequestOptions{Host: "127.0.0.1", Port: addr.Port, RawReq: raw}) + req := <-received + assert.NotContains(t, req, "Accept-Encoding") +} diff --git a/internal/logger/logger.go b/internal/logger/logger.go new file mode 100644 index 0000000..e0f6598 --- /dev/null +++ b/internal/logger/logger.go @@ -0,0 +1,49 @@ +// Package logger provides a minimal levelled logger for Shiv. +// Output is gated behind a verbose flag set at startup via Init. +// Only fatal startup messages use Always — everything else is verbose. +package logger + +import ( + "fmt" + "io" + "os" + "time" +) + +var verbose bool + +// Init configures the logger. Call once at startup. +func Init(isVerbose bool) { + verbose = isVerbose +} + +// Always logs a message that is always shown regardless of verbose flag. +// Use sparingly — only for critical startup/shutdown messages. +func Always(format string, args ...any) { + log(os.Stdout, "INFO", format, args...) +} + +// Info logs a message only when verbose mode is enabled. +func Info(format string, args ...any) { + if verbose { + log(os.Stdout, "INFO", format, args...) + } +} + +// Error logs an error that is always shown. +func Error(format string, args ...any) { + log(os.Stderr, "ERR ", format, args...) +} + +// Debug logs a message only when verbose mode is enabled. +func Debug(format string, args ...any) { + if verbose { + log(os.Stdout, "DBG ", format, args...) + } +} + +func log(w io.Writer, level, format string, args ...any) { + timestamp := time.Now().Format("15:04:05.000") + msg := fmt.Sprintf(format, args...) + fmt.Fprintf(w, "[shiv] %s %s %s\n", timestamp, level, msg) +} diff --git a/internal/plugin/api.go b/internal/plugin/api.go new file mode 100644 index 0000000..986c5d2 --- /dev/null +++ b/internal/plugin/api.go @@ -0,0 +1,313 @@ +package plugin + +import ( + "io" + "net/http" + "strings" + "time" + + lua "github.com/yuin/gopher-lua" + + "github.com/shiv/internal/store" +) + +// registerAPI registers all Lua API namespaces on the given VM. +// logFn is called when the Lua script calls log() — the engine routes +// these lines onto the bus after each hook call. +// Called once per plugin at load time. +func registerAPI(L *lua.LState, st *store.Store, logFn func(string)) { + registerLog(L, logFn) + registerHTTPClient(L) + registerDB(L, st) +} + +// ── log ─────────────────────────────────────────────────────────────────────── + +func registerLog(L *lua.LState, logFn func(string)) { + L.SetGlobal("log", L.NewFunction(func(L *lua.LState) int { + msg := L.CheckString(1) + logFn(msg) + return 0 + })) +} + +// ── http client ─────────────────────────────────────────────────────────────── + +var httpClient = &http.Client{Timeout: 15 * time.Second} + +func registerHTTPClient(L *lua.LState) { + tbl := L.NewTable() + + // http.get(url) → {status, headers, body} or nil, err + L.SetField(tbl, "get", L.NewFunction(func(L *lua.LState) int { + url := L.CheckString(1) + resp, err := httpClient.Get(url) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + defer resp.Body.Close() + return pushResponse(L, resp) + })) + + // http.post(url, body, content_type) → {status, headers, body} or nil, err + L.SetField(tbl, "post", L.NewFunction(func(L *lua.LState) int { + url := L.CheckString(1) + body := L.CheckString(2) + ct := L.OptString(3, "application/json") + resp, err := httpClient.Post(url, ct, strings.NewReader(body)) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + defer resp.Body.Close() + return pushResponse(L, resp) + })) + + // http.request(method, url, body, headers_table) → {status, headers, body} or nil, err + L.SetField(tbl, "request", L.NewFunction(func(L *lua.LState) int { + method := L.CheckString(1) + url := L.CheckString(2) + body := L.OptString(3, "") + headersTbl := L.OptTable(4, nil) + + var bodyReader io.Reader + if body != "" { + bodyReader = strings.NewReader(body) + } + + req, err := http.NewRequest(method, url, bodyReader) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + + if headersTbl != nil { + headersTbl.ForEach(func(k, v lua.LValue) { + req.Header.Set(k.String(), v.String()) + }) + } + + resp, err := httpClient.Do(req) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + defer resp.Body.Close() + return pushResponse(L, resp) + })) + + L.SetGlobal("http", tbl) +} + +func pushResponse(L *lua.LState, resp *http.Response) int { + body, _ := io.ReadAll(resp.Body) + + tbl := L.NewTable() + L.SetField(tbl, "status", lua.LNumber(resp.StatusCode)) + L.SetField(tbl, "body", lua.LString(body)) + + headers := L.NewTable() + for k, vals := range resp.Header { + L.SetField(headers, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "headers", headers) + + L.Push(tbl) + return 1 +} + +// ── db ──────────────────────────────────────────────────────────────────────── + +func registerDB(L *lua.LState, st *store.Store) { + tbl := L.NewTable() + + // db.history(limit) → array of transaction tables + L.SetField(tbl, "history", L.NewFunction(func(L *lua.LState) int { + limit := L.OptInt(1, 100) + txs, err := st.TransactionsPage(0, store.TransactionFilter{ShowOutScope: true}) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + if limit > 0 && len(txs) > limit { + txs = txs[:limit] + } + arr := L.NewTable() + for i, tx := range txs { + L.RawSetInt(arr, i+1, transactionToTable(L, tx)) + } + L.Push(arr) + return 1 + })) + + // db.history_filter(host, method, status) → array of transaction tables + L.SetField(tbl, "history_filter", L.NewFunction(func(L *lua.LState) int { + host := L.OptString(1, "") + method := L.OptString(2, "") + _ = L.OptInt(3, 0) // status filtering done client-side below + + filter := store.TransactionFilter{ + Host: host, + ShowOutScope: true, + } + if method != "" { + filter.Search = method + } + + txs, err := st.TransactionsPage(0, filter) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + arr := L.NewTable() + idx := 1 + for _, tx := range txs { + L.RawSetInt(arr, idx, transactionToTable(L, tx)) + idx++ + } + L.Push(arr) + return 1 + })) + + // db.scope_get() → array of {id, host} + L.SetField(tbl, "scope_get", L.NewFunction(func(L *lua.LState) int { + entries, err := st.AllScopeEntries() + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + arr := L.NewTable() + for i, e := range entries { + row := L.NewTable() + L.SetField(row, "id", lua.LNumber(e.ID)) + L.SetField(row, "host", lua.LString(e.Host)) + L.RawSetInt(arr, i+1, row) + } + L.Push(arr) + return 1 + })) + + // db.scope_add(host) → true or nil, err + L.SetField(tbl, "scope_add", L.NewFunction(func(L *lua.LState) int { + host := L.CheckString(1) + if err := st.AddScopeEntry(host); err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + L.Push(lua.LTrue) + return 1 + })) + + // db.scope_remove(id) → true or nil, err + L.SetField(tbl, "scope_remove", L.NewFunction(func(L *lua.LState) int { + id := L.CheckInt64(1) + if err := st.DeleteScopeEntry(id); err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + L.Push(lua.LTrue) + return 1 + })) + + // db.loot_get() → array of loot tables + L.SetField(tbl, "loot_get", L.NewFunction(func(L *lua.LState) int { + entries, err := st.AllLoot() + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + arr := L.NewTable() + for i, e := range entries { + row := L.NewTable() + L.SetField(row, "id", lua.LNumber(e.ID)) + L.SetField(row, "title", lua.LString(e.Title)) + L.SetField(row, "severity", lua.LString(e.Severity)) + L.SetField(row, "notes", lua.LString(e.Notes)) + L.SetField(row, "raw_request", lua.LString(e.RawRequest)) + L.SetField(row, "raw_response", lua.LString(e.RawResponse)) + L.RawSetInt(arr, i+1, row) + } + L.Push(arr) + return 1 + })) + + // db.loot_add(title, severity, notes, raw_request, raw_response) → id or nil, err + L.SetField(tbl, "loot_add", L.NewFunction(func(L *lua.LState) int { + title := L.CheckString(1) + severity := L.CheckString(2) + notes := L.OptString(3, "") + rawReq := L.OptString(4, "") + rawResp := L.OptString(5, "") + id, err := st.AddLoot(store.LootEntry{ + Title: title, + Severity: severity, + Notes: notes, + RawRequest: rawReq, + RawResponse: rawResp, + }) + if err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + L.Push(lua.LNumber(id)) + return 1 + })) + + // db.loot_delete(id) → true or nil, err + L.SetField(tbl, "loot_delete", L.NewFunction(func(L *lua.LState) int { + id := L.CheckInt64(1) + if err := st.DeleteLoot(id); err != nil { + L.Push(lua.LNil) + L.Push(lua.LString(err.Error())) + return 2 + } + L.Push(lua.LTrue) + return 1 + })) + + L.SetGlobal("db", tbl) +} + +// ── helpers ─────────────────────────────────────────────────────────────────── + +func transactionToTable(L *lua.LState, tx store.Transaction) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "id", lua.LNumber(tx.ID)) + L.SetField(tbl, "host", lua.LString(tx.Host)) + L.SetField(tbl, "method", lua.LString(tx.Method)) + L.SetField(tbl, "url", lua.LString(tx.URL)) + L.SetField(tbl, "proto", lua.LString(tx.Proto)) + L.SetField(tbl, "status", lua.LNumber(tx.StatusCode)) + L.SetField(tbl, "duration_ms", lua.LNumber(tx.DurationMs)) + L.SetField(tbl, "tls", lua.LBool(tx.TLS)) + L.SetField(tbl, "in_scope", lua.LBool(tx.InScope)) + L.SetField(tbl, "req_body", lua.LString(tx.ReqBody)) + L.SetField(tbl, "resp_body", lua.LString(tx.RespBody)) + L.SetField(tbl, "timestamp", lua.LString(tx.Timestamp.UTC().Format(time.RFC3339))) + + reqH := L.NewTable() + for k, vals := range tx.ReqHeaders { + L.SetField(reqH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "req_headers", reqH) + + respH := L.NewTable() + for k, vals := range tx.RespHeaders { + L.SetField(respH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "resp_headers", respH) + + return tbl +} diff --git a/internal/plugin/engine.go b/internal/plugin/engine.go new file mode 100644 index 0000000..0623d36 --- /dev/null +++ b/internal/plugin/engine.go @@ -0,0 +1,347 @@ +package plugin + +import ( + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "strings" + + lua "github.com/yuin/gopher-lua" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +// Engine loads all Lua plugins from a directory and implements +// events.RequestMiddleware, events.ResponseObserver, +// events.WebSocketFrameObserver, and events.PluginEnabledObserver +// so it can be registered directly on the bus. +type Engine struct { + plugins []*Plugin + bus *events.Bus + dir string + st *store.Store +} + +// NewEngine loads all *.lua files from dir, registers each in the store, +// restores persisted enabled state, and returns a ready Engine. +// Plugins that fail to load are skipped. +func NewEngine(dir string, st *store.Store, bus *events.Bus) (*Engine, error) { + plugins, err := loadDir(dir, st) + if err != nil { + return nil, err + } + + e := &Engine{plugins: plugins, bus: bus, dir: dir, st: st} + + for _, p := range plugins { + // Persist plugin record; preserves existing enabled state on conflict. + if err := st.RegisterPlugin(p.name, p.path); err != nil { + logger.Error("plugin: register %s: %v", p.name, err) + } + // Restore enabled state from previous session. + if enabled, err := st.PluginEnabled(p.name); err == nil { + p.enabled.Store(enabled) + } + // Drain any log() calls made during on_load so they reach the UI. + e.emitLogs(p) + } + + active := make([]string, 0, len(plugins)) + for _, p := range plugins { + active = append(active, p.name) + } + if err := st.ScanPlugins(active); err != nil { + logger.Error("plugin: scan: %v", err) + } + + return e, nil +} + +// emitLogs drains a plugin's log buffer and emits each line onto the bus. +// Called after every hook invocation, including on error, so no log line +// is ever lost. +func (e *Engine) emitLogs(p *Plugin) { + for _, line := range p.drainLogs() { + e.bus.EmitPluginLog(events.PluginLogEvent{Name: p.name, Message: line}) + } +} + +// HandleRequest implements events.RequestMiddleware. +// Each plugin's on_request(req) hook is called in load order. Table +// construction and the Lua call happen in a single locked section via +// callWithBuilder — no VM access escapes the lock at any point. +// A plugin may modify method, url, headers, body, or set drop=true. +// The first drop short-circuits all remaining plugins. +func (e *Engine) HandleRequest(ev events.RequestEvent) events.RequestResult { + req := ev.Request + body := ev.Body + + if req == nil || req.URL == nil { + return events.RequestResult{Drop: false, Request: req, Body: body} + } + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_request") { + continue + } + + // Snapshot Go values before entering the lock so the builder closure + // captures plain Go strings, not live request fields that could race. + method := req.Method + rawURL := req.URL.String() + bodySnap := string(body) + headers := req.Header + + ret, err := p.callWithBuilder("on_request", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "method", lua.LString(method)) + L.SetField(tbl, "url", lua.LString(rawURL)) + L.SetField(tbl, "body", lua.LString(bodySnap)) + L.SetField(tbl, "drop", lua.LFalse) + + hTbl := L.NewTable() + for k, vals := range headers { + L.SetField(hTbl, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "headers", hTbl) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_request: %v", p.name, err) + continue + } + if ret == nil { + continue + } + + // Read back modifications. ret is a Lua table owned by p.L — + // mu must be held while reading it. + p.mu.Lock() + drop := false + if d, ok := ret.RawGetString("drop").(lua.LBool); ok { + drop = bool(d) + } + if !drop { + if m, ok := ret.RawGetString("method").(lua.LString); ok && string(m) != "" { + req.Method = string(m) + } + if u, ok := ret.RawGetString("url").(lua.LString); ok && string(u) != "" { + if parsed, err := req.URL.Parse(string(u)); err == nil { + req.URL = parsed + } + } + if b, ok := ret.RawGetString("body").(lua.LString); ok { + body = []byte(b) + } + if hTbl, ok := ret.RawGetString("headers").(*lua.LTable); ok { + newHeaders := make(http.Header) + hTbl.ForEach(func(k, v lua.LValue) { + newHeaders.Set(k.String(), v.String()) + }) + req.Header = newHeaders + } + } + p.mu.Unlock() + + if drop { + return events.RequestResult{Drop: true, Request: req, Body: body} + } + } + + return events.RequestResult{Drop: false, Request: req, Body: body} +} + +// ObserveResponse implements events.ResponseObserver. +// Each plugin's on_response(resp) hook is called in load order. Table +// construction and the Lua call happen in a single locked section via +// callWithBuilder — no VM access escapes the lock at any point. +// The hook is read-only — return value is ignored. +func (e *Engine) ObserveResponse(ev events.ResponseEvent) { + // Snapshot all event fields once before the plugin loop. Each plugin + // call captures these plain Go values in its builder closure — no live + // event fields are accessed inside the lock. + host := ev.Host + // ev.Host arrives as host:port from the proxy; expose bare hostname to plugins. + if h, _, err := net.SplitHostPort(ev.Host); err == nil { + host = h + } + method := ev.Method + url := ev.URL + proto := ev.Proto + status := ev.StatusCode + durationMs := ev.DurationMs + tls := ev.TLS + reqBody := string(ev.ReqBody) + respBody := string(ev.RespBody) + reqHeaders := ev.ReqHeaders + respHeaders := ev.RespHeaders + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_response") { + continue + } + + _, err := p.callWithBuilder("on_response", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "host", lua.LString(host)) + L.SetField(tbl, "method", lua.LString(method)) + L.SetField(tbl, "url", lua.LString(url)) + L.SetField(tbl, "proto", lua.LString(proto)) + L.SetField(tbl, "status", lua.LNumber(status)) + L.SetField(tbl, "duration_ms", lua.LNumber(durationMs)) + L.SetField(tbl, "tls", lua.LBool(tls)) + L.SetField(tbl, "req_body", lua.LString(reqBody)) + L.SetField(tbl, "resp_body", lua.LString(respBody)) + + rqH := L.NewTable() + for k, vals := range reqHeaders { + L.SetField(rqH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "req_headers", rqH) + + rpH := L.NewTable() + for k, vals := range respHeaders { + L.SetField(rpH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "resp_headers", rpH) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_response: %v", p.name, err) + } + } +} + +// ObserveWebSocketFrame implements events.WebSocketFrameObserver. +// Each plugin's on_websocket_frame(frame) hook is called in load order. +// The hook may modify frame.payload — the last modification wins. +// Plugins cannot drop frames. +func (e *Engine) ObserveWebSocketFrame(ev events.WebSocketFrameEvent) events.WebSocketFrameResult { + payload := ev.Payload + + // Snapshot Go values before the plugin loop. + connID := ev.ConnectionID + direction := int(ev.Direction) + opcode := int(ev.Opcode) + payloadSnap := string(ev.Payload) + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_websocket_frame") { + continue + } + + ret, err := p.callWithBuilder("on_websocket_frame", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "connection_id", lua.LNumber(connID)) + L.SetField(tbl, "direction", lua.LNumber(direction)) + L.SetField(tbl, "opcode", lua.LNumber(opcode)) + L.SetField(tbl, "payload", lua.LString(payloadSnap)) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_websocket_frame: %v", p.name, err) + continue + } + if ret == nil { + continue + } + + p.mu.Lock() + if b, ok := ret.RawGetString("payload").(lua.LString); ok { + payload = []byte(b) + payloadSnap = string(payload) + } + p.mu.Unlock() + } + + return events.WebSocketFrameResult{Payload: payload} +} + +// ObservePluginEnabled implements events.PluginEnabledObserver. +// Toggles the in-memory enabled flag for the named plugin so the change +// takes effect immediately for all subsequent hook calls. +func (e *Engine) ObservePluginEnabled(ev events.SetPluginEnabledEvent) { + for _, p := range e.plugins { + if p.name == ev.Name { + p.enabled.Store(ev.Enabled) + return + } + } +} + +// ObserveLoadPlugin implements events.LoadPluginObserver. +// Copies the selected file into the plugins directory, loads it, registers +// it in the store, and appends it to the engine's plugin slice so it +// participates in all subsequent hook calls. +// Safe without a mutex — the bus is synchronous so this and all hook +// iteration methods are always called sequentially. +func (e *Engine) ObserveLoadPlugin(ev events.LoadPluginEvent) { + filename := filepath.Base(ev.SourcePath) + destPath := filepath.Join(e.dir, filename) + + if err := copyFile(ev.SourcePath, destPath); err != nil { + logger.Error("plugin: import %s: %v", filename, err) + return + } + + p, err := load(destPath, e.st) + if err != nil { + logger.Error("plugin: load %s: %v", filename, err) + return + } + + if err := e.st.RegisterPlugin(p.name, p.path); err != nil { + logger.Error("plugin: register %s: %v", p.name, err) + } + + e.emitLogs(p) + e.plugins = append(e.plugins, p) +} + +// copyFile copies src to dst, creating the destination directory if needed. +func copyFile(src, dst string) error { + in, err := os.Open(src) + if err != nil { + return fmt.Errorf("open: %w", err) + } + defer in.Close() + + if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil { + return fmt.Errorf("mkdir: %w", err) + } + + out, err := os.Create(dst) + if err != nil { + return fmt.Errorf("create: %w", err) + } + defer out.Close() + + if _, err := io.Copy(out, in); err != nil { + return fmt.Errorf("copy: %w", err) + } + return out.Sync() +} + +// Close shuts down all plugin VMs. +func (e *Engine) Close() { + for _, p := range e.plugins { + p.close() + } +} diff --git a/internal/plugin/engine_test.go b/internal/plugin/engine_test.go new file mode 100644 index 0000000..33a4f8a --- /dev/null +++ b/internal/plugin/engine_test.go @@ -0,0 +1,1003 @@ +package plugin_test + +import ( + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/plugin" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +func pluginDir(t *testing.T, scripts map[string]string) string { + t.Helper() + dir := t.TempDir() + for name, src := range scripts { + require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(src), 0644)) + } + return dir +} + +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-plugin-engine-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + st, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { st.Close() }) + return st +} + +// newEngineWithBus creates an Engine using a caller-supplied bus so tests can +// register custom observers (e.g. to capture log events). +func newEngineWithBus(t *testing.T, st *store.Store, bus *events.Bus, scripts map[string]string) *plugin.Engine { + t.Helper() + dir := pluginDir(t, scripts) + e, err := plugin.NewEngine(dir, st, bus) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + return e +} + +// newEngine creates an Engine with a no-op bus — sufficient for tests that do +// not need to observe log or enable events. +func newEngine(t *testing.T, st *store.Store, scripts map[string]string) *plugin.Engine { + t.Helper() + return newEngineWithBus(t, st, events.NewBus(), scripts) +} + +func requestEvent(t *testing.T, method, rawURL string, body []byte) events.RequestEvent { + t.Helper() + return events.RequestEvent{ + Request: httptest.NewRequest(method, rawURL, nil), + Body: body, + } +} + +func responseEvent() events.ResponseEvent { + return events.ResponseEvent{ + Timestamp: time.Now(), + Host: "example.com", + Proto: "HTTP/1.1", + Method: "GET", + URL: "http://example.com/path", + ReqHeaders: http.Header{"X-Req": []string{"req-val"}}, + ReqBody: []byte("req body"), + StatusCode: 200, + RespHeaders: http.Header{"X-Resp": []string{"resp-val"}}, + RespBody: []byte("resp body"), + DurationMs: 55, + TLS: true, + } +} + +func wsFrameEvent(connID uint64, payload []byte) events.WebSocketFrameEvent { + return events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: payload, + } +} + +// ── NewEngine ───────────────────────────────────────────────────────────────── + +func TestNewEngine_EmptyDir_ReturnsEngine(t *testing.T) { + st := newTestStore(t) + e, err := plugin.NewEngine(t.TempDir(), st, events.NewBus()) + require.NoError(t, err) + assert.NotNil(t, e) + e.Close() +} + +func TestNewEngine_NonExistentDir_ReturnsEngine(t *testing.T) { + st := newTestStore(t) + e, err := plugin.NewEngine("/does/not/exist", st, events.NewBus()) + require.NoError(t, err) + assert.NotNil(t, e) + e.Close() +} + +func TestNewEngine_InvalidScript_SkippedEngineStillCreated(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "bad.lua": `%%% not lua`, + "good.lua": `x = 1`, + }) + assert.NotPanics(t, func() { + e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + }) +} + +func TestNewEngine_RegistersPluginsInStore(t *testing.T) { + st := newTestStore(t) + newEngine(t, st, map[string]string{ + "a.lua": `x = 1`, + "b.lua": `y = 2`, + }) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Len(t, entries, 2) +} + +func TestNewEngine_RestoresEnabledStateFromStore(t *testing.T) { + st := newTestStore(t) + dir := pluginDir(t, map[string]string{ + "test.lua": `function on_request(r) r.drop = true; return r end`, + }) + + // First load — plugin is enabled by default and drops requests. + e1, err := plugin.NewEngine(dir, st, events.NewBus()) + require.NoError(t, err) + result := e1.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.True(t, result.Drop) + e1.Close() + + // Disable via store directly. + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + // Second load — must restore disabled state from store. + e2, err := plugin.NewEngine(dir, st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e2.Close() }) + result = e2.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.False(t, result.Drop, "plugin disabled in store must not fire on_request") +} + +func TestNewEngine_ScanRemovesStaleRecords(t *testing.T) { + st := newTestStore(t) + + // Manually register a plugin that won't be on disk. + require.NoError(t, st.RegisterPlugin("ghost.lua", "/ghost.lua")) + + dir := pluginDir(t, map[string]string{"real.lua": `x = 1`}) + _, err := plugin.NewEngine(dir, st, events.NewBus()) + require.NoError(t, err) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "real.lua", entries[0].Name) +} + +// ── HandleRequest — no hook ─────────────────────────────────────────────────── + +func TestHandleRequest_NoPlugins_PassThrough(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, nil) + ev := requestEvent(t, "GET", "http://example.com/path", []byte("body")) + result := e.HandleRequest(ev) + assert.False(t, result.Drop) + assert.Equal(t, ev.Request, result.Request) + assert.Equal(t, ev.Body, result.Body) +} + +func TestHandleRequest_PluginWithNoOnRequest_PassThrough(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{"noop.lua": `-- no hooks`}) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.False(t, result.Drop) +} + +func TestHandleRequest_NilRequest_PassThrough(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "hook.lua": `function on_request(r) r.drop = true; return r end`, + }) + result := e.HandleRequest(events.RequestEvent{Request: nil, Body: nil}) + assert.False(t, result.Drop) +} + +// ── HandleRequest — drop ────────────────────────────────────────────────────── + +func TestHandleRequest_Drop_ReturnsDrop(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "drop.lua": `function on_request(r) r.drop = true; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.True(t, result.Drop) +} + +func TestHandleRequest_FirstPluginDrops_SecondNotCalled(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_drop.lua": `function on_request(r) r.drop = true; return r end`, + "2_panic.lua": `function on_request(r) error("must not be called") end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.True(t, result.Drop) +} + +// ── HandleRequest — modifications ───────────────────────────────────────────── + +func TestHandleRequest_ModifiesBody(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "body.lua": `function on_request(r) r.body = "injected"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "POST", "http://example.com", []byte("original"))) + assert.False(t, result.Drop) + assert.Equal(t, []byte("injected"), result.Body) +} + +func TestHandleRequest_ModifiesMethod(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "method.lua": `function on_request(r) r.method = "POST"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.Equal(t, "POST", result.Request.Method) +} + +func TestHandleRequest_ModifiesURL(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "url.lua": `function on_request(r) r.url = "http://example.com/modified"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com/original", nil)) + assert.Equal(t, "/modified", result.Request.URL.Path) +} + +func TestHandleRequest_ModifiesHeaders(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "headers.lua": `function on_request(r) r.headers["X-Plugin"] = "injected"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.Equal(t, "injected", result.Request.Header.Get("X-Plugin")) +} + +func TestHandleRequest_ChainedPlugins_ModificationsAccumulate(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_first.lua": `function on_request(r) r.body = r.body .. "-first"; return r end`, + "2_second.lua": `function on_request(r) r.body = r.body .. "-second"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", []byte("start"))) + assert.Equal(t, []byte("start-first-second"), result.Body) +} + +func TestHandleRequest_LuaError_PluginSkipped_ChainContinues(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_error.lua": `function on_request(r) error("oops") end`, + "2_ok.lua": `function on_request(r) r.body = "reached"; return r end`, + }) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", []byte("original"))) + assert.False(t, result.Drop) + assert.Equal(t, []byte("reached"), result.Body) +} + +// ── HandleRequest — enabled flag ────────────────────────────────────────────── + +func TestHandleRequest_DisabledPlugin_HookNotCalled(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "drop.lua": `function on_request(r) r.drop = true; return r end`, + }) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "drop.lua", Enabled: false}) + + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.False(t, result.Drop, "disabled plugin must not fire on_request") +} + +func TestHandleRequest_ReenabledPlugin_HookCalledAgain(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "drop.lua": `function on_request(r) r.drop = true; return r end`, + }) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "drop.lua", Enabled: false}) + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.False(t, result.Drop) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "drop.lua", Enabled: true}) + result = e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.True(t, result.Drop, "re-enabled plugin must fire on_request again") +} + +func TestHandleRequest_DisabledPluginInChain_OthersStillRun(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_disabled.lua": `function on_request(r) r.body = r.body .. "-disabled"; return r end`, + "2_enabled.lua": `function on_request(r) r.body = r.body .. "-enabled"; return r end`, + }) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "1_disabled.lua", Enabled: false}) + + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", []byte("start"))) + assert.Equal(t, []byte("start-enabled"), result.Body) +} + +// ── ObservePluginEnabled ────────────────────────────────────────────────────── + +func TestObservePluginEnabled_UnknownPlugin_NoPanic(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, nil) + assert.NotPanics(t, func() { + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "ghost.lua", Enabled: false}) + }) +} + +func TestObservePluginEnabled_TakesEffectImmediately(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "drop.lua": `function on_request(r) r.drop = true; return r end`, + }) + + // Enabled by default — drops. + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.True(t, result.Drop) + + // Disable — next call must pass through. + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "drop.lua", Enabled: false}) + result = e.HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + assert.False(t, result.Drop) +} + +// ── ObserveResponse ─────────────────────────────────────────────────────────── + +func TestObserveResponse_NoPlugins_NoPanic(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, nil) + assert.NotPanics(t, func() { e.ObserveResponse(responseEvent()) }) +} + +func TestObserveResponse_ReceivesCorrectFields(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "observe.lua": ` + function on_response(r) + db.loot_add(r.host .. "|" .. r.method .. "|" .. tostring(r.status) .. "|" .. tostring(r.tls), "Info") + end + `, + }) + + ev := responseEvent() + e.ObserveResponse(ev) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, fmt.Sprintf("%s|%s|%d|%v", ev.Host, ev.Method, ev.StatusCode, ev.TLS), loot[0].Title) +} + +func TestObserveResponse_MultiplePluginsAllCalled(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "a.lua": `function on_response(r) db.loot_add("a", "Info") end`, + "b.lua": `function on_response(r) db.loot_add("b", "Info") end`, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + loot, err := st.AllLoot() + require.NoError(t, err) + assert.Len(t, loot, 2) +} + +func TestObserveResponse_LuaError_OtherPluginsContinue(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_error.lua": `function on_response(r) error("fail") end`, + "2_ok.lua": `function on_response(r) db.loot_add("ok", "Info") end`, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "ok", loot[0].Title) +} + +func TestObserveResponse_DisabledPlugin_HookNotCalled(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "obs.lua": `function on_response(r) db.loot_add("called", "Info") end`, + }) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "obs.lua", Enabled: false}) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + assert.Empty(t, loot, "disabled plugin must not fire on_response") +} + +// ── ObserveWebSocketFrame ───────────────────────────────────────────────────── + +func TestObserveWebSocketFrame_NoPlugins_ReturnsOriginalPayload(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, nil) + ev := wsFrameEvent(1, []byte("hello")) + result := e.ObserveWebSocketFrame(ev) + assert.Equal(t, []byte("hello"), result.Payload) +} + +func TestObserveWebSocketFrame_PluginWithNoHook_ReturnsOriginalPayload(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{"noop.lua": `-- no hooks`}) + result := e.ObserveWebSocketFrame(wsFrameEvent(1, []byte("original"))) + assert.Equal(t, []byte("original"), result.Payload) +} + +func TestObserveWebSocketFrame_PluginModifiesPayload(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "ws.lua": ` + function on_websocket_frame(f) + f.payload = f.payload .. "-modified" + return f + end + `, + }) + result := e.ObserveWebSocketFrame(wsFrameEvent(1, []byte("hello"))) + assert.Equal(t, []byte("hello-modified"), result.Payload) +} + +func TestObserveWebSocketFrame_ChainedPlugins_ModificationsAccumulate(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_first.lua": `function on_websocket_frame(f) f.payload = f.payload .. "-first"; return f end`, + "2_second.lua": `function on_websocket_frame(f) f.payload = f.payload .. "-second"; return f end`, + }) + result := e.ObserveWebSocketFrame(wsFrameEvent(1, []byte("start"))) + assert.Equal(t, []byte("start-first-second"), result.Payload) +} + +func TestObserveWebSocketFrame_ReceivesCorrectFields(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "ws.lua": ` + function on_websocket_frame(f) + db.loot_add( + tostring(f.connection_id) .. "|" .. tostring(f.direction) .. "|" .. tostring(f.opcode), + "Info" + ) + return f + end + `, + }) + + ev := events.WebSocketFrameEvent{ + ConnectionID: 42, + Timestamp: time.Now(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketBinary, + Payload: []byte("data"), + } + e.ObserveWebSocketFrame(ev) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "42|1|2", loot[0].Title) +} + +func TestObserveWebSocketFrame_DisabledPlugin_HookNotCalled(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "ws.lua": `function on_websocket_frame(f) f.payload = "modified"; return f end`, + }) + + e.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "ws.lua", Enabled: false}) + result := e.ObserveWebSocketFrame(wsFrameEvent(1, []byte("original"))) + assert.Equal(t, []byte("original"), result.Payload, "disabled plugin must not modify frame") +} + +func TestObserveWebSocketFrame_LuaError_ReturnsLastGoodPayload(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "1_ok.lua": `function on_websocket_frame(f) f.payload = "good"; return f end`, + "2_error.lua": `function on_websocket_frame(f) error("fail") end`, + }) + result := e.ObserveWebSocketFrame(wsFrameEvent(1, []byte("original"))) + assert.Equal(t, []byte("good"), result.Payload) +} + +// ── ObserveLoadPlugin ───────────────────────────────────────────────────────── + +func TestObserveLoadPlugin_LoadsAndRegistersPlugin(t *testing.T) { + st := newTestStore(t) + destDir := t.TempDir() + e, err := plugin.NewEngine(destDir, st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + // Write source file to a separate temp dir. + srcDir := t.TempDir() + srcPath := filepath.Join(srcDir, "new.lua") + require.NoError(t, os.WriteFile(srcPath, []byte(` + function on_request(r) + r.body = "from-new-plugin" + return r + end + `), 0644)) + + e.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: srcPath}) + + // Plugin must be active immediately. + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", []byte("original"))) + assert.Equal(t, []byte("from-new-plugin"), result.Body) +} + +func TestObserveLoadPlugin_CopiesFileToPluginsDir(t *testing.T) { + st := newTestStore(t) + destDir := t.TempDir() + e, err := plugin.NewEngine(destDir, st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + srcPath := filepath.Join(t.TempDir(), "copied.lua") + require.NoError(t, os.WriteFile(srcPath, []byte(`x = 1`), 0644)) + + e.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: srcPath}) + + _, err = os.Stat(filepath.Join(destDir, "copied.lua")) + assert.NoError(t, err, "file must be copied to plugins directory") +} + +func TestObserveLoadPlugin_RegistersInStore(t *testing.T) { + st := newTestStore(t) + destDir := t.TempDir() + e, err := plugin.NewEngine(destDir, st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + srcPath := filepath.Join(t.TempDir(), "registered.lua") + require.NoError(t, os.WriteFile(srcPath, []byte(`x = 1`), 0644)) + + e.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: srcPath}) + time.Sleep(50 * time.Millisecond) + + entries, err := st.AllPlugins() + require.NoError(t, err) + names := make([]string, 0, len(entries)) + for _, e := range entries { + names = append(names, e.Name) + } + assert.Contains(t, names, "registered.lua") +} + +func TestObserveLoadPlugin_BadSourcePath_NoPanic(t *testing.T) { + st := newTestStore(t) + e, err := plugin.NewEngine(t.TempDir(), st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + assert.NotPanics(t, func() { + e.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: "/does/not/exist.lua"}) + }) +} + +func TestObserveLoadPlugin_InvalidLua_NoPanic(t *testing.T) { + st := newTestStore(t) + destDir := t.TempDir() + e, err := plugin.NewEngine(destDir, st, events.NewBus()) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + srcPath := filepath.Join(t.TempDir(), "bad.lua") + require.NoError(t, os.WriteFile(srcPath, []byte(`%%% invalid lua`), 0644)) + + assert.NotPanics(t, func() { + e.ObserveLoadPlugin(events.LoadPluginEvent{SourcePath: srcPath}) + }) +} + +// ── log() routing ───────────────────────────────────────────────────────────── + +func TestLog_RoutedThroughBus(t *testing.T) { + st := newTestStore(t) + bus := events.NewBus() + + var mu sync.Mutex + var logLines []string + bus.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { + mu.Lock() + logLines = append(logLines, e.Name+": "+e.Message) + mu.Unlock() + })) + + newEngineWithBus(t, st, bus, map[string]string{ + "logger.lua": ` + function on_request(r) + log("intercepted request") + return r + end + `, + }) + + newEngineWithBus(t, st, bus, map[string]string{ + "logger.lua": ` + function on_request(r) + log("intercepted request") + return r + end + `, + }).HandleRequest(requestEvent(t, "GET", "http://example.com", nil)) + + mu.Lock() + found := false + for _, l := range logLines { + if l == "logger.lua: intercepted request" { + found = true + break + } + } + mu.Unlock() + assert.True(t, found, "log() call must emit PluginLogEvent on the bus") +} + +func TestLog_OnLoad_DrainedAfterNewEngine(t *testing.T) { + st := newTestStore(t) + bus := events.NewBus() + + var mu sync.Mutex + var logLines []string + bus.Register(events.PluginLogObserverFunc(func(e events.PluginLogEvent) { + mu.Lock() + logLines = append(logLines, e.Message) + mu.Unlock() + })) + + dir := pluginDir(t, map[string]string{ + "startup.lua": ` + function on_load() + log("loaded ok") + end + `, + }) + e, err := plugin.NewEngine(dir, st, bus) + require.NoError(t, err) + t.Cleanup(func() { e.Close() }) + + mu.Lock() + lines := append([]string{}, logLines...) + mu.Unlock() + assert.Contains(t, lines, "loaded ok", "on_load log() calls must be emitted before NewEngine returns") +} + +// ── Engine.Close ────────────────────────────────────────────────────────────── + +func TestClose_NoPanic(t *testing.T) { + st := newTestStore(t) + e, err := plugin.NewEngine(t.TempDir(), st, events.NewBus()) + require.NoError(t, err) + assert.NotPanics(t, func() { e.Close() }) +} + +// ── API: log ────────────────────────────────────────────────────────────────── + +func TestAPI_Log_IsAvailable(t *testing.T) { + st := newTestStore(t) + dir := pluginDir(t, map[string]string{ + "log.lua": `function on_load() log("hello from plugin") end`, + }) + e, err := plugin.NewEngine(dir, st, events.NewBus()) + require.NoError(t, err) + e.Close() +} + +// ── API: http.get ───────────────────────────────────────────────────────────── + +func TestAPI_HTTPGet_ReturnsStatusAndBody(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, "pong") + })) + defer upstream.Close() + + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "get.lua": fmt.Sprintf(` + function on_response(r) + local resp = http.get("%s") + if resp then db.loot_add(tostring(resp.status) .. "|" .. resp.body, "Info") end + end + `, upstream.URL), + }) + e.ObserveResponse(responseEvent()) + time.Sleep(100 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "200|pong", loot[0].Title) +} + +func TestAPI_HTTPGet_BadURL_ReturnsNilAndError(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "get_fail.lua": ` + function on_response(r) + local resp, err = http.get("http://127.0.0.1:1/") + if resp == nil and err ~= nil then db.loot_add("error_handled", "Info") end + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(200 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "error_handled", loot[0].Title) +} + +// ── API: http.post ──────────────────────────────────────────────────────────── + +func TestAPI_HTTPPost_SendsBodyAndReturnsResponse(t *testing.T) { + var receivedBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + buf := make([]byte, 512) + n, _ := r.Body.Read(buf) + receivedBody = string(buf[:n]) + w.WriteHeader(http.StatusCreated) + fmt.Fprint(w, "created") + })) + defer upstream.Close() + + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "post.lua": fmt.Sprintf(` + function on_response(r) + local resp = http.post("%s", '{"key":"val"}', "application/json") + if resp then db.loot_add(tostring(resp.status) .. "|" .. resp.body, "Info") end + end + `, upstream.URL), + }) + e.ObserveResponse(responseEvent()) + time.Sleep(100 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "201|created", loot[0].Title) + assert.Equal(t, `{"key":"val"}`, receivedBody) +} + +// ── API: http.request ───────────────────────────────────────────────────────── + +func TestAPI_HTTPRequest_SendsCustomMethodAndHeaders(t *testing.T) { + var receivedMethod, receivedHeader string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedMethod = r.Method + receivedHeader = r.Header.Get("X-Custom") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "request.lua": fmt.Sprintf(` + function on_response(r) + local resp = http.request("DELETE", "%s", "", {["X-Custom"] = "plugin-value"}) + if resp then db.loot_add(tostring(resp.status), "Info") end + end + `, upstream.URL), + }) + e.ObserveResponse(responseEvent()) + time.Sleep(100 * time.Millisecond) + + assert.Equal(t, "DELETE", receivedMethod) + assert.Equal(t, "plugin-value", receivedHeader) +} + +// ── API: db.history ─────────────────────────────────────────────────────────── + +func TestAPI_DBHistory_ReturnsTransactions(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(store.Transaction{ + Timestamp: time.Now(), Host: "target.com", Method: "GET", + URL: "http://target.com/", Proto: "HTTP/1.1", StatusCode: 200, InScope: true, + })) + time.Sleep(50 * time.Millisecond) + + e := newEngine(t, st, map[string]string{ + "history.lua": ` + function on_response(r) + local txs = db.history(10) + if txs then db.loot_add(tostring(#txs), "Info") end + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "1", loot[0].Title) +} + +func TestAPI_DBHistory_TransactionFieldsExposed(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(store.Transaction{ + Timestamp: time.Now(), Host: "check.com", Method: "POST", + URL: "http://check.com/api", Proto: "HTTP/1.1", StatusCode: 201, InScope: true, + })) + time.Sleep(50 * time.Millisecond) + + e := newEngine(t, st, map[string]string{ + "fields.lua": ` + function on_response(r) + local txs = db.history(1) + if txs and #txs > 0 then + local tx = txs[1] + db.loot_add(tx.host .. "|" .. tx.method .. "|" .. tostring(tx.status), "Info") + end + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "check.com|POST|201", loot[0].Title) +} + +// ── API: db.history_filter ──────────────────────────────────────────────────── + +func TestAPI_DBHistoryFilter_FiltersByHost(t *testing.T) { + st := newTestStore(t) + for _, host := range []string{"alpha.com", "beta.com", "alpha.com"} { + require.NoError(t, st.Log(store.Transaction{ + Timestamp: time.Now(), Host: host, Method: "GET", + URL: "http://" + host + "/", Proto: "HTTP/1.1", StatusCode: 200, InScope: true, + })) + } + time.Sleep(50 * time.Millisecond) + + e := newEngine(t, st, map[string]string{ + "filter.lua": ` + function on_response(r) + local txs = db.history_filter("alpha.com") + db.loot_add(tostring(#txs), "Info") + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "2", loot[0].Title) +} + +// ── API: db.scope ───────────────────────────────────────────────────────────── + +func TestAPI_DBScope_AddGetRemove(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "scope.lua": ` + function on_response(r) + db.scope_add("plugin-scope.com") + local entries = db.scope_get() + db.loot_add("count:" .. tostring(#entries), "Info") + if #entries > 0 then db.scope_remove(entries[1].id) end + local after = db.scope_get() + db.loot_add("after:" .. tostring(#after), "Info") + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 2) + titles := map[string]bool{} + for _, l := range loot { + titles[l.Title] = true + } + assert.True(t, titles["count:1"]) + assert.True(t, titles["after:0"]) +} + +// ── API: db.loot ────────────────────────────────────────────────────────────── + +func TestAPI_DBLoot_AllFieldsStored(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "loot.lua": ` + function on_response(r) + db.loot_add("my finding", "Critical", "detailed notes", "raw req", "raw resp") + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, loot, 1) + assert.Equal(t, "my finding", loot[0].Title) + assert.Equal(t, "Critical", loot[0].Severity) + assert.Equal(t, "detailed notes", loot[0].Notes) + assert.Equal(t, "raw req", loot[0].RawRequest) + assert.Equal(t, "raw resp", loot[0].RawResponse) +} + +func TestAPI_DBLoot_AddGetDelete(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "loot.lua": ` + function on_response(r) + local id = db.loot_add("finding", "High", "notes", "req", "resp") + local entries = db.loot_get() + db.loot_add("count:" .. tostring(#entries), "Info") + db.loot_delete(id) + local after = db.loot_get() + db.loot_add("after:" .. tostring(#after), "Info") + end + `, + }) + e.ObserveResponse(responseEvent()) + time.Sleep(50 * time.Millisecond) + + loot, err := st.AllLoot() + require.NoError(t, err) + titles := map[string]bool{} + for _, l := range loot { + titles[l.Title] = true + } + assert.True(t, titles["count:1"]) + assert.True(t, titles["after:1"]) +} + +// ── Concurrency ─────────────────────────────────────────────────────────────── + +func TestEngine_ConcurrentHandleRequest_Safe(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{ + "concurrent.lua": `function on_request(r) r.body = r.body .. "-ok"; return r end`, + }) + const goroutines = 30 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + result := e.HandleRequest(requestEvent(t, "GET", "http://example.com", []byte("body"))) + assert.False(t, result.Drop) + }() + } + wg.Wait() +} + +func TestEngine_ConcurrentObserveResponse_Safe(t *testing.T) { + st := newTestStore(t) + e := newEngine(t, st, map[string]string{"obs.lua": `function on_response(r) end`}) + const goroutines = 30 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + e.ObserveResponse(responseEvent()) + }() + } + wg.Wait() +} diff --git a/internal/plugin/plugin.go b/internal/plugin/plugin.go new file mode 100644 index 0000000..86496a1 --- /dev/null +++ b/internal/plugin/plugin.go @@ -0,0 +1,197 @@ +package plugin + +import ( + "context" + "fmt" + "os" + "path/filepath" + "sync" + "sync/atomic" + "time" + + lua "github.com/yuin/gopher-lua" + + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +const callTimeout = 500 * time.Millisecond + +// Plugin is a single loaded Lua script with its own isolated VM. +// mu serialises all VM access — gopher-lua LState is not goroutine-safe. +type Plugin struct { + name string + path string + L *lua.LState + mu sync.Mutex + + // enabled controls whether this plugin's hooks are called. + // Uses atomic.Bool so the engine can check it without acquiring mu. + enabled atomic.Bool + + // logBuf accumulates log() calls made during a Lua hook execution. + // Written only while mu is held (log() can only be called from within + // a Lua hook, which always runs under mu). Drained by the engine after + // each hook call via drainLogs, which also acquires mu. + logBuf []string +} + +// load reads the Lua file, registers the full API, executes the script, +// then calls on_load() if defined. +func load(path string, st *store.Store) (*Plugin, error) { + name := filepath.Base(path) + + L := lua.NewState(lua.Options{ + SkipOpenLibs: false, + }) + + p := &Plugin{name: name, path: path, L: L} + p.enabled.Store(true) + + // Register API before executing the script so on_load can use it. + // The logFn closure appends to p.logBuf. It is safe without acquiring + // mu because log() is only ever invoked from within a Lua hook, which + // always runs under mu via callNoReturn or callWithBuilder. + registerAPI(L, st, func(msg string) { + p.logBuf = append(p.logBuf, msg) + }) + + if err := L.DoFile(path); err != nil { + L.Close() + return nil, fmt.Errorf("plugin %s: load: %w", name, err) + } + + // Call on_load() if defined — errors are logged but not fatal. + if L.GetGlobal("on_load") != lua.LNil { + if err := p.callNoReturn("on_load", lua.LNil); err != nil { + logger.Error("plugin %s: on_load: %v", name, err) + } + } + + return p, nil +} + +// callNoReturn invokes a Lua function that returns nothing. +// Holds mu for the duration. Enforces callTimeout. Recovers panics. +func (p *Plugin) callNoReturn(fn string, arg lua.LValue) (retErr error) { + p.mu.Lock() + defer p.mu.Unlock() + + defer func() { + if r := recover(); r != nil { + retErr = fmt.Errorf("panic: %v", r) + } + }() + + f := p.L.GetGlobal(fn) + if f == lua.LNil { + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), callTimeout) + defer cancel() + p.L.SetContext(ctx) + defer p.L.SetContext(context.Background()) + + params := lua.P{Fn: f, NRet: 0, Protect: true} + if arg != lua.LNil { + return p.L.CallByParam(params, arg) + } + return p.L.CallByParam(params) +} + +// callWithBuilder builds a Lua table and calls fn with it in a single locked +// section, returning the result table. builder receives the VM and must +// construct and return the argument table — it runs under mu. The Lua call +// also runs under mu, so no VM access escapes the lock at any point. +// Returns the returned table, or nil if the function returned nothing. +func (p *Plugin) callWithBuilder(fn string, builder func(L *lua.LState) *lua.LTable) (ret *lua.LTable, retErr error) { + p.mu.Lock() + defer p.mu.Unlock() + + defer func() { + if r := recover(); r != nil { + retErr = fmt.Errorf("panic: %v", r) + } + }() + + f := p.L.GetGlobal(fn) + if f == lua.LNil { + return nil, nil + } + + arg := builder(p.L) + + ctx, cancel := context.WithTimeout(context.Background(), callTimeout) + defer cancel() + p.L.SetContext(ctx) + defer p.L.SetContext(context.Background()) + + if err := p.L.CallByParam(lua.P{Fn: f, NRet: 1, Protect: true}, arg); err != nil { + return nil, err + } + + top := p.L.Get(-1) + p.L.Pop(1) + + if tbl, ok := top.(*lua.LTable); ok { + return tbl, nil + } + return nil, nil +} + +// drainLogs returns all buffered log lines and clears the buffer. +// Called by the engine after each hook invocation so log lines can be +// emitted onto the bus. Acquires mu independently of the hook call. +func (p *Plugin) drainLogs() []string { + p.mu.Lock() + defer p.mu.Unlock() + if len(p.logBuf) == 0 { + return nil + } + lines := p.logBuf + p.logBuf = nil + return lines +} + +// close shuts down the Lua VM. +func (p *Plugin) close() { + p.mu.Lock() + defer p.mu.Unlock() + p.L.Close() +} + +// has returns true if the plugin defines the named global function. +func (p *Plugin) has(fn string) bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.L.GetGlobal(fn) != lua.LNil +} + +// loadDir scans dir for *.lua files and loads each as a Plugin. +// Files that fail to load are logged and skipped. +func loadDir(dir string, st *store.Store) ([]*Plugin, error) { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, fmt.Errorf("plugin: read dir %s: %w", dir, err) + } + + var plugins []*Plugin + for _, e := range entries { + if e.IsDir() || filepath.Ext(e.Name()) != ".lua" { + continue + } + path := filepath.Join(dir, e.Name()) + p, err := load(path, st) + if err != nil { + logger.Error("plugin: skip %s: %v", e.Name(), err) + continue + } + logger.Info("plugin: loaded %s", e.Name()) + plugins = append(plugins, p) + } + return plugins, nil +} diff --git a/internal/plugin/plugin_test.go b/internal/plugin/plugin_test.go new file mode 100644 index 0000000..d020c31 --- /dev/null +++ b/internal/plugin/plugin_test.go @@ -0,0 +1,391 @@ +package plugin + +// White-box tests. Package plugin (not plugin_test) so unexported types and +// functions — load, loadDir, Plugin.has, Plugin.callWithBuilder, +// Plugin.callNoReturn — are directly accessible. + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + lua "github.com/yuin/gopher-lua" + + "github.com/shiv/internal/store" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +// writeLua writes a Lua script to a temp file and returns its path. +// The file is removed when the test ends. +func writeLua(t *testing.T, src string) string { + t.Helper() + f, err := os.CreateTemp("", "shiv-plugin-*.lua") + require.NoError(t, err) + _, err = f.WriteString(src) + require.NoError(t, err) + require.NoError(t, f.Close()) + t.Cleanup(func() { os.Remove(f.Name()) }) + return f.Name() +} + +// newTestStore opens a temp SQLite store for plugin tests. +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-plugin-store-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + st, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { st.Close() }) + return st +} + +// loadPlugin is a convenience wrapper around load that also registers cleanup. +func loadPlugin(t *testing.T, src string, st *store.Store) *Plugin { + t.Helper() + path := writeLua(t, src) + p, err := load(path, st) + require.NoError(t, err) + t.Cleanup(func() { p.close() }) + return p +} + +// ── load ────────────────────────────────────────────────────────────────────── + +func TestLoad_ValidScript_ReturnsPlugin(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `-- empty script`, st) + assert.NotNil(t, p) +} + +func TestLoad_InvalidScript_ReturnsError(t *testing.T) { + st := newTestStore(t) + path := writeLua(t, `this is not valid lua %%%`) + _, err := load(path, st) + assert.Error(t, err) +} + +func TestLoad_NonExistentFile_ReturnsError(t *testing.T) { + st := newTestStore(t) + _, err := load("/does/not/exist.lua", st) + assert.Error(t, err) +} + +func TestLoad_CallsOnLoad(t *testing.T) { + st := newTestStore(t) + // on_load sets a global that we can check afterwards. + p := loadPlugin(t, ` + loaded = false + function on_load() + loaded = true + end + `, st) + + p.mu.Lock() + v := p.L.GetGlobal("loaded") + p.mu.Unlock() + + assert.Equal(t, lua.LTrue, v) +} + +func TestLoad_OnLoadError_PluginStillLoaded(t *testing.T) { + st := newTestStore(t) + // on_load errors are logged but must not prevent the plugin from loading. + p, err := func() (*Plugin, error) { + path := writeLua(t, ` + function on_load() + error("intentional on_load failure") + end + `) + return load(path, st) + }() + // load itself should succeed despite on_load failing. + require.NoError(t, err) + assert.NotNil(t, p) + p.close() +} + +func TestLoad_SetsName(t *testing.T) { + st := newTestStore(t) + path := writeLua(t, `-- empty`) + p, err := load(path, st) + require.NoError(t, err) + defer p.close() + assert.Equal(t, filepath.Base(path), p.name) +} + +// ── has ─────────────────────────────────────────────────────────────────────── + +func TestHas_DefinedFunction_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `function on_request(r) return r end`, st) + assert.True(t, p.has("on_request")) +} + +func TestHas_UndefinedFunction_ReturnsFalse(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `-- empty`, st) + assert.False(t, p.has("on_request")) +} + +func TestHas_NilGlobal_ReturnsFalse(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `on_request = nil`, st) + assert.False(t, p.has("on_request")) +} + +// ── callNoReturn ────────────────────────────────────────────────────────────── + +func TestCallNoReturn_UndefinedFunction_ReturnsNil(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `-- empty`, st) + err := p.callNoReturn("no_such_fn", lua.LNil) + assert.NoError(t, err) +} + +func TestCallNoReturn_FunctionExecutes(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + called = false + function my_fn() + called = true + end + `, st) + + err := p.callNoReturn("my_fn", lua.LNil) + require.NoError(t, err) + + p.mu.Lock() + v := p.L.GetGlobal("called") + p.mu.Unlock() + assert.Equal(t, lua.LTrue, v) +} + +func TestCallNoReturn_FunctionReceivesArg(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + received = nil + function my_fn(arg) + received = arg + end + `, st) + + p.mu.Lock() + arg := p.L.NewTable() + p.L.SetField(arg, "key", lua.LString("value")) + p.mu.Unlock() + + err := p.callNoReturn("my_fn", arg) + require.NoError(t, err) + + p.mu.Lock() + tbl, ok := p.L.GetGlobal("received").(*lua.LTable) + p.mu.Unlock() + require.True(t, ok) + + p.mu.Lock() + v := tbl.RawGetString("key") + p.mu.Unlock() + assert.Equal(t, lua.LString("value"), v) +} + +func TestCallNoReturn_LuaError_ReturnsError(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + function bad_fn() + error("intentional error") + end + `, st) + + err := p.callNoReturn("bad_fn", lua.LNil) + assert.Error(t, err) +} + +func TestCallNoReturn_Timeout_ReturnsError(t *testing.T) { + st := newTestStore(t) + // Infinite loop — must be killed by the context timeout. + p := loadPlugin(t, ` + function infinite() + while true do end + end + `, st) + + start := time.Now() + err := p.callNoReturn("infinite", lua.LNil) + elapsed := time.Since(start) + + assert.Error(t, err) + // Should terminate within a reasonable margin of callTimeout (500ms). + assert.Less(t, elapsed, callTimeout+300*time.Millisecond) +} + +// ── callWithBuilder ─────────────────────────────────────────────────────────── + +func TestCallWithBuilder_UndefinedFunction_ReturnsNilNoError(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `-- empty`, st) + + ret, err := p.callWithBuilder("no_such_fn", func(L *lua.LState) *lua.LTable { + return L.NewTable() + }) + + assert.NoError(t, err) + assert.Nil(t, ret) +} + +func TestCallWithBuilder_BuilderRunsUnderLock(t *testing.T) { + // Verify the builder is called (indirectly — if it panics the test fails). + st := newTestStore(t) + p := loadPlugin(t, `function fn(t) return t end`, st) + + builderCalled := false + ret, err := p.callWithBuilder("fn", func(L *lua.LState) *lua.LTable { + builderCalled = true + tbl := L.NewTable() + L.SetField(tbl, "x", lua.LString("hello")) + return tbl + }) + + require.NoError(t, err) + assert.True(t, builderCalled) + assert.NotNil(t, ret) +} + +func TestCallWithBuilder_FunctionCanModifyTable(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + function transform(t) + t.result = "modified" + return t + end + `, st) + + ret, err := p.callWithBuilder("transform", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "result", lua.LString("original")) + return tbl + }) + + require.NoError(t, err) + require.NotNil(t, ret) + + p.mu.Lock() + v := ret.RawGetString("result") + p.mu.Unlock() + assert.Equal(t, lua.LString("modified"), v) +} + +func TestCallWithBuilder_FunctionReturnsNil_ReturnsNilTable(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, `function fn(t) return nil end`, st) + + ret, err := p.callWithBuilder("fn", func(L *lua.LState) *lua.LTable { + return L.NewTable() + }) + + assert.NoError(t, err) + assert.Nil(t, ret) +} + +func TestCallWithBuilder_LuaError_ReturnsError(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + function bad(t) + error("builder error") + end + `, st) + + _, err := p.callWithBuilder("bad", func(L *lua.LState) *lua.LTable { + return L.NewTable() + }) + + assert.Error(t, err) +} + +func TestCallWithBuilder_Timeout_ReturnsError(t *testing.T) { + st := newTestStore(t) + p := loadPlugin(t, ` + function slow(t) + while true do end + end + `, st) + + start := time.Now() + _, err := p.callWithBuilder("slow", func(L *lua.LState) *lua.LTable { + return L.NewTable() + }) + elapsed := time.Since(start) + + assert.Error(t, err) + assert.Less(t, elapsed, callTimeout+300*time.Millisecond) +} + +// ── loadDir ─────────────────────────────────────────────────────────────────── + +func TestLoadDir_NonExistentDir_ReturnsNilSlice(t *testing.T) { + st := newTestStore(t) + plugins, err := loadDir("/does/not/exist", st) + assert.NoError(t, err) + assert.Nil(t, plugins) +} + +func TestLoadDir_EmptyDir_ReturnsEmpty(t *testing.T) { + st := newTestStore(t) + dir := t.TempDir() + plugins, err := loadDir(dir, st) + assert.NoError(t, err) + assert.Empty(t, plugins) +} + +func TestLoadDir_LoadsLuaFiles(t *testing.T) { + st := newTestStore(t) + dir := t.TempDir() + + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.lua"), []byte(`x = 1`), 0644)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "b.lua"), []byte(`y = 2`), 0644)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "readme.txt"), []byte(`not lua`), 0644)) + + plugins, err := loadDir(dir, st) + require.NoError(t, err) + assert.Len(t, plugins, 2) + for _, p := range plugins { + p.close() + } +} + +func TestLoadDir_SkipsDirectories(t *testing.T) { + st := newTestStore(t) + dir := t.TempDir() + + require.NoError(t, os.Mkdir(filepath.Join(dir, "subdir.lua"), 0755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "real.lua"), []byte(`z = 3`), 0644)) + + plugins, err := loadDir(dir, st) + require.NoError(t, err) + assert.Len(t, plugins, 1) + for _, p := range plugins { + p.close() + } +} + +func TestLoadDir_InvalidLuaSkipped_OthersLoaded(t *testing.T) { + st := newTestStore(t) + dir := t.TempDir() + + require.NoError(t, os.WriteFile(filepath.Join(dir, "good.lua"), []byte(`g = 1`), 0644)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "bad.lua"), []byte(`%%% invalid`), 0644)) + + plugins, err := loadDir(dir, st) + require.NoError(t, err) + // bad.lua skipped, good.lua loaded + assert.Len(t, plugins, 1) + assert.Equal(t, "good.lua", plugins[0].name) + for _, p := range plugins { + p.close() + } +} diff --git a/internal/proxy/forward.go b/internal/proxy/forward.go new file mode 100644 index 0000000..7e40d8e --- /dev/null +++ b/internal/proxy/forward.go @@ -0,0 +1,29 @@ +package proxy + +import ( + "fmt" + "net/http" + + internalhttp "github.com/shiv/internal/http" +) + +func forward(req *http.Request, client *http.Client) (*http.Response, error) { + out, err := http.NewRequestWithContext(req.Context(), req.Method, req.URL.String(), req.Body) + if err != nil { + return nil, fmt.Errorf("forward: build request: %w", err) + } + for k, vals := range req.Header { + if internalhttp.HopByHop[k] { + continue + } + out.Header[k] = vals + } + internalhttp.StripRequestCacheHeaders(out.Header) + + resp, err := client.Do(out) + if err != nil { + return nil, fmt.Errorf("forward: upstream request: %w", err) + } + internalhttp.StripResponseCacheHeaders(resp.Header) + return resp, nil +} diff --git a/internal/proxy/forward_test.go b/internal/proxy/forward_test.go new file mode 100644 index 0000000..88d4db4 --- /dev/null +++ b/internal/proxy/forward_test.go @@ -0,0 +1,197 @@ +package proxy + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "net/http" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" + internalhttp "github.com/shiv/internal/http" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func decompressBody(h http.Header, body []byte) []byte { return internalhttp.Decompress(h, body) } +func isBinary(h http.Header) bool { return internalhttp.IsBinary(h) } +func stripRequestCacheHeaders(h http.Header) { internalhttp.StripRequestCacheHeaders(h) } +func stripResponseCacheHeaders(h http.Header) { internalhttp.StripResponseCacheHeaders(h) } + +// ── decompressBody ──────────────────────────────────────────────────────────── + +func TestDecompressBody_NoEncoding(t *testing.T) { + body := []byte("plain text") + out := decompressBody(http.Header{}, body) + assert.Equal(t, body, out) +} + +func TestDecompressBody_Gzip(t *testing.T) { + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write([]byte("gzip content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"gzip"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "gzip content", string(out)) +} + +func TestDecompressBody_Gzip_CorruptData(t *testing.T) { + body := []byte("not gzip") + h := http.Header{"Content-Encoding": []string{"gzip"}} + out := decompressBody(h, body) + assert.Equal(t, body, out, "corrupt gzip should return original bytes") +} + +func TestDecompressBody_Deflate(t *testing.T) { + var buf bytes.Buffer + w, err := flate.NewWriter(&buf, flate.DefaultCompression) + require.NoError(t, err) + _, err = w.Write([]byte("deflate content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"deflate"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "deflate content", string(out)) +} + +func TestDecompressBody_Brotli(t *testing.T) { + var buf bytes.Buffer + w := brotli.NewWriter(&buf) + _, err := w.Write([]byte("brotli content")) + require.NoError(t, err) + w.Close() + + h := http.Header{"Content-Encoding": []string{"br"}} + out := decompressBody(h, buf.Bytes()) + assert.Equal(t, "brotli content", string(out)) +} + +func TestDecompressBody_Zstd(t *testing.T) { + enc, err := zstd.NewWriter(nil) + require.NoError(t, err) + compressed := enc.EncodeAll([]byte("zstd content"), nil) + + h := http.Header{"Content-Encoding": []string{"zstd"}} + out := decompressBody(h, compressed) + assert.Equal(t, "zstd content", string(out)) +} + +func TestDecompressBody_UnknownEncoding(t *testing.T) { + body := []byte("mystery") + h := http.Header{"Content-Encoding": []string{"identity"}} + out := decompressBody(h, body) + assert.Equal(t, body, out) +} + +// ── isBinary ───────────────────────────────────────────────────────────────── + +func TestIsBinary_PlainText(t *testing.T) { + h := http.Header{"Content-Type": []string{"text/plain"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_HTML(t *testing.T) { + h := http.Header{"Content-Type": []string{"text/html; charset=utf-8"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_JSON(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/json"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_XML(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/xml"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_JavaScript(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/javascript"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_FormEncoded(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}} + assert.False(t, isBinary(h)) +} + +func TestIsBinary_Image(t *testing.T) { + h := http.Header{"Content-Type": []string{"image/png"}} + assert.True(t, isBinary(h)) +} + +func TestIsBinary_OctetStream(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/octet-stream"}} + assert.True(t, isBinary(h)) +} + +func TestIsBinary_NoContentType(t *testing.T) { + assert.False(t, isBinary(http.Header{})) +} + +func TestIsBinary_ZstdEncoding(t *testing.T) { + h := http.Header{ + "Content-Encoding": []string{"zstd"}, + "Content-Type": []string{"text/plain"}, + } + assert.False(t, isBinary(h)) +} + +// ── stripRequestCacheHeaders ────────────────────────────────────────────────── + +func TestStripRequestCacheHeaders(t *testing.T) { + h := http.Header{ + "If-None-Match": []string{`"etag123"`}, + "If-Modified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "If-Range": []string{`"etag123"`}, + "If-Match": []string{`"etag123"`}, + "If-Unmodified-Since": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Authorization": []string{"Bearer token"}, + } + stripRequestCacheHeaders(h) + + assert.Empty(t, h.Get("If-None-Match")) + assert.Empty(t, h.Get("If-Modified-Since")) + assert.Empty(t, h.Get("If-Range")) + assert.Empty(t, h.Get("If-Match")) + assert.Empty(t, h.Get("If-Unmodified-Since")) + assert.Equal(t, "Bearer token", h.Get("Authorization")) +} + +func TestStripRequestCacheHeaders_AlreadyAbsent(t *testing.T) { + h := http.Header{"Content-Type": []string{"application/json"}} + assert.NotPanics(t, func() { stripRequestCacheHeaders(h) }) + assert.Equal(t, "application/json", h.Get("Content-Type")) +} + +// ── stripResponseCacheHeaders ───────────────────────────────────────────────── + +func TestStripResponseCacheHeaders(t *testing.T) { + h := http.Header{ + "Cache-Control": []string{"max-age=3600"}, + "Expires": []string{"Mon, 01 Jan 2025 00:00:00 GMT"}, + "ETag": []string{`"etag123"`}, + "Last-Modified": []string{"Mon, 01 Jan 2024 00:00:00 GMT"}, + "Pragma": []string{"no-cache"}, + "Content-Type": []string{"text/html"}, + } + stripResponseCacheHeaders(h) + + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) + assert.Empty(t, h.Get("Expires")) + assert.Empty(t, h.Get("ETag")) + assert.Empty(t, h.Get("Last-Modified")) + assert.Empty(t, h.Get("Pragma")) + assert.Equal(t, "text/html", h.Get("Content-Type")) +} + +func TestStripResponseCacheHeaders_SetsNoCacheDirective(t *testing.T) { + h := http.Header{} + stripResponseCacheHeaders(h) + assert.Equal(t, "no-store, no-cache, must-revalidate", h.Get("Cache-Control")) +} diff --git a/internal/proxy/mitm.go b/internal/proxy/mitm.go new file mode 100644 index 0000000..5ed63a1 --- /dev/null +++ b/internal/proxy/mitm.go @@ -0,0 +1,214 @@ +package proxy + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "net" + "net/http" + "time" + + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) { + defer recoverPanic("ServeHTTP " + r.URL.String()) + w.WriteHeader(http.StatusOK) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + + hijacker, ok := w.(http.Hijacker) + if !ok { + logger.Error("mitm: server does not support hijacking") + return + } + clientConn, _, err := hijacker.Hijack() + if err != nil { + logger.Error("mitm: hijack: %v", err) + return + } + p.trackConn(clientConn) + defer p.untrackConn(clientConn) + defer clientConn.Close() + + bareHost, _, err := net.SplitHostPort(r.Host) + if err != nil { + bareHost = r.Host + } + + tlsCert, err := p.certAuth.TLSCertForHost(bareHost) + if err != nil { + logger.Error("mitm: cert for %s: %v", bareHost, err) + return + } + + browserTLS := tls.Server(clientConn, &tls.Config{ + Certificates: []tls.Certificate{*tlsCert}, + NextProtos: []string{"h2", "http/1.1"}, + }) + if err := browserTLS.Handshake(); err != nil { + logger.Debug("mitm: browser TLS handshake for %s: %v", bareHost, err) + return + } + defer browserTLS.Close() + logger.Debug("mitm: negotiated protocol for %s: %q", bareHost, browserTLS.ConnectionState().NegotiatedProtocol) + + if browserTLS.ConnectionState().NegotiatedProtocol == "h2" { + p.handleConnectH2(browserTLS, r, bareHost) + return + } + + transport := &http.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := (&net.Dialer{Timeout: 10 * time.Second}).DialContext(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("mitm: dial %s: %w", addr, err) + } + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + }) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("mitm: upstream TLS handshake for %s: %w", bareHost, err) + } + return tlsConn, nil + }, + ResponseHeaderTimeout: 30 * time.Second, + } + defer transport.CloseIdleConnections() + + // No client-level Timeout — a per-connection timeout would terminate the + // entire keep-alive session after the first long request. Individual + // requests are bounded by ResponseHeaderTimeout + read deadlines instead. + upstreamClient := &http.Client{ + Transport: transport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + + browserReader := bufio.NewReader(browserTLS) + for { + req, err := http.ReadRequest(browserReader) + if err != nil { + logger.Debug("mitm: connection closed for %s: %v", bareHost, err) + return + } + + start := time.Now() + + reqBody, err := io.ReadAll(io.LimitReader(req.Body, maxBodySize)) + if err != nil { + logger.Error("mitm: read request body for %s: %v", bareHost, err) + return + } + req.Body.Close() + + req.URL.Scheme = "https" + // Preserve r.Host (with port) so non-443 HTTPS targets work correctly. + req.URL.Host = r.Host + + // Detect WebSocket upgrade before intercept. WebSocket connections are + // long-lived bidirectional streams — they cannot be intercepted the same + // way as regular HTTP requests. Hand off to the WebSocket handler which + // manages the entire session and returns when the connection closes. + if isWebSocketUpgrade(req) { + p.handleWebSocketTLS(browserTLS, browserReader, req, bareHost, r.Host) + return + } + + result := p.bus.EmitRequest(events.RequestEvent{Request: req, Body: reqBody}) + if result.Drop { + resp := &http.Response{ + Status: "403 Forbidden", + StatusCode: 403, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Body: http.NoBody, + Header: make(http.Header), + ContentLength: 0, + } + if err := resp.Write(browserTLS); err != nil { + logger.Error("mitm: write 403 to browser for %s: %v", bareHost, err) + } + continue + } + interceptedReq, reqBody := result.Request, result.Body + + outReq, err := http.NewRequest(interceptedReq.Method, interceptedReq.URL.String(), bytes.NewReader(reqBody)) + if err != nil { + logger.Error("mitm: build upstream request for %s: %v", bareHost, err) + return + } + outReq.Header = interceptedReq.Header.Clone() + outReq.Host = internalhttp.NormalizeHost(bareHost, true) + internalhttp.StripRequestCacheHeaders(outReq.Header) + + resp, err := upstreamClient.Do(outReq) + if err != nil { + // Distinguish client-cancelled requests (normal) from real upstream + // failures. Context cancellation here means the browser closed the + // connection before we could forward — log at Debug. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.Debug("mitm: stream cancelled for %s: %v", bareHost, err) + } else { + logger.Error("mitm: upstream request for %s: %v", bareHost, err) + } + return + } + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + internalhttp.StripResponseCacheHeaders(resp.Header) + resp.Body = io.NopCloser(reader) + + writeErr := resp.Write(browserTLS) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + + if writeErr != nil { + logger.Debug("mitm: write response to browser for %s: %v", bareHost, writeErr) + return + } + + respBody := respBuf.Bytes() + + if len(respBody) >= int(maxBodySize) { + logger.Debug("mitm: response body truncated at %d bytes for %s %s", + maxBodySize, interceptedReq.Method, interceptedReq.URL) + } + + elapsed := time.Since(start).Milliseconds() + logger.Info("%s %s %d %db %dms", interceptedReq.Method, interceptedReq.URL, resp.StatusCode, len(respBody), elapsed) + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: r.Host, + Proto: "HTTP/1.1", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: true, + }) + } +} diff --git a/internal/proxy/mitm2.go b/internal/proxy/mitm2.go new file mode 100644 index 0000000..fcf8e0f --- /dev/null +++ b/internal/proxy/mitm2.go @@ -0,0 +1,211 @@ +package proxy + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "errors" + "fmt" + "io" + "net" + "net/http" + "time" + + "golang.org/x/net/http2" + + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +const ( + h2UpstreamDialTimeout = 10 * time.Second + h2UpstreamHeaderTimeout = 30 * time.Second + // h2StreamRequestTimeout is a per-stream deadline, not a connection-level + // timeout. Using a context timeout per stream instead of a client-level + // Timeout means a single slow stream cannot kill the entire H2 connection. + h2StreamRequestTimeout = 60 * time.Second + // h2MaxConcurrentStreams matches the HTTP/2 spec default (RFC 7540 §6.5.2). + h2MaxConcurrentStreams = 250 +) + +func (p *Proxy) handleConnectH2(browserTLS *tls.Conn, connectReq *http.Request, bareHost string) { + defer recoverPanic("handleConnectH2 " + bareHost) + + // Use the port-aware host from the original CONNECT request for upstream + // dialling so non-443 HTTPS ports are preserved. + upstreamHost := connectReq.Host + + upstreamTransport := &http.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := (&net.Dialer{Timeout: h2UpstreamDialTimeout}).DialContext(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("mitm/h2: dial %s: %w", addr, err) + } + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"h2", "http/1.1"}, + VerifyPeerCertificate: func(rawCerts [][]byte, _ [][]*x509.Certificate) error { + if len(rawCerts) == 0 { + return nil + } + cert, err := x509.ParseCertificate(rawCerts[0]) + if err != nil { + return nil + } + now := time.Now() + if now.After(cert.NotAfter) { + logger.Info("mitm/h2: EXPIRED cert for %s (expired %s, issuer: %s)", + bareHost, cert.NotAfter.Format("2006-01-02"), cert.Issuer.CommonName) + } else if now.Before(cert.NotBefore) { + logger.Info("mitm/h2: NOT YET VALID cert for %s (valid from %s, issuer: %s)", + bareHost, cert.NotBefore.Format("2006-01-02"), cert.Issuer.CommonName) + } + if err := cert.VerifyHostname(bareHost); err != nil { + logger.Info("mitm/h2: hostname mismatch for %s: %v (CN: %s, SANs: %v)", + bareHost, err, cert.Subject.CommonName, cert.DNSNames) + } + return nil + }, + }) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("mitm/h2: upstream TLS handshake for %s: %w", addr, err) + } + return tlsConn, nil + }, + ResponseHeaderTimeout: h2UpstreamHeaderTimeout, + } + if err := http2.ConfigureTransport(upstreamTransport); err != nil { + logger.Error("mitm/h2: configure upstream transport: %v", err) + } + defer upstreamTransport.CloseIdleConnections() + + // No client-level Timeout — that would kill the entire H2 connection on + // any slow stream. Per-stream timeouts are applied via request contexts. + upstreamClient := &http.Client{ + Transport: upstreamTransport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + + streamHandler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + defer recoverPanic("h2 stream " + bareHost) + + req.URL.Scheme = "https" + req.URL.Host = upstreamHost + + reqBody, err := io.ReadAll(io.LimitReader(req.Body, maxBodySize)) + if err != nil { + logger.Error("mitm/h2: read request body for %s: %v", bareHost, err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + req.Body.Close() + + start := time.Now() + + result := p.bus.EmitRequest(events.RequestEvent{Request: req, Body: reqBody}) + if result.Drop { + // Stream is dropped by intercept rules; the H2 connection stays alive. + http.Error(w, "request dropped", http.StatusForbidden) + return + } + interceptedReq, reqBody := result.Request, result.Body + + streamCtx, cancel := context.WithTimeout(req.Context(), h2StreamRequestTimeout) + defer cancel() + + outReq, err := http.NewRequestWithContext(streamCtx, interceptedReq.Method, interceptedReq.URL.String(), bytes.NewReader(reqBody)) + if err != nil { + logger.Error("mitm/h2: build upstream request for %s: %v", bareHost, err) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + outReq.Header = interceptedReq.Header.Clone() + outReq.Host = internalhttp.NormalizeHost(bareHost, true) + // Strip H2-illegal headers before forwarding. + outReq.Header.Del("Transfer-Encoding") + internalhttp.StripRequestCacheHeaders(outReq.Header) + + resp, err := upstreamClient.Do(outReq) + if err != nil { + // context.Canceled means the browser cancelled the stream (RST_STREAM) + // or we hit our per-stream deadline. This is normal browser behaviour + // (prefetch, autocomplete, navigation away) — log at Debug, not Error. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.Debug("mitm/h2: stream cancelled/timed out for %s: %v", bareHost, err) + } else { + logger.Error("mitm/h2: upstream request for %s: %v", bareHost, err) + } + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + // Strip Transfer-Encoding — illegal in HTTP/2 responses (RFC 7540 §8.1.2.2). + // Strip caching headers so the browser always fetches fresh content. + resp.Header.Del("Transfer-Encoding") + internalhttp.StripResponseCacheHeaders(resp.Header) + + for k, vals := range resp.Header { + for _, v := range vals { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + + if _, err := io.Copy(w, reader); err != nil { + logger.Debug("mitm/h2: write response to browser for %s: %v", bareHost, err) + // Drain whatever remains into respBuf so the log body is as complete + // as possible even if the browser disconnected mid-stream. + _, _ = io.Copy(&respBuf, resp.Body) + } + + respBody := respBuf.Bytes() + resp.Body.Close() + + if len(respBody) >= int(maxBodySize) { + logger.Debug("mitm/h2: response body truncated at %d bytes for %s %s", + maxBodySize, interceptedReq.Method, interceptedReq.URL) + } + + elapsed := time.Since(start).Milliseconds() + logger.Info("%s %s %d %db %dms (h2)", interceptedReq.Method, interceptedReq.URL, resp.StatusCode, len(respBody), elapsed) + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: connectReq.Host, + Proto: "HTTP/2", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: true, + }) + }) + + h2srv := &http2.Server{ + MaxConcurrentStreams: h2MaxConcurrentStreams, + } + // ServeConn blocks until the browser closes the underlying TLS connection. + // It handles all stream multiplexing internally. No idle timer is needed — + // the connection lifetime is naturally tied to browserTLS. + h2srv.ServeConn(browserTLS, &http2.ServeConnOpts{ + Handler: streamHandler, + }) +} diff --git a/internal/proxy/mitm2_test.go b/internal/proxy/mitm2_test.go new file mode 100644 index 0000000..a1d4c40 --- /dev/null +++ b/internal/proxy/mitm2_test.go @@ -0,0 +1,280 @@ +package proxy_test + +import ( + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "golang.org/x/net/http2" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// h2UpstreamServer starts an HTTP/2-enabled TLS test server. +func h2UpstreamServer(t *testing.T, handler http.Handler) *httptest.Server { + t.Helper() + srv := httptest.NewUnstartedServer(handler) + srv.TLS = &tls.Config{ + NextProtos: []string{"h2", "http/1.1"}, + } + h2s := &http2.Server{} + require.NoError(t, http2.ConfigureServer(srv.Config, h2s)) + srv.StartTLS() + t.Cleanup(srv.Close) + return srv +} + +// h2ClientOverTunnel wraps a raw conn in a TLS client that negotiates HTTP/2 +// and returns an *http.Client that routes all requests over that connection. +func h2ClientOverTunnel(t *testing.T, conn net.Conn, serverName string) *http.Client { + t.Helper() + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: serverName, + InsecureSkipVerify: true, //nolint:gosec — test only + NextProtos: []string{"h2"}, + }) + require.NoError(t, tlsConn.Handshake()) + assert.Equal(t, "h2", tlsConn.ConnectionState().NegotiatedProtocol, "proxy must negotiate h2 with the client") + t.Cleanup(func() { tlsConn.Close() }) + + transport := &http2.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) { + return tlsConn, nil + }, + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec + } + + return &http.Client{ + Transport: transport, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } +} + +// ── tests ───────────────────────────────────────────────────────────────────── + +func TestMITM_H2_ForwardsGET(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + io.WriteString(w, "h2 works") + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + b, _ := io.ReadAll(resp.Body) + assert.Equal(t, "h2 works", string(b)) +} + +func TestMITM_H2_LogsTransaction(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + io.WriteString(w, "logged") + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/log-test") + require.NoError(t, err) + resp.Body.Close() + + tx := waitForTransaction(t, st) + assert.True(t, tx.TLS) + assert.Equal(t, "GET", tx.Method) + assert.Equal(t, http.StatusOK, tx.StatusCode) +} + +func TestMITM_H2_PostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + body := `{"key":"value"}` + req, err := http.NewRequest(http.MethodPost, "https://"+targetHost+"/api", strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode) + assert.Equal(t, body, receivedBody) +} + +func TestMITM_H2_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/") + require.NoError(t, err) + resp.Body.Close() + + assert.Empty(t, resp.Header.Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", resp.Header.Get("Cache-Control")) +} + +func TestMITM_H2_InterceptDrop_Returns403(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.WriteString(w, "should not reach") + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/secret") + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestMITM_H2_InterceptForward_ModifiesRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "h2-injected") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + resp, err := client.Get("https://" + targetHost + "/inject") + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "h2-injected", receivedHeader) +} + +func TestMITM_H2_MultipleConcurrentStreams(t *testing.T) { + st := newTestStore(t) + + upstream := h2UpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, "path:%s", r.URL.Path) + })) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + client := h2ClientOverTunnel(t, conn, host) + + results := make(chan string, 3) + for _, path := range []string{"/a", "/b", "/c"} { + p := path + go func() { + resp, err := client.Get("https://" + targetHost + p) + if err != nil { + results <- "error" + return + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + results <- string(b) + }() + } + + got := make(map[string]bool) + for i := 0; i < 3; i++ { + got[<-results] = true + } + assert.True(t, got["path:/a"]) + assert.True(t, got["path:/b"]) + assert.True(t, got["path:/c"]) +} diff --git a/internal/proxy/mitm_test.go b/internal/proxy/mitm_test.go new file mode 100644 index 0000000..c683c70 --- /dev/null +++ b/internal/proxy/mitm_test.go @@ -0,0 +1,281 @@ +package proxy_test + +import ( + "bufio" + "bytes" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// startProxyServer starts the proxy as a real TCP listener and returns its address. +// httptest.NewRecorder does not support hijacking so we need a real listener for CONNECT. +func startProxyServer(t *testing.T, st *store.Store) string { + t.Helper() + bus := newTestBus(t, st) + p := newTestProxy(t, bus) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + srv := &http.Server{Handler: p} + go srv.Serve(ln) + t.Cleanup(func() { srv.Close() }) + + return ln.Addr().String() +} + +// tlsUpstream starts an httptest TLS server. +func tlsUpstream(t *testing.T, status int, body string) *httptest.Server { + t.Helper() + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + io.WriteString(w, body) + })) + t.Cleanup(srv.Close) + return srv +} + +// doConnectTunnel sends CONNECT to the proxy and returns the raw conn after the 200 response. +func doConnectTunnel(t *testing.T, proxyAddr, targetHost string) net.Conn { + t.Helper() + conn, err := net.DialTimeout("tcp", proxyAddr, 5*time.Second) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + fmt.Fprintf(conn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", targetHost, targetHost) + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "proxy must respond 200 to CONNECT") + + return conn +} + +// tlsOverTunnel wraps conn in a TLS client trusting the proxy's self-signed cert. +func tlsOverTunnel(t *testing.T, conn net.Conn, serverName string) *tls.Conn { + t.Helper() + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: serverName, + InsecureSkipVerify: true, //nolint:gosec + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + return tlsConn +} + +// sendHTTPOverConn writes a raw HTTP/1.1 request and reads the response. +func sendHTTPOverConn(t *testing.T, conn io.ReadWriter, method, path, host, body string) *http.Response { + t.Helper() + var req strings.Builder + fmt.Fprintf(&req, "%s %s HTTP/1.1\r\n", method, path) + fmt.Fprintf(&req, "Host: %s\r\n", host) + if body != "" { + fmt.Fprintf(&req, "Content-Length: %d\r\n", len(body)) + } + req.WriteString("Connection: close\r\n\r\n") + if body != "" { + req.WriteString(body) + } + + _, err := io.WriteString(conn, req.String()) + require.NoError(t, err) + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + return resp +} + +// waitForTransaction blocks until a transaction appears on st.Updates or times out. +func waitForTransaction(t *testing.T, st *store.Store) store.Transaction { + t.Helper() + select { + case tx := <-st.Updates: + return tx + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for transaction to be logged") + return store.Transaction{} + } +} + +// ── tests ───────────────────────────────────────────────────────────────────── + +func TestMITM_ConnectTunnel_ForwardsRequest(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "mitm works") + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/", targetHost, "") + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + b, _ := io.ReadAll(resp.Body) + assert.Equal(t, "mitm works", string(b)) +} + +func TestMITM_ConnectTunnel_LogsTransactionWithTLSTrue(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "logged") + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/log-test", targetHost, "") + resp.Body.Close() + + tx := waitForTransaction(t, st) + assert.True(t, tx.TLS, "transaction must be flagged as TLS") + assert.Equal(t, "GET", tx.Method) + assert.Equal(t, http.StatusOK, tx.StatusCode) +} + +func TestMITM_ConnectTunnel_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/", targetHost, "") + resp.Body.Close() + + assert.Empty(t, resp.Header.Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", resp.Header.Get("Cache-Control")) +} + +func TestMITM_ConnectTunnel_InterceptDrop_Returns403(t *testing.T) { + st := newTestStore(t) + upstream := tlsUpstream(t, http.StatusOK, "should not reach") + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/secret", targetHost, "") + resp.Body.Close() + + assert.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestMITM_ConnectTunnel_InterceptForward_ModifiesRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + st.Intercept.SetEnabled(true) + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "mitm-injected") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + resp := sendHTTPOverConn(t, tlsConn, "GET", "/inject", targetHost, "") + resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "mitm-injected", receivedHeader) +} + +func TestMITM_ConnectTunnel_PostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + t.Cleanup(upstream.Close) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tlsOverTunnel(t, conn, host) + + body := `{"key":"value"}` + var req bytes.Buffer + fmt.Fprintf(&req, "POST /api HTTP/1.1\r\n") + fmt.Fprintf(&req, "Host: %s\r\n", targetHost) + fmt.Fprintf(&req, "Content-Type: application/json\r\n") + fmt.Fprintf(&req, "Content-Length: %d\r\n", len(body)) + fmt.Fprintf(&req, "Connection: close\r\n\r\n") + fmt.Fprintf(&req, "%s", body) + _, err = tlsConn.Write(req.Bytes()) + require.NoError(t, err) + + resp, err := http.ReadResponse(bufio.NewReader(tlsConn), nil) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode) + assert.Equal(t, body, receivedBody) +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go new file mode 100644 index 0000000..3e9845a --- /dev/null +++ b/internal/proxy/proxy.go @@ -0,0 +1,200 @@ +package proxy + +import ( + "bytes" + "context" + "fmt" + "io" + "net" + "net/http" + "sync" + "time" + + "github.com/shiv/internal/cert" + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" +) + +type Proxy struct { + addr string + certAuth *cert.CA + bus *events.Bus + client *http.Client + + mu sync.Mutex + srv *http.Server + conns sync.Map // active hijacked net.Conn +} + +const maxBodySize = 10 << 20 + +func New(addr string, bus *events.Bus) (*Proxy, error) { + ca, err := cert.Load() + if err != nil { + return nil, fmt.Errorf("proxy: load CA: %w", err) + } + p := &Proxy{ + addr: addr, + certAuth: ca, + bus: bus, + client: &http.Client{ + Transport: &http.Transport{ + ResponseHeaderTimeout: 30 * time.Second, + }, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + }, + } + bus.Register(p) + return p, nil +} + +func (p *Proxy) CA() *cert.CA { + return p.certAuth +} + +func (p *Proxy) trackConn(c net.Conn) { + p.conns.Store(c, struct{}{}) +} + +func (p *Proxy) untrackConn(c net.Conn) { + p.conns.Delete(c) +} + +func (p *Proxy) closeAllConns() { + p.conns.Range(func(k, _ any) bool { + k.(net.Conn).SetDeadline(time.Now()) + return true + }) +} + +func (p *Proxy) Start(newAddr string) error { + ln, err := net.Listen("tcp", newAddr) + if err != nil { + return fmt.Errorf("proxy: listen on %s: %w", newAddr, err) + } + + newSrv := &http.Server{ + Handler: p, + ReadTimeout: 60 * time.Second, + WriteTimeout: 60 * time.Second, + IdleTimeout: 120 * time.Second, + } + + p.mu.Lock() + oldSrv := p.srv + p.addr = newAddr + p.srv = newSrv + p.mu.Unlock() + + if oldSrv != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = oldSrv.Shutdown(ctx) + } + p.closeAllConns() + + go func() { + if err := newSrv.Serve(ln); err != nil && err != http.ErrServerClosed { + logger.Error("proxy: Serve: %v", err) + } + }() + + return nil +} + +func (p *Proxy) Stop() { + p.mu.Lock() + srv := p.srv + p.mu.Unlock() + + if srv != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = srv.Shutdown(ctx) + + p.mu.Lock() + p.srv = nil + p.mu.Unlock() + } + + p.closeAllConns() +} + +// ObserveProxyRestart implements events.ProxyCommandObserver. +func (p *Proxy) ObserveProxyRestart(e events.ProxyRestartEvent) { + _ = p.Start(e.Addr) +} + +// ObserveProxyStop implements events.ProxyCommandObserver. +func (p *Proxy) ObserveProxyStop(_ events.ProxyStopEvent) { + p.Stop() +} + +func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodConnect { + p.handleConnect(w, r) + return + } + + reqBody, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize)) + if err != nil { + http.Error(w, "failed to read request body", http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(reqBody)) + + result := p.bus.EmitRequest(events.RequestEvent{Request: r, Body: reqBody}) + if result.Drop { + http.Error(w, "request dropped", http.StatusForbidden) + return + } + interceptedReq, reqBody := result.Request, result.Body + interceptedReq.Body = io.NopCloser(bytes.NewReader(reqBody)) + + start := time.Now() + + resp, err := forward(interceptedReq, p.client) + if err != nil { + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + defer resp.Body.Close() + + var respBuf bytes.Buffer + reader := io.TeeReader(io.LimitReader(resp.Body, maxBodySize), &respBuf) + + for k, vals := range resp.Header { + for _, v := range vals { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + + _, _ = io.Copy(w, reader) + + respBody := respBuf.Bytes() + elapsed := time.Since(start).Milliseconds() + + logBody := internalhttp.Decompress(resp.Header, respBody) + if internalhttp.IsBinary(resp.Header) { + logBody = nil + } + + p.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: interceptedReq.Host, + Proto: "HTTP/1.1", + Method: interceptedReq.Method, + URL: interceptedReq.URL.String(), + ReqHeaders: interceptedReq.Header, + ReqBody: reqBody, + StatusCode: resp.StatusCode, + RespHeaders: resp.Header, + RespBody: logBody, + DurationMs: elapsed, + TLS: false, + }) +} diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go new file mode 100644 index 0000000..ab5d69e --- /dev/null +++ b/internal/proxy/proxy_test.go @@ -0,0 +1,218 @@ +package proxy_test + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/proxy" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestStore opens a temp SQLite store and registers cleanup. +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-proxy-test-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + projectStore, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { projectStore.Close() }) + return projectStore +} + +// newTestBus creates a bus wired to the given store — mirrors production wiring. +func newTestBus(t *testing.T, st *store.Store) *events.Bus { + t.Helper() + bus := events.NewBus() + bus.Register(st.Intercept) + bus.Register(st) + return bus +} + +// newTestProxy creates a Proxy wired to an events.Bus. +// The proxy is wired as an http.Handler so we can call it via httptest without +// binding a real port. +func newTestProxy(t *testing.T, bus *events.Bus) *proxy.Proxy { + t.Helper() + p, err := proxy.New("127.0.0.1:0", bus) + require.NoError(t, err) + return p +} + +// upstreamServer starts a simple upstream that echoes status + body. +func upstreamServer(t *testing.T, status int, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + io.WriteString(w, body) + })) + t.Cleanup(srv.Close) + return srv +} + +// ── HTTP forwarding ─────────────────────────────────────────────────────────── + +func TestServeHTTP_ForwardsRequestAndReturnsResponse(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "hello proxy") + p := newTestProxy(t, newTestBus(t, st)) + + req := httptest.NewRequest(http.MethodGet, upstream.URL+"/test", nil) + req.RequestURI = upstream.URL + "/test" + w := httptest.NewRecorder() + + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "hello proxy", w.Body.String()) +} + +func TestServeHTTP_ForwardsPostWithBody(t *testing.T) { + st := newTestStore(t) + + var receivedBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + receivedBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + body := `{"name":"test"}` + req := httptest.NewRequest(http.MethodPost, upstream.URL+"/api", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusCreated, w.Code) + assert.Equal(t, body, receivedBody) +} + +func TestServeHTTP_LogsTransactionToStore(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "logged") + p := newTestProxy(t, newTestBus(t, st)) + + req := httptest.NewRequest(http.MethodGet, upstream.URL+"/logged", nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + // Give the store a moment to write (it's async via channel) + time.Sleep(50 * time.Millisecond) + + txs, err := st.TransactionsPage(0, store.TransactionFilter{ShowOutScope: true}) + require.NoError(t, err) + require.Len(t, txs, 1) + assert.Equal(t, "GET", txs[0].Method) + assert.Equal(t, http.StatusOK, txs[0].StatusCode) +} + +func TestServeHTTP_ResponseHeadersPassedThrough(t *testing.T) { + st := newTestStore(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Custom-Header", "custom-value") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, "custom-value", w.Header().Get("X-Custom-Header")) +} + +func TestServeHTTP_CacheHeadersStripped(t *testing.T) { + st := newTestStore(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"abc123"`) + w.Header().Set("Cache-Control", "max-age=3600") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Empty(t, w.Header().Get("ETag")) + assert.Equal(t, "no-store, no-cache, must-revalidate", w.Header().Get("Cache-Control")) +} + +func TestServeHTTP_UpstreamError_ReturnsBadGateway(t *testing.T) { + st := newTestStore(t) + p := newTestProxy(t, newTestBus(t, st)) + + // Point at a port that is not listening + req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1:1/", nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) +} + +// ── Intercept gate ──────────────────────────────────────────────────────────── + +func TestServeHTTP_InterceptDrop_ReturnsForbidden(t *testing.T) { + st := newTestStore(t) + upstream := upstreamServer(t, http.StatusOK, "should not reach") + p := newTestProxy(t, newTestBus(t, st)) + + st.Intercept.SetEnabled(true) + + // Drain the intercept queue and drop the request in a goroutine + go func() { + pending := <-st.Intercept.Queue() + pending.Reply <- store.Decision{Forward: false} + }() + + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusForbidden, w.Code) +} + +func TestServeHTTP_InterceptForward_ModifiedRequest(t *testing.T) { + st := newTestStore(t) + + var receivedHeader string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Injected") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + p := newTestProxy(t, newTestBus(t, st)) + st.Intercept.SetEnabled(true) + + go func() { + pending := <-st.Intercept.Queue() + pending.Request.Header.Set("X-Injected", "intercepted") + pending.Reply <- store.Decision{ + Forward: true, + Request: pending.Request, + Body: pending.Body, + } + }() + + req := httptest.NewRequest(http.MethodGet, upstream.URL, nil) + w := httptest.NewRecorder() + p.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "intercepted", receivedHeader) +} diff --git a/internal/proxy/recover.go b/internal/proxy/recover.go new file mode 100644 index 0000000..27338e9 --- /dev/null +++ b/internal/proxy/recover.go @@ -0,0 +1,13 @@ +package proxy + +import ( + "runtime/debug" + + "github.com/shiv/internal/logger" +) + +func recoverPanic(context string) { + if r := recover(); r != nil { + logger.Error("panic in %s: %v\n%s", context, r, debug.Stack()) + } +} diff --git a/internal/proxy/websocket.go b/internal/proxy/websocket.go new file mode 100644 index 0000000..b02470e --- /dev/null +++ b/internal/proxy/websocket.go @@ -0,0 +1,225 @@ +package proxy + +import ( + "bufio" + "crypto/tls" + "fmt" + "net" + "net/http" + "net/url" + "time" + + "github.com/gorilla/websocket" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +// isWebSocketUpgrade returns true if the request is a WebSocket upgrade. +func isWebSocketUpgrade(req *http.Request) bool { + return websocket.IsWebSocketUpgrade(req) +} + +// handleWebSocketTLS handles a WebSocket upgrade over an already-established +// TLS connection. +// +// browserReader is the *bufio.Reader that was used to parse the upgrade +// request in the H1 loop. It must be passed here — not discarded — because +// bufio.Reader reads ahead and may have buffered bytes beyond the end of the +// HTTP request. Discarding it and creating a new bufio.Reader over the raw +// conn would silently lose those bytes, corrupting the WebSocket stream. +// +// Design: +// 1. Dial the upstream as a WebSocket client. +// 2. Give gorilla's Upgrader a hijackWriter that wraps (tlsConn, browserReader) +// so gorilla writes the 101 and takes back the conn via Hijack() — +// the read side of Hijack uses the existing browserReader, preserving +// any buffered bytes. Zero replay. Zero data loss. +// 3. Proxy frames bidirectionally with gorilla on both sides. +func (p *Proxy) handleWebSocketTLS( + tlsConn *tls.Conn, + browserReader *bufio.Reader, + req *http.Request, + bareHost string, + hostWithPort string, +) { + defer recoverPanic("websocket " + bareHost) + + upstreamURL := &url.URL{ + Scheme: "wss", + Host: hostWithPort, + Path: req.URL.Path, + RawQuery: req.URL.RawQuery, + } + + // ── 1. Dial upstream ────────────────────────────────────────────────────── + + upstreamHeaders := http.Header{} + for k, vals := range req.Header { + switch k { + case "Upgrade", "Connection", "Sec-Websocket-Key", + "Sec-Websocket-Version", "Sec-Websocket-Extensions", + "Sec-Websocket-Protocol": + // Let gorilla's dialer handle these. + default: + upstreamHeaders[k] = vals + } + } + + dialer := websocket.Dialer{ + TLSClientConfig: &tls.Config{ + ServerName: bareHost, + InsecureSkipVerify: true, //nolint:gosec + }, + HandshakeTimeout: 10 * time.Second, + } + + upstreamConn, upstreamResp, err := dialer.Dial(upstreamURL.String(), upstreamHeaders) + if err != nil { + logger.Error("ws: dial upstream %s: %v", upstreamURL, err) + writeWSError(tlsConn, http.StatusBadGateway) + return + } + defer upstreamConn.Close() + + // ── 2. Upgrade browser conn ─────────────────────────────────────────────── + // + // hijackWriter implements http.ResponseWriter + http.Hijacker over + // (tlsConn, browserReader). When gorilla's Upgrader calls Hijack(), it + // gets back the existing bufio.ReadWriter — the read side IS browserReader, + // so any bytes already buffered by the H1 loop are not lost. + + responseHeader := http.Header{} + if proto := upstreamResp.Header.Get("Sec-Websocket-Protocol"); proto != "" { + responseHeader.Set("Sec-WebSocket-Protocol", proto) + } + + hw := newHijackWriter(tlsConn, browserReader) + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + ReadBufferSize: 1 << 20, + WriteBufferSize: 1 << 20, + } + browserConn, err := upgrader.Upgrade(hw, req, responseHeader) + if err != nil { + logger.Error("ws: upgrade browser for %s: %v", bareHost, err) + return + } + defer browserConn.Close() + const wsMaxMessageSize = 16 << 20 // 16 MB + browserConn.SetReadLimit(wsMaxMessageSize) + upstreamConn.SetReadLimit(wsMaxMessageSize) + + // ── 3. Emit connection event ────────────────────────────────────────────── + + connID := p.bus.EmitWebSocketConnection(events.WebSocketConnectionEvent{ + Host: hostWithPort, + URL: upstreamURL.String(), + TLS: true, + Timestamp: time.Now(), + }) + + logger.Info("ws: connected %s%s", bareHost, req.URL.Path) + + // ── 4. Bidirectional frame proxy ────────────────────────────────────────── + + done := make(chan struct{}, 2) + + go func() { + defer func() { done <- struct{}{} }() + for { + msgType, payload, err := browserConn.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.Debug("ws: browser→upstream read for %s: %v", bareHost, err) + } + upstreamConn.Close() + return + } + result := p.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketOpcode(msgType), + Payload: payload, + }) + if err := upstreamConn.WriteMessage(msgType, result.Payload); err != nil { + logger.Debug("ws: browser→upstream write for %s: %v", bareHost, err) + return + } + } + }() + + go func() { + defer func() { done <- struct{}{} }() + for { + msgType, payload, err := upstreamConn.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.Debug("ws: upstream→browser read for %s: %v", bareHost, err) + } + browserConn.Close() + return + } + result := p.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: connID, + Timestamp: time.Now(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketOpcode(msgType), + Payload: payload, + }) + if err := browserConn.WriteMessage(msgType, result.Payload); err != nil { + logger.Debug("ws: upstream→browser write for %s: %v", bareHost, err) + return + } + } + }() + + <-done + <-done + logger.Info("ws: closed %s%s", bareHost, req.URL.Path) +} + +// writeWSError writes a minimal HTTP error response directly to a raw conn. +// Used when the upstream dial fails before the 101 is sent. +func writeWSError(conn net.Conn, code int) { + fmt.Fprintf(conn, "HTTP/1.1 %d %s\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + code, http.StatusText(code)) +} + +// hijackWriter implements http.ResponseWriter and http.Hijacker over a raw +// net.Conn + an existing *bufio.Reader. +// +// The critical invariant: the bufio.Reader passed here MUST be the same one +// used to parse the HTTP upgrade request. gorilla's Upgrader calls Hijack() +// to take ownership of the connection — the read side of the returned +// bufio.ReadWriter is this same reader, so any bytes already buffered in it +// (read-ahead past the end of the HTTP request) are not lost. +// +// The write side is a fresh bufio.Writer over the raw conn — gorilla flushes +// it after writing the 101 response. +type hijackWriter struct { + conn net.Conn + reader *bufio.Reader + header http.Header +} + +func newHijackWriter(conn net.Conn, reader *bufio.Reader) *hijackWriter { + return &hijackWriter{ + conn: conn, + reader: reader, + header: make(http.Header), + } +} + +func (hw *hijackWriter) Header() http.Header { return hw.header } +func (hw *hijackWriter) WriteHeader(_ int) {} +func (hw *hijackWriter) Write(b []byte) (int, error) { return hw.conn.Write(b) } + +// Hijack returns the raw conn and a bufio.ReadWriter whose read side is the +// existing browserReader — preserving any buffered bytes — and whose write +// side is a fresh writer over the conn. +func (hw *hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + brw := bufio.NewReadWriter(hw.reader, bufio.NewWriter(hw.conn)) + return hw.conn, brw, nil +} diff --git a/internal/proxy/websocket_test.go b/internal/proxy/websocket_test.go new file mode 100644 index 0000000..6ea6024 --- /dev/null +++ b/internal/proxy/websocket_test.go @@ -0,0 +1,764 @@ +package proxy_test + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── helpers ─────────────────────────────────────────────────────────────────── + +// wsUpstreamServer starts a TLS WebSocket echo server. +func wsUpstreamServer(t *testing.T) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + msgType, payload, err := conn.ReadMessage() + if err != nil { + return + } + if err := conn.WriteMessage(msgType, payload); err != nil { + return + } + } + })) + t.Cleanup(srv.Close) + return srv +} + +// wsUpstreamServerWithHandler starts a TLS WebSocket server with a custom handler. +func wsUpstreamServerWithHandler(t *testing.T, handler func(*websocket.Conn)) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + handler(conn) + })) + t.Cleanup(srv.Close) + return srv +} + +// wsClientOverTunnel dials a WebSocket through the proxy CONNECT tunnel. +func wsClientOverTunnel(t *testing.T, proxyAddr, targetHost string) *websocket.Conn { + t.Helper() + conn := doConnectTunnel(t, proxyAddr, targetHost) + + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: strings.Split(targetHost, ":")[0], + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"http/1.1"}, + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + + dialer := websocket.Dialer{ + NetDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return tlsConn, nil + }, + } + wsConn, _, err := dialer.Dial("wss://"+targetHost+"/", nil) + require.NoError(t, err) + t.Cleanup(func() { wsConn.Close() }) + return wsConn +} + +// sendHTTPKeepAlive writes an HTTP/1.1 request with Connection: keep-alive +// and reads the response, leaving the connection open for reuse. +func sendHTTPKeepAlive(t *testing.T, conn io.ReadWriter, method, path, host string) *http.Response { + t.Helper() + var req strings.Builder + fmt.Fprintf(&req, "%s %s HTTP/1.1\r\n", method, path) + fmt.Fprintf(&req, "Host: %s\r\n", host) + req.WriteString("Connection: keep-alive\r\n\r\n") + _, err := io.WriteString(conn, req.String()) + require.NoError(t, err) + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + return resp +} +func wsClientOverTunnelWithPath(t *testing.T, proxyAddr, targetHost, path string) *websocket.Conn { + t.Helper() + conn := doConnectTunnel(t, proxyAddr, targetHost) + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: strings.Split(targetHost, ":")[0], + InsecureSkipVerify: true, //nolint:gosec + NextProtos: []string{"http/1.1"}, + }) + require.NoError(t, tlsConn.Handshake()) + t.Cleanup(func() { tlsConn.Close() }) + + dialer := websocket.Dialer{ + NetDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return tlsConn, nil + }, + } + wsConn, _, err := dialer.Dial("wss://"+targetHost+path, nil) + require.NoError(t, err) + t.Cleanup(func() { wsConn.Close() }) + return wsConn +} + +// waitForFrames polls until the store has at least n frames for the first +// connection, or fails after 3 seconds. +func waitForFrames(t *testing.T, st interface { + AllWebSocketConnections() (interface{ Len() int }, error) +}, n int) { + t.Helper() + // implemented inline in each test for flexibility +} + +// ── basic correctness ───────────────────────────────────────────────────────── + +func TestMITM_WebSocket_EchoesMessages(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("hello"))) + msgType, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.TextMessage, msgType) + assert.Equal(t, "hello", string(payload)) +} + +func TestMITM_WebSocket_LogsConnection(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("ping"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + loaded, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(loaded) > 0 { + assert.Equal(t, targetHost, loaded[0].Host) + assert.True(t, loaded[0].TLS) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for WebSocket connection to be logged") +} + +func TestMITM_WebSocket_LogsFrames(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + for _, msg := range []string{"frame1", "frame2"} { + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(msg))) + _, _, err := wsConn.ReadMessage() + require.NoError(t, err) + } + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 4 { + assert.Equal(t, 4, len(frames)) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for WebSocket frames to be logged") +} + +// ── hijackWriter / gorilla coupling canary ──────────────────────────────────── + +// TestMITM_WebSocket_HijackCalledByUpgrader verifies that gorilla's Upgrader +// calls Hijack() on our hijackWriter. If gorilla ever changes this internal +// behaviour the test will fail before it hits production. +func TestMITM_WebSocket_HijackCalledByUpgrader(t *testing.T) { + hijackCalled := false + + // trackingWriter wraps a real http.ResponseWriter and records Hijack calls. + type trackingWriter struct { + http.ResponseWriter + conn net.Conn + brw *bufio.ReadWriter + called *bool + } + tw := &trackingWriter{} + tw.called = &hijackCalled + + // We test this via a real proxy round-trip — if the WebSocket connection + // succeeds, Hijack() was called. If gorilla stops calling Hijack(), the + // upgrade fails and the test fails. + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + // If the connection succeeds, Hijack was called correctly. + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("hijack-test"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "hijack-test", string(payload), "gorilla Upgrader must use Hijack to take over the connection") +} + +// ── buffered reader integrity ───────────────────────────────────────────────── + +// TestMITM_WebSocket_BufferedReaderNotLost verifies that the proxy correctly +// uses the existing bufio.Reader when handing off to the WebSocket handler. +// This is tested indirectly: if browserReader were discarded and a new one +// created over the raw conn, the WS framing would desync on high-traffic +// connections. We verify correctness by sending many rapid messages immediately +// after connection, which exercises the read buffer aggressively. +func TestMITM_WebSocket_BufferedReaderNotLost(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // Send 50 messages immediately without waiting for echo. + // If the bufio.Reader is recreated (losing buffered bytes), framing + // desyncs and ReadMessage returns an error. + const n = 50 + done := make(chan error, 1) + go func() { + for i := 0; i < n; i++ { + if err := wsConn.WriteMessage(websocket.TextMessage, []byte("buf")); err != nil { + done <- err + return + } + } + done <- nil + }() + + wsConn.SetReadDeadline(time.Now().Add(10 * time.Second)) + for i := 0; i < n; i++ { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err, "frame %d: bufio.Reader must not lose bytes", i) + assert.Equal(t, "buf", string(payload)) + } + require.NoError(t, <-done) +} + +// ── message types ───────────────────────────────────────────────────────────── + +func TestMITM_WebSocket_BinaryMessage(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + payload := []byte{0x00, 0x01, 0x02, 0xFE, 0xFF} + require.NoError(t, wsConn.WriteMessage(websocket.BinaryMessage, payload)) + msgType, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.BinaryMessage, msgType) + assert.Equal(t, payload, received) +} + +func TestMITM_WebSocket_EmptyPayload(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(""))) + msgType, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.TextMessage, msgType) + assert.Equal(t, []byte{}, received) +} + +func TestMITM_WebSocket_LargePayload(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // 512KB — well under wsFrameSizeLimit but large enough to exercise buffering. + payload := bytes.Repeat([]byte("A"), 512*1024) + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, payload)) + _, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, payload, received) +} + +func TestMITM_WebSocket_PayloadAtSizeLimit(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + // Exactly at wsFrameSizeLimit (1MB) — should proxy fine, and logged payload + // should be exactly wsFrameSizeLimit bytes (no truncation needed). + payload := bytes.Repeat([]byte("B"), 1<<20) + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, payload)) + _, received, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, len(payload), len(received), "payload at size limit must be proxied intact") +} + +// ── message ordering ────────────────────────────────────────────────────────── + +func TestMITM_WebSocket_MessageOrderingPreserved(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + messages := []string{"first", "second", "third", "fourth", "fifth"} + for _, msg := range messages { + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte(msg))) + } + for _, expected := range messages { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, expected, string(payload), "messages must arrive in send order") + } +} + +// ── frame logging correctness ───────────────────────────────────────────────── + +func TestMITM_WebSocket_FrameDirectionsCorrect(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("dirtest"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 2 { + // First frame must be client→server, second server→client. + assert.Equal(t, 0, int(frames[0].Direction), "first frame must be client→server") + assert.Equal(t, 1, int(frames[1].Direction), "second frame must be server→client") + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for frames with correct directions") +} + +func TestMITM_WebSocket_FramePayloadsLogged(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("payload-check"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) == 0 { + time.Sleep(50 * time.Millisecond) + continue + } + frames, err := st.FramesForConnection(conns[0].ID) + require.NoError(t, err) + if len(frames) >= 1 { + assert.Equal(t, "payload-check", string(frames[0].Payload)) + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for frame payload to be logged") +} + +// ── connection lifecycle ────────────────────────────────────────────────────── + +func TestMITM_WebSocket_CleanClose(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("bye"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "bye", string(payload)) + + err = wsConn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done")) + require.NoError(t, err) + wsConn.Close() +} + +func TestMITM_WebSocket_UpstreamDropsConnection(t *testing.T) { + st := newTestStore(t) + + // Upstream closes the connection after receiving the first message. + upstream := wsUpstreamServerWithHandler(t, func(conn *websocket.Conn) { + _, _, _ = conn.ReadMessage() + // Close without sending a close frame — simulates abrupt upstream drop. + conn.Close() + }) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("drop-me"))) + + // Browser side should get an error when upstream drops. + wsConn.SetReadDeadline(time.Now().Add(3 * time.Second)) + _, _, err = wsConn.ReadMessage() + assert.Error(t, err, "browser conn must get an error when upstream drops") +} + +func TestMITM_WebSocket_BrowserDropsConnection(t *testing.T) { + st := newTestStore(t) + + received := make(chan struct{}) + upstream := wsUpstreamServerWithHandler(t, func(conn *websocket.Conn) { + // Signal that upstream is ready, then try to read — should fail when + // browser drops. + close(received) + conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + _, _, _ = conn.ReadMessage() + }) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + + // Wait for upstream to be ready, then abruptly close browser side. + <-received + wsConn.Close() + + // Upstream read should fail — the proxy must propagate the close. + // The test passes if it completes without hanging (the upstream handler returns). +} + +// ── concurrent connections ──────────────────────────────────────────────────── + +func TestMITM_WebSocket_ConcurrentConnections(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + const nConns = 5 + var wg sync.WaitGroup + errors := make(chan error, nConns) + + for i := 0; i < nConns; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + msg := []byte("concurrent") + if err := wsConn.WriteMessage(websocket.TextMessage, msg); err != nil { + errors <- err + return + } + _, payload, err := wsConn.ReadMessage() + if err != nil { + errors <- err + return + } + if string(payload) != string(msg) { + errors <- nil + } + }(i) + } + + wg.Wait() + close(errors) + for err := range errors { + require.NoError(t, err) + } +} + +// ── HTTP and WebSocket through same proxy ───────────────────────────────────── + +// TestMITM_WebSocket_AfterHTTPOnSameProxy verifies that the proxy correctly +// handles both HTTP and WebSocket connections. HTTP and WS use separate CONNECT +// tunnels (which is how browsers work — each is a separate TCP connection). +func TestMITM_WebSocket_AfterHTTPOnSameProxy(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + + // First: a plain HTTPS request through the proxy. + httpConn := doConnectTunnel(t, proxyAddr, targetHost) + httpTLS := tlsOverTunnel(t, httpConn, host) + resp := sendHTTPOverConn(t, httpTLS, "GET", "/", targetHost, "") + resp.Body.Close() + // The upstream is a WS-only server so it won't return 200, + // but the proxy must forward the request and get a response. + assert.NotZero(t, resp.StatusCode) + + // Then: a WebSocket connection through the same proxy. + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("after-http"))) + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "after-http", string(payload), "WebSocket must work correctly after a prior HTTP request through the same proxy") +} + +// ── URL and path preservation ───────────────────────────────────────────────── + +func TestMITM_WebSocket_PathPreserved(t *testing.T) { + st := newTestStore(t) + + receivedPath := make(chan string, 1) + var wsUpgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath <- r.URL.Path + "?" + r.URL.RawQuery + conn, err := wsUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + conn.Close() + })) + t.Cleanup(srv.Close) + + upstreamAddr := strings.TrimPrefix(srv.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + _ = wsClientOverTunnelWithPath(t, proxyAddr, targetHost, "/chat/room?token=abc123") + + select { + case path := <-receivedPath: + assert.Equal(t, "/chat/room?token=abc123", path, "path and query must be preserved through proxy") + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for upstream to receive request") + } +} + +// ── store URL logged correctly ──────────────────────────────────────────────── + +func TestMITM_WebSocket_URLLoggedCorrectly(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnelWithPath(t, proxyAddr, targetHost, "/api/stream") + defer wsConn.Close() + + require.NoError(t, wsConn.WriteMessage(websocket.TextMessage, []byte("url-test"))) + _, _, err = wsConn.ReadMessage() + require.NoError(t, err) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + if len(conns) > 0 { + assert.Contains(t, conns[0].URL, "/api/stream", "URL must be logged with path") + assert.Contains(t, conns[0].URL, "wss://", "URL must use wss scheme") + return + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("timed out waiting for connection URL to be logged") +} + +// ── rapid fire messages ─────────────────────────────────────────────────────── + +func TestMITM_WebSocket_RapidFireMessages(t *testing.T) { + st := newTestStore(t) + upstream := wsUpstreamServer(t) + + upstreamAddr := strings.TrimPrefix(upstream.URL, "https://") + host, port, err := net.SplitHostPort(upstreamAddr) + require.NoError(t, err) + targetHost := net.JoinHostPort(host, port) + + proxyAddr := startProxyServer(t, st) + wsConn := wsClientOverTunnel(t, proxyAddr, targetHost) + defer wsConn.Close() + + const n = 100 + var sent atomic.Int64 + + // Send 100 messages without waiting for echo. + go func() { + for range n { + if err := wsConn.WriteMessage(websocket.TextMessage, []byte("rapid")); err != nil { + return + } + sent.Add(1) + } + }() + + // Receive all 100 echoes. + wsConn.SetReadDeadline(time.Now().Add(10 * time.Second)) + for range n { + _, payload, err := wsConn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "rapid", string(payload)) + } +} diff --git a/internal/store/annotation.go b/internal/store/annotation.go new file mode 100644 index 0000000..84b0c48 --- /dev/null +++ b/internal/store/annotation.go @@ -0,0 +1,58 @@ +package store + +import "fmt" + +// Annotation holds optional comment and highlight colour for a history row. +type Annotation struct { + HistoryID uint64 + Comment string + Color string // empty means no highlight; one of: "red", "orange", "yellow", "green", "blue", "purple" +} + +// SetAnnotation creates or replaces the annotation for the given history row. +func (s *Store) SetAnnotation(a Annotation) error { + return s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO annotations (history_id, comment, color) + VALUES (?, ?, ?) + ON CONFLICT(history_id) DO UPDATE SET + comment = excluded.comment, + color = excluded.color`, + a.HistoryID, a.Comment, a.Color, + ) + if err != nil { + return fmt.Errorf("store: set annotation %d: %w", a.HistoryID, err) + } + return nil + }) +} + +// DeleteAnnotation removes the annotation for the given history row. +func (s *Store) DeleteAnnotation(historyID uint64) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM annotations WHERE history_id = ?`, historyID) + if err != nil { + return fmt.Errorf("store: delete annotation %d: %w", historyID, err) + } + return nil + }) +} + +// AllAnnotations returns every annotation row. Called once at startup to +// populate the in-memory annotation map. +func (s *Store) AllAnnotations() ([]Annotation, error) { + rows, err := s.db.Query(`SELECT history_id, comment, color FROM annotations`) + if err != nil { + return nil, fmt.Errorf("store: all annotations: %w", err) + } + defer rows.Close() + var out []Annotation + for rows.Next() { + var a Annotation + if err := rows.Scan(&a.HistoryID, &a.Comment, &a.Color); err != nil { + return nil, fmt.Errorf("store: scan annotation: %w", err) + } + out = append(out, a) + } + return out, rows.Err() +} diff --git a/internal/store/annotation_test.go b/internal/store/annotation_test.go new file mode 100644 index 0000000..72c79e1 --- /dev/null +++ b/internal/store/annotation_test.go @@ -0,0 +1,92 @@ +package store_test + +import ( + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func logTxForAnnotation(t *testing.T, st *store.Store) uint64 { + t.Helper() + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + txs, err := st.AllTransactions() + require.NoError(t, err) + require.NotEmpty(t, txs) + return txs[0].ID +} + +func TestSetAnnotation_PersistsCommentAndColor(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + a := store.Annotation{HistoryID: id, Comment: "suspicious", Color: "red"} + require.NoError(t, st.SetAnnotation(a)) + + all, err := st.AllAnnotations() + require.NoError(t, err) + require.Len(t, all, 1) + assert.Equal(t, id, all[0].HistoryID) + assert.Equal(t, "suspicious", all[0].Comment) + assert.Equal(t, "red", all[0].Color) +} + +func TestSetAnnotation_UpdatesExisting(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "first", Color: "blue"})) + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "second", Color: "green"})) + + all, err := st.AllAnnotations() + require.NoError(t, err) + require.Len(t, all, 1) + assert.Equal(t, "second", all[0].Comment) + assert.Equal(t, "green", all[0].Color) +} + +func TestDeleteAnnotation_RemovesRow(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "note", Color: "yellow"})) + require.NoError(t, st.DeleteAnnotation(id)) + + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Empty(t, all) +} + +func TestDeleteAnnotation_NonExistent_NoError(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.DeleteAnnotation(9999)) +} + +func TestAllAnnotations_Empty(t *testing.T) { + st := newTestStore(t) + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Empty(t, all) +} + +func TestAllAnnotations_MultipleRows(t *testing.T) { + st := newTestStore(t) + + tx1 := makeTx("GET", "a.com", "https://a.com/", 200) + tx2 := makeTx("POST", "b.com", "https://b.com/", 201) + require.NoError(t, st.Log(tx1)) + require.NoError(t, st.Log(tx2)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 2) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: txs[0].ID, Comment: "c1", Color: "red"})) + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: txs[1].ID, Comment: "c2", Color: "blue"})) + + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Len(t, all, 2) +} diff --git a/internal/store/har.go b/internal/store/har.go new file mode 100644 index 0000000..047300a --- /dev/null +++ b/internal/store/har.go @@ -0,0 +1,239 @@ +package store + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +// HAR 1.2 — https://w3c.github.io/web-performance/specs/HAR/Overview.html + +type HAR struct { + Log HARLog `json:"log"` +} + +type HARLog struct { + Version string `json:"version"` + Creator HARCreator `json:"creator"` + Entries []HAREntry `json:"entries"` +} + +type HARCreator struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type HAREntry struct { + StartedDateTime string `json:"startedDateTime"` + Time int64 `json:"time"` + Request HARRequest `json:"request"` + Response HARResponse `json:"response"` + Timings HARTimings `json:"timings"` +} + +type HARRequest struct { + Method string `json:"method"` + URL string `json:"url"` + HTTPVersion string `json:"httpVersion"` + Headers []HARNameValue `json:"headers"` + QueryString []HARNameValue `json:"queryString"` + PostData *HARPostData `json:"postData,omitempty"` + HeadersSize int `json:"headersSize"` + BodySize int `json:"bodySize"` +} + +type HARResponse struct { + Status int `json:"status"` + StatusText string `json:"statusText"` + HTTPVersion string `json:"httpVersion"` + Headers []HARNameValue `json:"headers"` + Content HARContent `json:"content"` + RedirectURL string `json:"redirectURL"` + HeadersSize int `json:"headersSize"` + BodySize int `json:"bodySize"` +} + +type HARNameValue struct { + Name string `json:"name"` + Value string `json:"value"` +} + +type HARPostData struct { + MimeType string `json:"mimeType"` + Text string `json:"text"` +} + +type HARContent struct { + Size int `json:"size"` + MimeType string `json:"mimeType"` + Text string `json:"text,omitempty"` +} + +type HARTimings struct { + Send int64 `json:"send"` + Wait int64 `json:"wait"` + Receive int64 `json:"receive"` +} + +// ExportHAR serialises the given transactions into a HAR 1.2 JSON byte slice. +// Pass only the transactions you want exported — scope and out-of-scope +// filtering is the caller's responsibility. +func ExportHAR(txs []Transaction) ([]byte, error) { + har := HAR{ + Log: HARLog{ + Version: "1.2", + Creator: HARCreator{ + Name: "Shiv", + Version: "1.0", + }, + Entries: make([]HAREntry, 0, len(txs)), + }, + } + + for _, tx := range txs { + entry := HAREntry{ + StartedDateTime: tx.Timestamp.UTC().Format(time.RFC3339), + Time: tx.DurationMs, + Request: buildHARRequest(tx), + Response: buildHARResponse(tx), + Timings: HARTimings{ + Send: 0, + Wait: tx.DurationMs, + Receive: 0, + }, + } + har.Log.Entries = append(har.Log.Entries, entry) + } + + data, err := json.MarshalIndent(har, "", " ") + if err != nil { + return nil, fmt.Errorf("har: marshal: %w", err) + } + return data, nil +} + +func stripDefaultPort(rawURL string) string { + for _, prefix := range []string{"https://", "http://"} { + if !strings.HasPrefix(rawURL, prefix) { + continue + } + rest := rawURL[len(prefix):] + slash := strings.Index(rest, "/") + if slash < 0 { + slash = len(rest) + } + host := rest[:slash] + path := rest[slash:] + h, port, err := net.SplitHostPort(host) + if err != nil { + return rawURL + } + if (prefix == "https://" && port == "443") || (prefix == "http://" && port == "80") { + return prefix + h + path + } + return rawURL + } + return rawURL +} + +func buildHARRequest(tx Transaction) HARRequest { + headers := make([]HARNameValue, 0, len(tx.ReqHeaders)) + for k, vals := range tx.ReqHeaders { + for _, v := range vals { + headers = append(headers, HARNameValue{Name: k, Value: v}) + } + } + + httpVersion := tx.Proto + if httpVersion == "" { + httpVersion = "HTTP/1.1" + } + + req := HARRequest{ + Method: tx.Method, + URL: stripDefaultPort(tx.URL), + HTTPVersion: httpVersion, + Headers: headers, + QueryString: extractQueryString(tx.URL), + HeadersSize: -1, + BodySize: len(tx.ReqBody), + } + + if len(tx.ReqBody) > 0 { + mimeType := tx.ReqHeaders.Get("Content-Type") + if mimeType == "" { + mimeType = "application/octet-stream" + } + req.PostData = &HARPostData{ + MimeType: mimeType, + Text: string(tx.ReqBody), + } + } + + return req +} + +func buildHARResponse(tx Transaction) HARResponse { + headers := make([]HARNameValue, 0, len(tx.RespHeaders)) + for k, vals := range tx.RespHeaders { + for _, v := range vals { + headers = append(headers, HARNameValue{Name: k, Value: v}) + } + } + + httpVersion := tx.Proto + if httpVersion == "" { + httpVersion = "HTTP/1.1" + } + + mimeType := tx.RespHeaders.Get("Content-Type") + if mimeType == "" { + mimeType = "application/octet-stream" + } + + redirectURL := tx.RespHeaders.Get("Location") + + return HARResponse{ + Status: tx.StatusCode, + StatusText: statusText(tx.StatusCode), + HTTPVersion: httpVersion, + Headers: headers, + Content: HARContent{ + Size: len(tx.RespBody), + MimeType: mimeType, + Text: string(tx.RespBody), + }, + RedirectURL: redirectURL, + HeadersSize: -1, + BodySize: len(tx.RespBody), + } +} + +func extractQueryString(rawURL string) []HARNameValue { + idx := strings.Index(rawURL, "?") + if idx < 0 { + return []HARNameValue{} + } + parsed, err := url.ParseQuery(rawURL[idx+1:]) + if err != nil || len(parsed) == 0 { + return []HARNameValue{} + } + pairs := make([]HARNameValue, 0, len(parsed)) + for k, vals := range parsed { + for _, v := range vals { + pairs = append(pairs, HARNameValue{Name: k, Value: v}) + } + } + return pairs +} + +func statusText(code int) string { + if t := http.StatusText(code); t != "" { + return t + } + return fmt.Sprintf("%d", code) +} diff --git a/internal/store/history.go b/internal/store/history.go new file mode 100644 index 0000000..066c7c2 --- /dev/null +++ b/internal/store/history.go @@ -0,0 +1,401 @@ +package store + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "strconv" + "strings" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +// Transaction is a matched HTTP request/response pair. +// Host is always a pure hostname — never host:port. +// Port holds the destination port as an integer. +type Transaction struct { + ID uint64 + Timestamp time.Time + Host string + Port int + Method string + URL string + Proto string // "HTTP/2" or "HTTP/1.1" + ReqHeaders http.Header + ReqBody []byte + StatusCode int + RespHeaders http.Header + RespBody []byte + RespSize int + DurationMs int64 + TLS bool + InScope bool +} + +// TransactionFilter holds filter criteria for paginated history queries. +// Zero values mean "no filter" for that field. +type TransactionFilter struct { + Search string // matched against host + url + method + status (case-insensitive) + Host string // exact host match (from site map selection) + Port int // exact port match; 0 means no filter + PathPrefix string // URL path prefix filter (from site map selection) + ShowOutScope bool // if false, only in_scope=1 rows are returned + Methods []string // whitelist of HTTP methods; empty = all + StatusMin int // lower bound on status_code; 0 = no lower bound + StatusMax int // upper bound on status_code; 0 = no upper bound + HideStatic bool // if true, filter out common static asset extensions in memory +} + +// PageSize is the number of rows returned per page. +const PageSize = 100 + +// SiteMapEntry is a minimal host+port+url triplet used to build the site map +// at startup. Port is required to differentiate apps on the same host at +// different ports (e.g. localhost:3000 vs localhost:8080). +type SiteMapEntry struct { + Host string + Port int + URL string +} + +// SiteMapEntries returns the host, port, and url for every row in history. +// Used once at startup to populate the site map with all historical +// paths before the first page of transactions is displayed. +func (s *Store) SiteMapEntries() ([]SiteMapEntry, error) { + rows, err := s.db.Query(`SELECT host, port, url FROM history ORDER BY id ASC`) + if err != nil { + return nil, fmt.Errorf("store: site map entries: %w", err) + } + defer rows.Close() + var entries []SiteMapEntry + for rows.Next() { + var e SiteMapEntry + if err := rows.Scan(&e.Host, &e.Port, &e.URL); err != nil { + return nil, fmt.Errorf("store: scan site map entry: %w", err) + } + entries = append(entries, e) + } + return entries, rows.Err() +} + +func (s *Store) GetTransaction(id uint64) (*Transaction, error) { + var transaction Transaction + var timestampStr, reqHeaders, respHeaders string + var tlsFlag, scopeFlag int + err := s.db.QueryRow(` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, req_body, status_code, resp_headers, + resp_body, duration_ms, tls, in_scope, resp_size + FROM history WHERE id = ?`, id, + ).Scan( + &transaction.ID, ×tampStr, &transaction.Host, &transaction.Port, + &transaction.Method, &transaction.URL, &transaction.Proto, + &reqHeaders, &transaction.ReqBody, &transaction.StatusCode, &respHeaders, + &transaction.RespBody, &transaction.DurationMs, &tlsFlag, &scopeFlag, + &transaction.RespSize, + ) + if err != nil { + return nil, fmt.Errorf("store: get transaction %d: %w", id, err) + } + transaction.Timestamp, _ = time.Parse(time.RFC3339, timestampStr) + transaction.TLS = tlsFlag == 1 + transaction.InScope = scopeFlag == 1 + if err := json.Unmarshal([]byte(reqHeaders), &transaction.ReqHeaders); err != nil { + transaction.ReqHeaders = http.Header{} + } + if err := json.Unmarshal([]byte(respHeaders), &transaction.RespHeaders); err != nil { + transaction.RespHeaders = http.Header{} + } + return &transaction, nil +} + +// Log writes every transaction as a unique row — no deduplication. +func (s *Store) Log(t Transaction) error { + return s.write(func() error { + reqH, err := json.Marshal(t.ReqHeaders) + if err != nil { + return fmt.Errorf("store: marshal req headers: %w", err) + } + respH, err := json.Marshal(t.RespHeaders) + if err != nil { + return fmt.Errorf("store: marshal resp headers: %w", err) + } + proto := t.Proto + if proto == "" { + proto = "HTTP/1.1" + } + respSize := len(t.RespBody) + res, err := s.db.Exec(` + INSERT INTO history + (timestamp, host, port, method, url, proto, req_headers, req_body, + status_code, resp_headers, resp_body, duration_ms, tls, in_scope, resp_size) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + t.Timestamp.UTC().Format(time.RFC3339), + t.Host, t.Port, t.Method, t.URL, proto, + string(reqH), t.ReqBody, + t.StatusCode, string(respH), t.RespBody, + t.DurationMs, + boolToInt(t.TLS), + boolToInt(t.InScope), + respSize, + ) + if err != nil { + return fmt.Errorf("store: log transaction: %w", err) + } + id, _ := res.LastInsertId() + t.ID = uint64(id) + t.Proto = proto + t.RespSize = respSize + // Guard against sending on a closed channel when the store is shutting down. + select { + case <-s.done: + default: + select { + case s.Updates <- t: + default: + } + } + return nil + }) +} + +// TransactionsPage returns up to PageSize transactions with id < beforeID, +// ordered by id descending (newest first), applying filter criteria. +// Pass beforeID=0 to start from the newest row. +// Bodies are omitted; use GetTransaction for the full row. +func (s *Store) TransactionsPage(beforeID uint64, filter TransactionFilter) ([]Transaction, error) { + var conditions []string + var args []any + + if beforeID > 0 { + conditions = append(conditions, "id < ?") + args = append(args, beforeID) + } + + if !filter.ShowOutScope { + conditions = append(conditions, "in_scope = 1") + } + + if filter.Host != "" { + conditions = append(conditions, "host = ?") + args = append(args, filter.Host) + if filter.Port != 0 { + conditions = append(conditions, "port = ?") + args = append(args, filter.Port) + } + } + + if filter.PathPrefix != "" { + conditions = append(conditions, "(url LIKE ? OR url LIKE ? OR url LIKE ?)") + prefix := filter.PathPrefix + args = append(args, + "%"+prefix, + "%"+prefix+"?%", + "%"+prefix+"/%", + ) + } + + if filter.Search != "" { + terms := strings.Fields(strings.ToLower(filter.Search)) + for _, term := range terms { + conditions = append(conditions, "(LOWER(host) LIKE ? OR LOWER(url) LIKE ? OR LOWER(method) LIKE ? OR CAST(status_code AS TEXT) LIKE ?)") + like := "%" + term + "%" + args = append(args, like, like, like, like) + } + } + + if len(filter.Methods) > 0 { + placeholders := make([]string, len(filter.Methods)) + for i, m := range filter.Methods { + placeholders[i] = "?" + args = append(args, strings.ToUpper(m)) + } + conditions = append(conditions, "method IN ("+strings.Join(placeholders, ",")+")") + } + + if filter.StatusMin > 0 { + conditions = append(conditions, "status_code >= ?") + args = append(args, filter.StatusMin) + } + if filter.StatusMax > 0 { + conditions = append(conditions, "status_code <= ?") + args = append(args, filter.StatusMax) + } + + query := ` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, '' as req_body, status_code, resp_headers, + '' as resp_body, duration_ms, tls, in_scope, resp_size + FROM history` + + if len(conditions) > 0 { + query += " WHERE " + strings.Join(conditions, " AND ") + } + query += fmt.Sprintf(" ORDER BY id DESC LIMIT %d", PageSize) + + rows, err := s.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("store: query history page: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// SearchTransactions performs a full-text LIKE scan across request/response bodies +// and/or headers. At least one of inReqBody, inRespBody, inHeaders must be true. +// Results are ordered by id descending and capped at 200 rows. +func (s *Store) SearchTransactions(term string, inReqBody, inRespBody, inHeaders bool) ([]Transaction, error) { + if term == "" || (!inReqBody && !inRespBody && !inHeaders) { + return nil, nil + } + like := "%" + term + "%" + var clauses []string + var args []any + if inReqBody { + clauses = append(clauses, "CAST(req_body AS TEXT) LIKE ?") + args = append(args, like) + } + if inRespBody { + clauses = append(clauses, "CAST(resp_body AS TEXT) LIKE ?") + args = append(args, like) + } + if inHeaders { + clauses = append(clauses, "req_headers LIKE ?") + args = append(args, like) + clauses = append(clauses, "resp_headers LIKE ?") + args = append(args, like) + } + query := ` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, '' as req_body, status_code, resp_headers, + '' as resp_body, duration_ms, tls, in_scope, resp_size + FROM history + WHERE (` + strings.Join(clauses, " OR ") + `) + ORDER BY id DESC LIMIT 200` + rows, err := s.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("store: search transactions: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// AllTransactions returns all transactions ordered by id descending, capped at +// 500 rows. Used primarily in tests; production UI uses TransactionsPage. +func (s *Store) AllTransactions() ([]Transaction, error) { + rows, err := s.db.Query(` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, req_body, status_code, resp_headers, + resp_body, duration_ms, tls, in_scope, resp_size + FROM history + ORDER BY id DESC LIMIT 500`) + if err != nil { + return nil, fmt.Errorf("store: query all transactions: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// ClearHistory deletes all rows from the history table. +func (s *Store) ClearHistory() error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM history`) + if err != nil { + return fmt.Errorf("store: clear history: %w", err) + } + _, err = s.db.Exec(`VACUUM`) + if err != nil { + return fmt.Errorf("store: VACUUM history: %w", err) + } + + return nil + }) +} + +func scanTransactions(rows interface { + Next() bool + Scan(...any) error + Err() error +}) ([]Transaction, error) { + var txs []Transaction + for rows.Next() { + var tx Transaction + var ts, reqH, respH string + var tlsInt, scopeInt int + if err := rows.Scan( + &tx.ID, &ts, &tx.Host, &tx.Port, &tx.Method, &tx.URL, &tx.Proto, + &reqH, &tx.ReqBody, &tx.StatusCode, &respH, + &tx.RespBody, &tx.DurationMs, &tlsInt, &scopeInt, + &tx.RespSize, + ); err != nil { + return nil, fmt.Errorf("store: scan transaction: %w", err) + } + tx.Timestamp, _ = time.Parse(time.RFC3339, ts) + tx.TLS = tlsInt == 1 + tx.InScope = scopeInt == 1 + if err := json.Unmarshal([]byte(reqH), &tx.ReqHeaders); err != nil { + tx.ReqHeaders = http.Header{} + } + if err := json.Unmarshal([]byte(respH), &tx.RespHeaders); err != nil { + tx.RespHeaders = http.Header{} + } + txs = append(txs, tx) + } + return txs, rows.Err() +} + +// ObserveResponse implements events.ResponseObserver. It splits e.Host +// (which arrives as host:port from the proxy) into a pure hostname and integer +// port once here — no downstream consumer ever parses host/port again. +func (s *Store) ObserveResponse(e events.ResponseEvent) { + host, portStr, err := net.SplitHostPort(e.Host) + if err != nil { + host = e.Host + if e.TLS { + portStr = "443" + } else { + portStr = "80" + } + } + port, _ := strconv.Atoi(portStr) + + if err := s.Log(Transaction{ + Timestamp: e.Timestamp, + Host: host, + Port: port, + Proto: e.Proto, + Method: e.Method, + URL: e.URL, + ReqHeaders: e.ReqHeaders, + ReqBody: e.ReqBody, + StatusCode: e.StatusCode, + RespHeaders: e.RespHeaders, + RespBody: e.RespBody, + DurationMs: e.DurationMs, + TLS: e.TLS, + InScope: s.InScope(host), + }); err != nil { + logger.Error("store: observe response: %v", err) + } +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +// DeleteTransactionsByHostPort deletes all history rows for the given host and port. +func (s *Store) DeleteTransactionsByHostPort(host string, port int) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM history WHERE host = ? AND port = ?`, host, port) + if err != nil { + return fmt.Errorf("store: delete transactions for %s:%d: %w", host, port, err) + } + return nil + }) +} diff --git a/internal/store/history_test.go b/internal/store/history_test.go new file mode 100644 index 0000000..d3a8e99 --- /dev/null +++ b/internal/store/history_test.go @@ -0,0 +1,190 @@ +package store_test + +import ( + "fmt" + "net/http" + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeTx(method, host, url string, status int) store.Transaction { + return store.Transaction{ + Timestamp: time.Now().UTC(), + Host: host, + Method: method, + URL: url, + ReqHeaders: http.Header{"Content-Type": []string{"text/plain"}}, + ReqBody: []byte("request body"), + StatusCode: status, + RespHeaders: http.Header{"Content-Type": []string{"text/html"}}, + RespBody: []byte("response body"), + DurationMs: 10, + TLS: false, + InScope: true, + } +} + +func TestLog_Insert(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + assert.Equal(t, "GET", txs[0].Method) + assert.Equal(t, "example.com", txs[0].Host) + assert.Equal(t, 200, txs[0].StatusCode) + assert.True(t, txs[0].InScope) +} + +func TestLog_NoDedup_SameMethodHostURLStatus(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2, "no dedup — every request is a unique row") +} + +func TestLog_NoDedupOnDifferentStatus(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 404))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_NoDedupOnDifferentMethod(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 200))) + require.NoError(t, st.Log(makeTx("POST", "example.com", "https://example.com/", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_NoDedupOnDifferentURL(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/a", 200))) + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/b", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_TLSFlagPersisted(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + tx.TLS = true + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + assert.True(t, txs[0].TLS) +} + +func TestAllTransactions_OrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "first.com", "https://first.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "second.com", "https://second.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "third.com", "https://third.com/", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 3) + assert.Equal(t, "third.com", txs[0].Host) + assert.Equal(t, "second.com", txs[1].Host) + assert.Equal(t, "first.com", txs[2].Host) +} + +func TestAllTransactions_CappedAt500(t *testing.T) { + st := newTestStore(t) + + for i := 0; i < 510; i++ { + url := fmt.Sprintf("https://example.com/path/%d", i) + require.NoError(t, st.Log(makeTx("GET", "example.com", url, 200))) + } + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.LessOrEqual(t, len(txs), 500) +} + +func TestGetTransaction_ReturnsFullBodies(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("POST", "example.com", "https://example.com/api", 201) + tx.ReqBody = []byte(`{"input":"data"}`) + tx.RespBody = []byte(`{"output":"result"}`) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + + full, err := st.GetTransaction(txs[0].ID) + require.NoError(t, err) + assert.Equal(t, []byte(`{"input":"data"}`), full.ReqBody) + assert.Equal(t, []byte(`{"output":"result"}`), full.RespBody) + assert.Equal(t, "POST", full.Method) + assert.Equal(t, 201, full.StatusCode) +} + +func TestGetTransaction_HeadersDeserialized(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + tx.ReqHeaders = http.Header{"X-Request-Id": []string{"req-123"}} + tx.RespHeaders = http.Header{"X-Trace-Id": []string{"trace-456"}} + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + + full, err := st.GetTransaction(txs[0].ID) + require.NoError(t, err) + assert.Equal(t, "req-123", full.ReqHeaders.Get("X-Request-Id")) + assert.Equal(t, "trace-456", full.RespHeaders.Get("X-Trace-Id")) +} + +func TestGetTransaction_NotFound(t *testing.T) { + st := newTestStore(t) + _, err := st.GetTransaction(99999) + assert.Error(t, err) +} + +func TestClearHistory_RemovesAllRows(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "b.com", "https://b.com/", 200))) + require.NoError(t, st.ClearHistory()) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Empty(t, txs) +} + +func TestClearHistory_EmptyStoreNoError(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.ClearHistory()) +} diff --git a/internal/store/intercept.go b/internal/store/intercept.go new file mode 100644 index 0000000..e86c683 --- /dev/null +++ b/internal/store/intercept.go @@ -0,0 +1,233 @@ +package store + +import ( + "net" + "net/http" + "sync" + + "github.com/shiv/internal/events" +) + +// PendingRequest is a request held at the intercept gate waiting for a +// user decision. The proxy goroutine blocks on Reply until the UI responds. +// done is closed when the request is cancelled via bypass so watchQueue +// can skip stale entries without deadlocking. +type PendingRequest struct { + Request *http.Request + Body []byte + Reply chan Decision + done chan struct{} +} + +func newPendingRequest(req *http.Request, body []byte) *PendingRequest { + return &PendingRequest{ + Request: req, + Body: body, + Reply: make(chan Decision, 1), + done: make(chan struct{}), + } +} + +// Cancel marks this pending request as cancelled. Safe to call multiple times. +func (p *PendingRequest) Cancel() { + select { + case <-p.done: + // already cancelled + default: + close(p.done) + } +} + +// IsDone returns true if this request was cancelled before the UI processed it. +func (p *PendingRequest) IsDone() bool { + select { + case <-p.done: + return true + default: + return false + } +} + +// Decision is the user's response to a pending request. +type Decision struct { + Forward bool + Request *http.Request + Body []byte +} + +// InterceptGate manages intercept state. When enabled, every request enters +// a serialising semaphore — one request is shown to the UI at a time, others +// wait in order. ForwardAll releases all currently waiting goroutines by +// closing the bypass channel and immediately replacing it so future requests +// are intercepted normally again. +type InterceptGate struct { + mu sync.RWMutex + enabled bool + + queue chan *PendingRequest // UI reads from here + sem chan struct{} // capacity 1 — one request shown at a time + + bypassMu sync.Mutex + bypass chan struct{} // closed to release all waiting goroutines + bypassClosed bool // tracks whether current bypass is already closed + + scopeMu sync.RWMutex + scopeCheckFn func(host string) bool // nil = intercept all; non-nil = intercept only if returns true +} + +// NewInterceptGate creates a new gate. Intercept is off by default. +func NewInterceptGate() *InterceptGate { + return &InterceptGate{ + queue: make(chan *PendingRequest, 1024), + sem: make(chan struct{}, 1), + bypass: make(chan struct{}), + } +} + +// Enabled returns whether intercept is currently on. +func (g *InterceptGate) Enabled() bool { + g.mu.RLock() + defer g.mu.RUnlock() + return g.enabled +} + +// SetEnabled turns intercept on or off. +// Turning off triggers bypass so all goroutines currently blocked in Hold +// are released immediately without deadlocking. +// Turning on resets the bypass channel so future requests are intercepted. +// g.mu and bypassMu are never held simultaneously to avoid deadlock with Hold. +func (g *InterceptGate) SetEnabled(v bool) { + g.mu.Lock() + g.enabled = v + g.mu.Unlock() + + if !v { + g.bypassMu.Lock() + if !g.bypassClosed { + close(g.bypass) + g.bypassClosed = true + } + g.bypassMu.Unlock() + } else { + g.bypassMu.Lock() + g.bypass = make(chan struct{}) + g.bypassClosed = false + g.bypassMu.Unlock() + } +} + +// SetScopeFilter sets a predicate that gates which requests enter the queue. +// Pass nil to intercept all requests. Pass a function returning true for +// in-scope hosts to intercept scope-only traffic. +func (g *InterceptGate) SetScopeFilter(fn func(host string) bool) { + g.scopeMu.Lock() + g.scopeCheckFn = fn + g.scopeMu.Unlock() +} + +// Hold blocks until the user makes a decision on this request. +// If intercept is off, returns immediately. +// Requests acquire the semaphore in arrival order so the UI sees exactly +// one request at a time. If bypass fires at any point during Hold, the +// goroutine exits immediately and the request is passed through unmodified. +// Any pending already in the queue is marked cancelled so watchQueue skips it. +func (g *InterceptGate) Hold(req *http.Request, body []byte) (*http.Request, []byte, bool) { + if !g.Enabled() { + return req, body, true + } + + g.scopeMu.RLock() + fn := g.scopeCheckFn + g.scopeMu.RUnlock() + if fn != nil { + host := req.Host + if h, _, err := net.SplitHostPort(host); err == nil { + host = h + } + if !fn(host) { + return req, body, true + } + } + + // Capture the current bypass reference once. This goroutine watches this + // exact channel throughout — if it gets closed (ForwardAll or SetEnabled(false)), + // we see it no matter where in Hold we are. + g.bypassMu.Lock() + bypass := g.bypass + g.bypassMu.Unlock() + + // Wait for semaphore (one slot) or bypass signal. + select { + case g.sem <- struct{}{}: + case <-bypass: + return req, body, true + } + + // Re-check enabled — may have changed while we waited for the semaphore. + if !g.Enabled() { + <-g.sem + return req, body, true + } + + // Re-check bypass — may have been closed while we raced for the semaphore. + select { + case <-bypass: + <-g.sem + return req, body, true + default: + } + + // We hold the semaphore and bypass is open. Put the request in the queue. + pending := newPendingRequest(req, body) + select { + case g.queue <- pending: + case <-bypass: + pending.Cancel() + <-g.sem + return req, body, true + } + + // Wait for the user's decision or a bypass signal. + select { + case decision := <-pending.Reply: + <-g.sem + return decision.Request, decision.Body, decision.Forward + case <-bypass: + // Bypass fired while we were waiting. Mark the pending as cancelled + // so watchQueue skips this stale queue entry, then release the semaphore. + pending.Cancel() + <-g.sem + return req, body, true + } +} + +// ForwardAll releases all goroutines currently blocked in Hold, passing their +// requests through unmodified. Future requests are intercepted normally — +// the bypass channel is replaced immediately after closing the old one. +func (g *InterceptGate) ForwardAll() { + g.bypassMu.Lock() + if !g.bypassClosed { + close(g.bypass) + g.bypassClosed = true + } + // Replace immediately so future requests are intercepted normally. + g.bypass = make(chan struct{}) + g.bypassClosed = false + g.bypassMu.Unlock() +} + +// Queue returns the channel the UI reads pending requests from. +func (g *InterceptGate) Queue() <-chan *PendingRequest { + return g.queue +} + +// HandleRequest implements events.RequestMiddleware. +// This is a thin adapter: it calls Hold and maps the result to RequestResult. +func (g *InterceptGate) HandleRequest(e events.RequestEvent) events.RequestResult { + req, body, forward := g.Hold(e.Request, e.Body) + return events.RequestResult{ + Drop: !forward, + Request: req, + Body: body, + } +} diff --git a/internal/store/intruder.go b/internal/store/intruder.go new file mode 100644 index 0000000..4968f5d --- /dev/null +++ b/internal/store/intruder.go @@ -0,0 +1,72 @@ +package store + +import "fmt" + +type IntruderConfig struct { + DelayMs int + StopOnStatus int + MaxRedirects int + FollowRedirects string // "never", "always", "in-scope" + TimeoutMs int + RawRequest string + Payloads string +} + +func DefaultIntruderConfig() IntruderConfig { + return IntruderConfig{ + DelayMs: 0, + StopOnStatus: 0, + MaxRedirects: 10, + FollowRedirects: "never", + TimeoutMs: 30000, + RawRequest: "", + Payloads: "", + } +} + +func (s *Store) LoadIntruderConfig() IntruderConfig { + config := DefaultIntruderConfig() + err := s.db.QueryRow(` + SELECT delay_ms, stop_on_status, max_redirects, follow_redirects, timeout_ms, raw_request, payloads + FROM intruder_config WHERE id = 1`).Scan( + &config.DelayMs, + &config.StopOnStatus, + &config.MaxRedirects, + &config.FollowRedirects, + &config.TimeoutMs, + &config.RawRequest, + &config.Payloads, + ) + if err != nil { + return config + } + return config +} + +func (s *Store) SaveIntruderConfig(config IntruderConfig) error { + return s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO intruder_config (id, delay_ms, stop_on_status, max_redirects, follow_redirects, timeout_ms, raw_request, payloads) + VALUES (1, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + delay_ms = excluded.delay_ms, + stop_on_status = excluded.stop_on_status, + max_redirects = excluded.max_redirects, + follow_redirects = excluded.follow_redirects, + timeout_ms = excluded.timeout_ms, + raw_request = excluded.raw_request, + payloads = excluded.payloads`, + config.DelayMs, + config.StopOnStatus, + config.MaxRedirects, + config.FollowRedirects, + config.TimeoutMs, + config.RawRequest, + config.Payloads, + ) + if err != nil { + return fmt.Errorf("store: save intruder config: %w", err) + } + return nil + }) +} diff --git a/internal/store/loot.go b/internal/store/loot.go new file mode 100644 index 0000000..c610db1 --- /dev/null +++ b/internal/store/loot.go @@ -0,0 +1,103 @@ +package store + +import ( + "fmt" + "time" +) + +type LootEntry struct { + ID int64 + Title string + Severity string + Notes string + HistoryID *uint64 + RawRequest string + RawResponse string + CreatedAt time.Time +} + +func (s *Store) AllLoot() ([]LootEntry, error) { + rows, err := s.db.Query(` + SELECT id, title, severity, notes, history_id, raw_request, raw_response, created_at + FROM loot ORDER BY + CASE severity + WHEN 'Critical' THEN 1 + WHEN 'High' THEN 2 + WHEN 'Medium' THEN 3 + WHEN 'Low' THEN 4 + WHEN 'Info' THEN 5 + ELSE 6 + END ASC, id DESC`) + if err != nil { + return nil, fmt.Errorf("store: query loot: %w", err) + } + defer rows.Close() + + var entries []LootEntry + for rows.Next() { + var entry LootEntry + var histID *int64 + var timestampStr string + if err := rows.Scan(&entry.ID, &entry.Title, &entry.Severity, &entry.Notes, &histID, &entry.RawRequest, &entry.RawResponse, ×tampStr); err != nil { + return nil, fmt.Errorf("store: scan loot: %w", err) + } + if histID != nil { + histIDUint := uint64(*histID) + entry.HistoryID = &histIDUint + } + entry.CreatedAt, _ = time.Parse(time.RFC3339, timestampStr) + entries = append(entries, entry) + } + return entries, rows.Err() +} + +func (s *Store) AddLoot(entry LootEntry) (int64, error) { + var id int64 + err := s.write(func() error { + var histID any + if entry.HistoryID != nil { + histID = *entry.HistoryID + } + result, err := s.db.Exec(` + INSERT INTO loot (title, severity, notes, history_id, raw_request, raw_response, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + entry.Title, entry.Severity, entry.Notes, histID, + entry.RawRequest, entry.RawResponse, + time.Now().UTC().Format(time.RFC3339), + ) + if err != nil { + return fmt.Errorf("store: add loot: %w", err) + } + id, err = result.LastInsertId() + return err + }) + if err != nil { + return 0, err + } + s.notifyLoot() + return id, nil +} + +func (s *Store) DeleteLoot(id int64) error { + err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM loot WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete loot: %w", err) + } + return nil + }) + if err != nil { + return err + } + s.notifyLoot() + return nil +} + +// notifyLoot signals the UI that loot state has changed. +// Non-blocking — a full buffer means the UI is already scheduled to reload. +func (s *Store) notifyLoot() { + select { + case s.LootEntries <- struct{}{}: + default: + } +} diff --git a/internal/store/loot_test.go b/internal/store/loot_test.go new file mode 100644 index 0000000..59bd4fe --- /dev/null +++ b/internal/store/loot_test.go @@ -0,0 +1,155 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeLoot(title, severity, notes string) store.LootEntry { + return store.LootEntry{ + Title: title, + Severity: severity, + Notes: notes, + } +} + +func TestAddLoot_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + entry := store.LootEntry{ + Title: "Admin panel exposed", + Severity: "High", + Notes: "Found at /admin with no auth", + } + id, err := st.AddLoot(entry) + require.NoError(t, err) + assert.Positive(t, id) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + + got := entries[0] + assert.Equal(t, id, got.ID) + assert.Equal(t, "Admin panel exposed", got.Title) + assert.Equal(t, "High", got.Severity) + assert.Equal(t, "Found at /admin with no auth", got.Notes) + assert.Nil(t, got.HistoryID) + assert.WithinDuration(t, time.Now(), got.CreatedAt, 5*time.Second) +} + +func TestAddLoot_WithLinkedHistoryID(t *testing.T) { + st := newTestStore(t) + + // first log a transaction to get a valid history ID + tx := makeTx("GET", "example.com", "https://example.com/secret", 200) + require.NoError(t, st.Log(tx)) + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + histID := txs[0].ID + + entry := store.LootEntry{ + Title: "Secret endpoint", + Severity: "Critical", + Notes: "Returns sensitive data", + HistoryID: &histID, + } + id, err := st.AddLoot(entry) + require.NoError(t, err) + assert.Positive(t, id) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + require.NotNil(t, entries[0].HistoryID) + assert.Equal(t, histID, *entries[0].HistoryID) +} + +func TestAllLoot_OrderedBySeverityThenIDDesc(t *testing.T) { + st := newTestStore(t) + + // Insert in reverse severity order to confirm sorting + _, err := st.AddLoot(makeLoot("Info finding", "Info", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Low finding", "Low", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Critical finding", "Critical", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("High finding", "High", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Medium finding", "Medium", "")) + require.NoError(t, err) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 5) + + assert.Equal(t, "Critical", entries[0].Severity) + assert.Equal(t, "High", entries[1].Severity) + assert.Equal(t, "Medium", entries[2].Severity) + assert.Equal(t, "Low", entries[3].Severity) + assert.Equal(t, "Info", entries[4].Severity) +} + +func TestAllLoot_MultipleSameSeverityOrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + id1, err := st.AddLoot(makeLoot("First High", "High", "")) + require.NoError(t, err) + id2, err := st.AddLoot(makeLoot("Second High", "High", "")) + require.NoError(t, err) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 2) + + // same severity — ordered by id DESC so id2 comes first + assert.Equal(t, id2, entries[0].ID) + assert.Equal(t, id1, entries[1].ID) +} + +func TestDeleteLoot(t *testing.T) { + st := newTestStore(t) + + id, err := st.AddLoot(makeLoot("To delete", "Low", "")) + require.NoError(t, err) + + require.NoError(t, st.DeleteLoot(id)) + + entries, err := st.AllLoot() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestDeleteLoot_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteLoot(9999)) +} + +func TestDeleteLoot_DoesNotAffectOtherEntries(t *testing.T) { + st := newTestStore(t) + + idKeep, err := st.AddLoot(makeLoot("Keep", "High", "")) + require.NoError(t, err) + idDel, err := st.AddLoot(makeLoot("Delete", "Low", "")) + require.NoError(t, err) + + require.NoError(t, st.DeleteLoot(idDel)) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, idKeep, entries[0].ID) +} + +func TestAllLoot_Empty(t *testing.T) { + st := newTestStore(t) + entries, err := st.AllLoot() + require.NoError(t, err) + assert.Empty(t, entries) +} diff --git a/internal/store/plugin.go b/internal/store/plugin.go new file mode 100644 index 0000000..545d9b4 --- /dev/null +++ b/internal/store/plugin.go @@ -0,0 +1,156 @@ +package store + +import ( + "fmt" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +const maxPluginLogLines = 500 + +// PluginEntry represents a loaded plugin and its persistent enabled state. +type PluginEntry struct { + Name string + Path string + Enabled bool +} + +// RegisterPlugin records a plugin in the database. If the plugin is new, +// enabled defaults to true. If it already exists, only the path is updated — +// the existing enabled state is preserved. Notifies the UI via PluginEntries. +func (s *Store) RegisterPlugin(name, path string) error { + if err := s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO plugins (name, path, enabled) VALUES (?, ?, 1) + ON CONFLICT(name) DO UPDATE SET path = excluded.path`, + name, path, + ) + if err != nil { + return fmt.Errorf("store: register plugin %s: %w", name, err) + } + return nil + }); err != nil { + return err + } + + select { + case s.PluginEntries <- struct{}{}: + default: + } + return nil +} + +// Scans for plugins on the disk, removes them from DB if they are deleted from disk. +func (s *Store) ScanPlugins(active []string) error { + set := make(map[string]struct{}, len(active)) + for _, name := range active { + set[name] = struct{}{} + } + all, err := s.AllPlugins() + if err != nil { + return err + } + for _, e := range all { + if _, ok := set[e.Name]; ok { + continue + } + if err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM plugins WHERE name = ?`, e.Name) + return err + }); err != nil { + return fmt.Errorf("store: scan plugins remove %s: %w", e.Name, err) + } + } + select { + case s.PluginEntries <- struct{}{}: + default: + } + return nil +} + +// PluginEnabled returns the persisted enabled state for a plugin. +// Returns true if the plugin record is not found (safe default for new plugins). +func (s *Store) PluginEnabled(name string) (bool, error) { + var enabled int + err := s.db.QueryRow(`SELECT enabled FROM plugins WHERE name = ?`, name).Scan(&enabled) + if err != nil { + return true, err + } + return enabled == 1, nil +} + +// AllPlugins returns all known plugins ordered by name. +func (s *Store) AllPlugins() ([]PluginEntry, error) { + rows, err := s.db.Query(`SELECT name, path, enabled FROM plugins ORDER BY name ASC`) + if err != nil { + return nil, fmt.Errorf("store: query plugins: %w", err) + } + defer rows.Close() + + var entries []PluginEntry + for rows.Next() { + var e PluginEntry + var enabledInt int + if err := rows.Scan(&e.Name, &e.Path, &enabledInt); err != nil { + return nil, fmt.Errorf("store: scan plugin: %w", err) + } + e.Enabled = enabledInt == 1 + entries = append(entries, e) + } + return entries, rows.Err() +} + +// PluginLogs returns the in-memory log lines for a plugin, oldest first. +// Returns a copy to avoid races with the writer. +func (s *Store) PluginLogs(name string) []string { + s.pluginLogsMu.RLock() + defer s.pluginLogsMu.RUnlock() + lines := s.pluginLogs[name] + if len(lines) == 0 { + return nil + } + result := make([]string, len(lines)) + copy(result, lines) + return result +} + +// ObservePluginLog implements events.PluginLogObserver. +// Appends the log line to the in-memory ring buffer and notifies the UI. +func (s *Store) ObservePluginLog(e events.PluginLogEvent) { + s.pluginLogsMu.Lock() + lines := s.pluginLogs[e.Name] + lines = append(lines, e.Message) + if len(lines) > maxPluginLogLines { + lines = lines[len(lines)-maxPluginLogLines:] + } + s.pluginLogs[e.Name] = lines + s.pluginLogsMu.Unlock() + + select { + case s.PluginEntries <- struct{}{}: + default: + } +} + +// ObservePluginEnabled implements events.PluginEnabledObserver. +// Persists the new enabled state to the database and notifies the UI. +func (s *Store) ObservePluginEnabled(e events.SetPluginEnabledEvent) { + if err := s.write(func() error { + _, err := s.db.Exec(`UPDATE plugins SET enabled = ? WHERE name = ?`, + boolToInt(e.Enabled), e.Name, + ) + if err != nil { + return fmt.Errorf("store: set plugin enabled %s: %w", e.Name, err) + } + return nil + }); err != nil { + logger.Error("store: observe plugin enabled: %v", err) + return + } + + select { + case s.PluginEntries <- struct{}{}: + default: + } +} diff --git a/internal/store/plugin_test.go b/internal/store/plugin_test.go new file mode 100644 index 0000000..b640628 --- /dev/null +++ b/internal/store/plugin_test.go @@ -0,0 +1,368 @@ +package store_test + +import ( + "fmt" + "sync" + "testing" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── RegisterPlugin ──────────────────────────────────────────────────────────── + +func TestRegisterPlugin_NewPlugin_DefaultsToEnabled(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "test.lua", entries[0].Name) + assert.Equal(t, "/plugins/test.lua", entries[0].Path) + assert.True(t, entries[0].Enabled) +} + +func TestRegisterPlugin_ExistingPlugin_PreservesEnabledState(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + // Re-register with a new path — enabled state must be preserved. + require.NoError(t, st.RegisterPlugin("test.lua", "/new/path/test.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "/new/path/test.lua", entries[0].Path) + assert.False(t, entries[0].Enabled, "enabled state must survive re-registration") +} + +func TestRegisterPlugin_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + + // Drain any startup notifications. + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after RegisterPlugin") + } +} + +func TestRegisterPlugin_MultiplePlugins(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("a.lua", "/plugins/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/plugins/b.lua")) + require.NoError(t, st.RegisterPlugin("c.lua", "/plugins/c.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Len(t, entries, 3) +} + +// ── PluginEnabled ───────────────────────────────────────────────────────────── + +func TestPluginEnabled_UnknownPlugin_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + + enabled, err := st.PluginEnabled("nonexistent.lua") + // Error is expected (no row) but the safe default is true. + assert.True(t, enabled) + _ = err +} + +func TestPluginEnabled_NewPlugin_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +func TestPluginEnabled_AfterDisable_ReturnsFalse(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.False(t, enabled) +} + +func TestPluginEnabled_AfterReEnable_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +// ── AllPlugins ──────────────────────────────────────────────────────────────── + +func TestAllPlugins_Empty(t *testing.T) { + st := newTestStore(t) + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestAllPlugins_OrderedByNameAsc(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("zebra.lua", "/z.lua")) + require.NoError(t, st.RegisterPlugin("alpha.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("mango.lua", "/m.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 3) + assert.Equal(t, "alpha.lua", entries[0].Name) + assert.Equal(t, "mango.lua", entries[1].Name) + assert.Equal(t, "zebra.lua", entries[2].Name) +} + +func TestAllPlugins_EnabledFieldReflectsState(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("on.lua", "/on.lua")) + require.NoError(t, st.RegisterPlugin("off.lua", "/off.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "off.lua", Enabled: false}) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 2) + + byName := make(map[string]store.PluginEntry) + for _, e := range entries { + byName[e.Name] = e + } + assert.True(t, byName["on.lua"].Enabled) + assert.False(t, byName["off.lua"].Enabled) +} + +// ── ScanPlugins ─────────────────────────────────────────────────────────────── + +func TestScanPlugins_RemovesStaleRecords(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("keep.lua", "/keep.lua")) + require.NoError(t, st.RegisterPlugin("gone.lua", "/gone.lua")) + + require.NoError(t, st.ScanPlugins([]string{"keep.lua"})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "keep.lua", entries[0].Name) +} + +func TestScanPlugins_EmptyActive_RemovesAll(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("a.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/b.lua")) + + require.NoError(t, st.ScanPlugins([]string{})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestScanPlugins_AllPresent_NothingRemoved(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("a.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/b.lua")) + + require.NoError(t, st.ScanPlugins([]string{"a.lua", "b.lua"})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Len(t, entries, 2) +} + +func TestScanPlugins_EmptyStore_NoPanic(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.ScanPlugins([]string{"a.lua"})) +} + +func TestScanPlugins_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("gone.lua", "/gone.lua")) + + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + require.NoError(t, st.ScanPlugins([]string{})) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ScanPlugins") + } +} + +// ── PluginLogs ──────────────────────────────────────────────────────────────── + +func TestPluginLogs_UnknownPlugin_ReturnsNil(t *testing.T) { + st := newTestStore(t) + assert.Nil(t, st.PluginLogs("nonexistent.lua")) +} + +func TestPluginLogs_ReturnsLinesInOrder(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 1"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 2"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 3"}) + + lines := st.PluginLogs("test.lua") + require.Len(t, lines, 3) + assert.Equal(t, "line 1", lines[0]) + assert.Equal(t, "line 2", lines[1]) + assert.Equal(t, "line 3", lines[2]) +} + +func TestPluginLogs_ReturnsACopy(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "original"}) + + lines := st.PluginLogs("test.lua") + lines[0] = "mutated" + + // Original must be unchanged. + assert.Equal(t, "original", st.PluginLogs("test.lua")[0]) +} + +func TestPluginLogs_IsolatedPerPlugin(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "a.lua", Message: "from a"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "b.lua", Message: "from b"}) + + assert.Equal(t, []string{"from a"}, st.PluginLogs("a.lua")) + assert.Equal(t, []string{"from b"}, st.PluginLogs("b.lua")) +} + +func TestPluginLogs_RingBuffer_CapsAtMaxLines(t *testing.T) { + st := newTestStore(t) + // Write more than maxPluginLogLines (500) entries. + for i := 0; i < 510; i++ { + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: fmt.Sprintf("line %d", i)}) + } + + lines := st.PluginLogs("test.lua") + assert.Len(t, lines, 500) + // The oldest 10 lines (0-9) were evicted; buffer starts at line 10. + assert.Equal(t, "line 10", lines[0]) + assert.Equal(t, "line 509", lines[499]) +} + +// ── ObservePluginLog ────────────────────────────────────────────────────────── + +func TestObservePluginLog_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "hello"}) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ObservePluginLog") + } +} + +func TestObservePluginLog_ChannelNonBlocking_WhenFull(t *testing.T) { + st := newTestStore(t) + // Fill the channel to capacity. + for i := 0; i < cap(st.PluginEntries); i++ { + st.PluginEntries <- struct{}{} + } + // Must not block even when channel is full. + assert.NotPanics(t, func() { + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "msg"}) + }) +} + +func TestObservePluginLog_ConcurrentWrites_Safe(t *testing.T) { + st := newTestStore(t) + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func(i int) { + defer wg.Done() + st.ObservePluginLog(events.PluginLogEvent{ + Name: "test.lua", + Message: fmt.Sprintf("msg %d", i), + }) + }(i) + } + wg.Wait() + + lines := st.PluginLogs("test.lua") + assert.Len(t, lines, goroutines) +} + +// ── ObservePluginEnabled ────────────────────────────────────────────────────── + +func TestObservePluginEnabled_PersistsDisabled(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.False(t, enabled) +} + +func TestObservePluginEnabled_PersistsEnabled(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +func TestObservePluginEnabled_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ObservePluginEnabled") + } +} + +func TestObservePluginEnabled_UnknownPlugin_NoError(t *testing.T) { + st := newTestStore(t) + // Toggling a plugin not in the DB must not panic or crash. + assert.NotPanics(t, func() { + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "ghost.lua", Enabled: false}) + }) +} diff --git a/internal/store/project.go b/internal/store/project.go new file mode 100644 index 0000000..e379f32 --- /dev/null +++ b/internal/store/project.go @@ -0,0 +1,245 @@ +package store + +import ( + "database/sql" + "fmt" + "sync" + + _ "modernc.org/sqlite" +) + +type Store struct { + db *sql.DB + writeCh chan func() error + done chan struct{} + + Updates chan Transaction + Intercept *InterceptGate + WebSocketConnections chan WebSocketConnection + WebSocketFrames chan WebSocketFrame + LootEntries chan struct{} + PluginEntries chan struct{} + + // pluginLogs holds in-memory log lines per plugin name. + // Protected by pluginLogsMu — separate from the SQLite write serialiser + // because log appends must never block on DB I/O. + pluginLogsMu sync.RWMutex + pluginLogs map[string][]string + + // scopeCache is an in-memory copy of the scope table. + // Populated lazily on first InScope call and invalidated on every write. + // Protected by scopeCacheMu — separate from the write serialiser so + // InScope reads never block on DB I/O. + scopeCacheMu sync.RWMutex + scopeCache []ScopeEntry + scopeCacheSet bool +} + +func Open(path string) (*Store, error) { + db, err := sql.Open("sqlite", path) + if err != nil { + return nil, fmt.Errorf("store: open %s: %w", path, err) + } + db.SetMaxOpenConns(1) + if _, err := db.Exec(`PRAGMA journal_mode=WAL`); err != nil { + return nil, fmt.Errorf("store: enable WAL: %w", err) + } + + projectStore := &Store{ + db: db, + writeCh: make(chan func() error, 256), + done: make(chan struct{}), + Updates: make(chan Transaction, 256), + Intercept: NewInterceptGate(), + WebSocketConnections: make(chan WebSocketConnection, 256), + WebSocketFrames: make(chan WebSocketFrame, 256), + LootEntries: make(chan struct{}, 256), + PluginEntries: make(chan struct{}, 256), + pluginLogs: make(map[string][]string), + } + if err := projectStore.migrate(); err != nil { + db.Close() + return nil, err + } + + go projectStore.writeLoop() + return projectStore, nil +} + +func (s *Store) Close() error { + close(s.done) + close(s.Updates) + close(s.WebSocketConnections) + close(s.WebSocketFrames) + close(s.LootEntries) + close(s.PluginEntries) + close(s.Intercept.queue) + return s.db.Close() +} + +func (s *Store) write(writeFunc func() error) error { + errorChan := make(chan error, 1) + select { + case s.writeCh <- func() error { + err := writeFunc() + errorChan <- err + return err + }: + case <-s.done: + return fmt.Errorf("store: closed") + } + select { + case err := <-errorChan: + return err + case <-s.done: + return fmt.Errorf("store: closed") + } +} + +func (s *Store) writeLoop() { + for { + select { + case fn := <-s.writeCh: + _ = fn() + case <-s.done: + + for { + select { + case fn := <-s.writeCh: + _ = fn() + default: + return + } + } + } + } +} + +func (s *Store) migrate() error { + _, err := s.db.Exec(schema) + if err != nil { + return fmt.Errorf("store: migrate: %w", err) + } + // Add new columns to existing tables if they don't exist. + // Errors are ignored — column may already exist. + migrations := []string{ + `ALTER TABLE loot ADD COLUMN raw_request TEXT DEFAULT ''`, + `ALTER TABLE loot ADD COLUMN raw_response TEXT DEFAULT ''`, + `ALTER TABLE history ADD COLUMN proto TEXT NOT NULL DEFAULT 'HTTP/1.1'`, + `ALTER TABLE history ADD COLUMN port INTEGER NOT NULL DEFAULT 443`, + `ALTER TABLE repeater_tabs ADD COLUMN tab_type TEXT NOT NULL DEFAULT 'http'`, + `ALTER TABLE history ADD COLUMN resp_size INTEGER NOT NULL DEFAULT 0`, + `CREATE INDEX IF NOT EXISTS idx_history_host_port_id ON history(host, port, id DESC)`, + } + for _, migration := range migrations { + s.db.Exec(migration) + } + return nil +} + +const schema = ` +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + timestamp TEXT NOT NULL, + host TEXT NOT NULL, + port INTEGER NOT NULL DEFAULT 443, + method TEXT NOT NULL, + url TEXT NOT NULL, + proto TEXT NOT NULL DEFAULT 'HTTP/1.1', + req_headers TEXT NOT NULL, + req_body BLOB, + status_code INTEGER, + resp_headers TEXT, + resp_body BLOB, + resp_size INTEGER NOT NULL DEFAULT 0, + duration_ms INTEGER, + tls INTEGER DEFAULT 0, + in_scope INTEGER DEFAULT 1 +); + +CREATE INDEX IF NOT EXISTS idx_history_host ON history(host); +CREATE INDEX IF NOT EXISTS idx_history_timestamp ON history(timestamp); +CREATE INDEX IF NOT EXISTS idx_history_scope ON history(in_scope); +CREATE INDEX IF NOT EXISTS idx_history_dedup ON history(method, host, url, status_code); +CREATE INDEX IF NOT EXISTS idx_history_host_id ON history(host, id DESC); +CREATE INDEX IF NOT EXISTS idx_history_scope_id ON history(in_scope, id DESC); + +CREATE TABLE IF NOT EXISTS repeater_tabs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL DEFAULT 'Tab', + host TEXT NOT NULL, + port INTEGER NOT NULL DEFAULT 80, + tls INTEGER DEFAULT 0, + raw_request TEXT NOT NULL, + last_response TEXT, + position INTEGER DEFAULT 0 +); + +CREATE TABLE IF NOT EXISTS loot ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + title TEXT NOT NULL, + severity TEXT NOT NULL, + notes TEXT, + history_id INTEGER REFERENCES history(id), + raw_request TEXT DEFAULT '', + raw_response TEXT DEFAULT '', + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS scope ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + host TEXT NOT NULL UNIQUE +); + +CREATE TABLE IF NOT EXISTS intruder_config ( + id INTEGER PRIMARY KEY, + delay_ms INTEGER DEFAULT 0, + stop_on_status INTEGER DEFAULT 0, + max_redirects INTEGER DEFAULT 10, + follow_redirects TEXT DEFAULT 'never', + timeout_ms INTEGER DEFAULT 30000, + raw_request TEXT DEFAULT '', + payloads TEXT DEFAULT '' +); + +CREATE TABLE IF NOT EXISTS websocket_connections ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + history_id INTEGER REFERENCES history(id), + host TEXT NOT NULL, + url TEXT NOT NULL, + tls INTEGER DEFAULT 0, + in_scope INTEGER DEFAULT 1, + timestamp TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_ws_connections_host ON websocket_connections(host); +CREATE INDEX IF NOT EXISTS idx_ws_connections_scope ON websocket_connections(in_scope); + +CREATE TABLE IF NOT EXISTS websocket_frames ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + connection_id INTEGER NOT NULL REFERENCES websocket_connections(id) ON DELETE CASCADE, + timestamp TEXT NOT NULL, + direction INTEGER NOT NULL, -- 0 = client→server, 1 = server→client + opcode INTEGER NOT NULL, + payload BLOB +); + +CREATE INDEX IF NOT EXISTS idx_ws_frames_connection ON websocket_frames(connection_id); + +CREATE TABLE IF NOT EXISTS plugins ( + name TEXT PRIMARY KEY, + path TEXT NOT NULL DEFAULT '', + enabled INTEGER NOT NULL DEFAULT 1 +); + +CREATE TABLE IF NOT EXISTS annotations ( + history_id INTEGER PRIMARY KEY REFERENCES history(id) ON DELETE CASCADE, + comment TEXT NOT NULL DEFAULT '', + color TEXT NOT NULL DEFAULT '' +); +` diff --git a/internal/store/repeater.go b/internal/store/repeater.go new file mode 100644 index 0000000..b433559 --- /dev/null +++ b/internal/store/repeater.go @@ -0,0 +1,103 @@ +package store + +import ( + "fmt" +) + +// RepeaterTab represents a saved repeater tab. +type RepeaterTab struct { + ID int64 + Name string + Host string + Port int + TLS bool + TabType string // "http" or "websocket" + RawRequest string + LastResponse string + Position int +} + +// AllRepeaterTabs returns all saved repeater tabs ordered by position. +func (s *Store) AllRepeaterTabs() ([]RepeaterTab, error) { + rows, err := s.db.Query(` + SELECT id, name, host, port, tls, COALESCE(tab_type, 'http'), raw_request, COALESCE(last_response, ''), position + FROM repeater_tabs + ORDER BY position ASC`) + if err != nil { + return nil, fmt.Errorf("store: query repeater tabs: %w", err) + } + defer rows.Close() + + var tabs []RepeaterTab + for rows.Next() { + var tab RepeaterTab + var tlsFlag int + if err := rows.Scan(&tab.ID, &tab.Name, &tab.Host, &tab.Port, &tlsFlag, &tab.TabType, &tab.RawRequest, &tab.LastResponse, &tab.Position); err != nil { + return nil, fmt.Errorf("store: scan repeater tab: %w", err) + } + tab.TLS = tlsFlag == 1 + if tab.TabType == "" { + tab.TabType = "http" + } + tabs = append(tabs, tab) + } + return tabs, rows.Err() +} + +// SaveRepeaterTab inserts a new repeater tab and returns its ID. +func (s *Store) SaveRepeaterTab(tab RepeaterTab) (int64, error) { + tabType := tab.TabType + if tabType == "" { + tabType = "http" + } + var id int64 + err := s.write(func() error { + result, err := s.db.Exec(` + INSERT INTO repeater_tabs (name, host, port, tls, tab_type, raw_request, last_response, position) + VALUES (?, ?, ?, ?, ?, ?, ?, (SELECT COALESCE(MAX(position), 0) + 1 FROM repeater_tabs))`, + tab.Name, tab.Host, tab.Port, boolToInt(tab.TLS), tabType, tab.RawRequest, tab.LastResponse, + ) + if err != nil { + return fmt.Errorf("store: save repeater tab: %w", err) + } + id, err = result.LastInsertId() + return err + }) + return id, err +} + +// UpdateRepeaterTab updates the request and response for an existing tab. +func (s *Store) UpdateRepeaterTab(id int64, rawRequest, lastResponse string) error { + return s.write(func() error { + _, err := s.db.Exec(` + UPDATE repeater_tabs SET raw_request = ?, last_response = ? WHERE id = ?`, + rawRequest, lastResponse, id, + ) + if err != nil { + return fmt.Errorf("store: update repeater tab: %w", err) + } + return nil + }) +} + +// RenameRepeaterTab updates the name of a tab. +func (s *Store) RenameRepeaterTab(id int64, name string) error { + return s.write(func() error { + _, err := s.db.Exec(`UPDATE repeater_tabs SET name = ? WHERE id = ?`, name, id) + if err != nil { + return fmt.Errorf("store: rename repeater tab: %w", err) + } + return nil + }) +} + +// DeleteRepeaterTab removes a repeater tab by ID. +func (s *Store) DeleteRepeaterTab(id int64) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM repeater_tabs WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete repeater tab: %w", err) + } + return nil + }) +} diff --git a/internal/store/repeater_test.go b/internal/store/repeater_test.go new file mode 100644 index 0000000..c85ad83 --- /dev/null +++ b/internal/store/repeater_test.go @@ -0,0 +1,265 @@ +package store_test + +import ( + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeTab(name, host string, port int, tls bool) store.RepeaterTab { + return store.RepeaterTab{ + Name: name, + Host: host, + Port: port, + TLS: tls, + TabType: "http", + RawRequest: "GET / HTTP/1.1\r\nHost: " + host + "\r\n\r\n", + } +} + +func makeWSTab(name, url string, tls bool) store.RepeaterTab { + return store.RepeaterTab{ + Name: name, + Host: url, + Port: 0, + TLS: tls, + TabType: "websocket", + RawRequest: url, + } +} + +// ── HTTP tab ────────────────────────────────────────────────────────────────── + +func TestSaveRepeaterTab_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + tab := store.RepeaterTab{ + Name: "My Tab", + Host: "example.com", + Port: 443, + TLS: true, + TabType: "http", + RawRequest: "GET /api HTTP/1.1\r\nHost: example.com\r\n\r\n", + } + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + + got := tabs[0] + assert.Equal(t, id, got.ID) + assert.Equal(t, "My Tab", got.Name) + assert.Equal(t, "example.com", got.Host) + assert.Equal(t, 443, got.Port) + assert.True(t, got.TLS) + assert.Equal(t, "http", got.TabType) + assert.Equal(t, tab.RawRequest, got.RawRequest) +} + +func TestSaveRepeaterTab_TLSFalse(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("HTTP Tab", "example.com", 80, false)) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.False(t, tabs[0].TLS) + assert.Equal(t, "http", tabs[0].TabType) +} + +func TestSaveRepeaterTab_DefaultTabType(t *testing.T) { + st := newTestStore(t) + + // TabType omitted — should default to "http". + tab := store.RepeaterTab{ + Name: "No Type", + Host: "example.com", + Port: 80, + RawRequest: "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n", + } + _, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "http", tabs[0].TabType) +} + +func TestAllRepeaterTabs_OrderedByPosition(t *testing.T) { + st := newTestStore(t) + + for _, name := range []string{"First", "Second", "Third"} { + _, err := st.SaveRepeaterTab(makeTab(name, "example.com", 80, false)) + require.NoError(t, err) + } + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 3) + assert.Equal(t, "First", tabs[0].Name) + assert.Equal(t, "Second", tabs[1].Name) + assert.Equal(t, "Third", tabs[2].Name) +} + +func TestAllRepeaterTabs_PositionAutoIncrements(t *testing.T) { + st := newTestStore(t) + + _, err := st.SaveRepeaterTab(makeTab("A", "a.com", 80, false)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeTab("B", "b.com", 80, false)) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 2) + assert.Less(t, tabs[0].Position, tabs[1].Position) +} + +func TestUpdateRepeaterTab_UpdatesRequestAndResponse(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Tab", "example.com", 80, false)) + require.NoError(t, err) + + newReq := "POST /submit HTTP/1.1\r\nHost: example.com\r\n\r\nbody" + newResp := "HTTP/1.1 200 OK\r\n\r\nok" + require.NoError(t, st.UpdateRepeaterTab(id, newReq, newResp)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, newReq, tabs[0].RawRequest) + assert.Equal(t, newResp, tabs[0].LastResponse) +} + +func TestUpdateRepeaterTab_NonExistentIDNoError(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.UpdateRepeaterTab(9999, "req", "resp")) +} + +func TestRenameRepeaterTab(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Old", "example.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.RenameRepeaterTab(id, "New Name")) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "New Name", tabs[0].Name) +} + +func TestDeleteRepeaterTab(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Tab", "example.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(id)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + assert.Empty(t, tabs) +} + +func TestDeleteRepeaterTab_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteRepeaterTab(9999)) +} + +func TestDeleteRepeaterTab_DoesNotAffectOtherTabs(t *testing.T) { + st := newTestStore(t) + + idKeep, err := st.SaveRepeaterTab(makeTab("Keep", "keep.com", 80, false)) + require.NoError(t, err) + idDel, err := st.SaveRepeaterTab(makeTab("Delete", "del.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(idDel)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, idKeep, tabs[0].ID) + assert.Equal(t, "Keep", tabs[0].Name) +} + +// ── WebSocket tab ───────────────────────────────────────────────────────────── + +func TestSaveRepeaterTab_WebSocketType(t *testing.T) { + st := newTestStore(t) + + tab := makeWSTab("WS Tab", "wss://echo.websocket.org:443/", true) + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "websocket", tabs[0].TabType) + assert.Equal(t, "wss://echo.websocket.org:443/", tabs[0].RawRequest) + assert.True(t, tabs[0].TLS) +} + +func TestAllRepeaterTabs_MixedTypes(t *testing.T) { + st := newTestStore(t) + + _, err := st.SaveRepeaterTab(makeTab("HTTP", "example.com", 80, false)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeWSTab("WS", "wss://example.com/ws", true)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeTab("HTTP2", "other.com", 443, true)) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 3) + assert.Equal(t, "http", tabs[0].TabType) + assert.Equal(t, "websocket", tabs[1].TabType) + assert.Equal(t, "http", tabs[2].TabType) +} + +func TestSaveRepeaterTab_WebSocketPreservesURL(t *testing.T) { + st := newTestStore(t) + + url := "wss://api.example.com:8443/realtime?token=abc123" + tab := makeWSTab("API WS", url, true) + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + + full, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, full, 1) + assert.Equal(t, id, full[0].ID) + assert.Equal(t, url, full[0].RawRequest) + assert.Equal(t, "websocket", full[0].TabType) +} + +func TestDeleteRepeaterTab_WebSocket(t *testing.T) { + st := newTestStore(t) + + httpID, err := st.SaveRepeaterTab(makeTab("HTTP", "example.com", 80, false)) + require.NoError(t, err) + wsID, err := st.SaveRepeaterTab(makeWSTab("WS", "wss://example.com/ws", true)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(wsID)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, httpID, tabs[0].ID) + assert.Equal(t, "http", tabs[0].TabType) +} diff --git a/internal/store/scope.go b/internal/store/scope.go new file mode 100644 index 0000000..8c7c4a0 --- /dev/null +++ b/internal/store/scope.go @@ -0,0 +1,116 @@ +package store + +import ( + "fmt" + "net" + "strings" +) + +// ScopeEntry represents a single scope target. +type ScopeEntry struct { + ID int64 + Host string +} + +// AllScopeEntries returns all scope entries. +func (s *Store) AllScopeEntries() ([]ScopeEntry, error) { + rows, err := s.db.Query(`SELECT id, host FROM scope ORDER BY id ASC`) + if err != nil { + return nil, fmt.Errorf("store: query scope: %w", err) + } + defer rows.Close() + + var entries []ScopeEntry + for rows.Next() { + var e ScopeEntry + if err := rows.Scan(&e.ID, &e.Host); err != nil { + return nil, fmt.Errorf("store: scan scope: %w", err) + } + entries = append(entries, e) + } + return entries, rows.Err() +} + +// AddScopeEntry adds a new host to scope. +func (s *Store) AddScopeEntry(host string) error { + err := s.write(func() error { + _, err := s.db.Exec(`INSERT INTO scope (host) VALUES (?)`, host) + if err != nil { + return fmt.Errorf("store: add scope entry: %w", err) + } + return nil + }) + if err == nil { + s.invalidateScopeCache() + } + return err +} + +// DeleteScopeEntry removes a scope entry by ID. +func (s *Store) DeleteScopeEntry(id int64) error { + err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM scope WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete scope entry: %w", err) + } + return nil + }) + if err == nil { + s.invalidateScopeCache() + } + return err +} + +// InScope returns true if the given host matches any scope entry. +// example.com matches example.com and all subdomains. +// Uses an in-memory cache populated on first call and invalidated on scope writes. +func (s *Store) InScope(host string) bool { + // Strip port if present; net.SplitHostPort handles IPv6 brackets correctly. + if h, _, err := net.SplitHostPort(host); err == nil { + host = h + } + + s.scopeCacheMu.RLock() + if s.scopeCacheSet { + entries := s.scopeCache + s.scopeCacheMu.RUnlock() + return matchesScope(entries, host) + } + s.scopeCacheMu.RUnlock() + + // Cache miss — populate under write lock. + s.scopeCacheMu.Lock() + // Re-check after acquiring write lock — another goroutine may have populated it. + if !s.scopeCacheSet { + entries, err := s.AllScopeEntries() + if err == nil { + s.scopeCache = entries + s.scopeCacheSet = true + } + } + entries := s.scopeCache + s.scopeCacheMu.Unlock() + + return matchesScope(entries, host) +} + +// invalidateScopeCache clears the scope cache so the next InScope call reloads from DB. +func (s *Store) invalidateScopeCache() { + s.scopeCacheMu.Lock() + s.scopeCache = nil + s.scopeCacheSet = false + s.scopeCacheMu.Unlock() +} + +// matchesScope returns true if host matches any entry in entries. +func matchesScope(entries []ScopeEntry, host string) bool { + if len(entries) == 0 { + return true // no scope defined, everything is in scope + } + for _, e := range entries { + if host == e.Host || strings.HasSuffix(host, "."+e.Host) { + return true + } + } + return false +} diff --git a/internal/store/scope_test.go b/internal/store/scope_test.go new file mode 100644 index 0000000..b1267c8 --- /dev/null +++ b/internal/store/scope_test.go @@ -0,0 +1,135 @@ +package store_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAddScopeEntry_AndRetrieve(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "example.com", entries[0].Host) + assert.Positive(t, entries[0].ID) +} + +func TestAddScopeEntry_Multiple(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + require.NoError(t, st.AddScopeEntry("other.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + assert.Len(t, entries, 2) +} + +func TestAddScopeEntry_Duplicate_Errors(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + err := st.AddScopeEntry("example.com") + assert.Error(t, err, "duplicate host must be rejected by UNIQUE constraint") +} + +func TestDeleteScopeEntry(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 1) + + require.NoError(t, st.DeleteScopeEntry(entries[0].ID)) + + entries, err = st.AllScopeEntries() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestDeleteScopeEntry_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteScopeEntry(9999)) +} + +func TestAllScopeEntries_OrderedByIDAsc(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("alpha.com")) + require.NoError(t, st.AddScopeEntry("beta.com")) + require.NoError(t, st.AddScopeEntry("gamma.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 3) + assert.Equal(t, "alpha.com", entries[0].Host) + assert.Equal(t, "beta.com", entries[1].Host) + assert.Equal(t, "gamma.com", entries[2].Host) +} + +func TestInScope_EmptyScopeAllowsEverything(t *testing.T) { + st := newTestStore(t) + assert.True(t, st.InScope("anything.com")) + assert.True(t, st.InScope("other.net")) +} + +func TestInScope_ExactMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("example.com")) +} + +func TestInScope_SubdomainMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("sub.example.com")) + assert.True(t, st.InScope("deep.sub.example.com")) +} + +func TestInScope_NoMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.False(t, st.InScope("other.com")) + assert.False(t, st.InScope("notexample.com")) + assert.False(t, st.InScope("example.com.evil.com")) +} + +func TestInScope_StripsPort(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("example.com:8080")) + assert.True(t, st.InScope("example.com:443")) +} + +func TestInScope_AfterDeleteScopeBecomesOpen(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.NoError(t, st.DeleteScopeEntry(entries[0].ID)) + + // scope is now empty — everything is in scope + assert.True(t, st.InScope("example.com")) + assert.True(t, st.InScope("completely-unrelated.com")) +} + +func TestInScope_MultipleEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + require.NoError(t, st.AddScopeEntry("target.io")) + + assert.True(t, st.InScope("example.com")) + assert.True(t, st.InScope("api.target.io")) + assert.False(t, st.InScope("other.com")) +} diff --git a/internal/store/search_test.go b/internal/store/search_test.go new file mode 100644 index 0000000..be21b3e --- /dev/null +++ b/internal/store/search_test.go @@ -0,0 +1,105 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func logTxWithBodies(t *testing.T, st *store.Store, method, host, url string, status int, reqBody, respBody string) { + t.Helper() + tx := store.Transaction{ + Timestamp: time.Now().UTC(), + Host: host, + Method: method, + URL: url, + ReqBody: []byte(reqBody), + RespBody: []byte(respBody), + StatusCode: status, + InScope: true, + } + require.NoError(t, st.Log(tx)) +} + +func TestSearchTransactions_FindsInRespBody(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "secret_token_here") + logTxWithBodies(t, st, "GET", "b.com", "https://b.com/", 200, "", "nothing special") + + results, err := st.SearchTransactions("secret_token", false, true, false) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "a.com", results[0].Host) +} + +func TestSearchTransactions_FindsInReqBody(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "POST", "a.com", "https://a.com/login", 200, "password=hunter2", "ok") + logTxWithBodies(t, st, "GET", "b.com", "https://b.com/", 200, "", "ok") + + results, err := st.SearchTransactions("hunter2", true, false, false) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "a.com", results[0].Host) +} + +func TestSearchTransactions_EmptyTerm_ReturnsNil(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "data") + + results, err := st.SearchTransactions("", true, true, false) + require.NoError(t, err) + assert.Nil(t, results) +} + +func TestSearchTransactions_NoFieldsSelected_ReturnsNil(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "data") + + results, err := st.SearchTransactions("data", false, false, false) + require.NoError(t, err) + assert.Nil(t, results) +} + +func TestSearchTransactions_NoMatch_ReturnsEmpty(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "nothing here") + + results, err := st.SearchTransactions("xyzzy_not_found", false, true, false) + require.NoError(t, err) + assert.Empty(t, results) +} + +func TestTransactionsPage_MethodFilter(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/", 200))) + require.NoError(t, st.Log(makeTx("POST", "a.com", "https://a.com/login", 200))) + require.NoError(t, st.Log(makeTx("PUT", "a.com", "https://a.com/item", 200))) + + results, err := st.TransactionsPage(0, store.TransactionFilter{ + ShowOutScope: true, + Methods: []string{"POST"}, + }) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "POST", results[0].Method) +} + +func TestTransactionsPage_StatusRangeFilter(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/ok", 200))) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/redir", 302))) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/err", 404))) + + results, err := st.TransactionsPage(0, store.TransactionFilter{ + ShowOutScope: true, + StatusMin: 400, + StatusMax: 499, + }) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, 404, results[0].StatusCode) +} diff --git a/internal/store/store_test.go b/internal/store/store_test.go new file mode 100644 index 0000000..6b9d012 --- /dev/null +++ b/internal/store/store_test.go @@ -0,0 +1,22 @@ +package store_test + +import ( + "os" + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/require" +) + +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-store-test-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + + st, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { st.Close() }) + return st +} diff --git a/internal/store/websocket.go b/internal/store/websocket.go new file mode 100644 index 0000000..562c487 --- /dev/null +++ b/internal/store/websocket.go @@ -0,0 +1,198 @@ +package store + +import ( + "fmt" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +const wsFrameSizeLimit = 1 << 20 // 1 MB per frame + +// WebSocketFrame is a single frame captured during a WebSocket session. +type WebSocketFrame struct { + ID uint64 + ConnectionID uint64 // foreign key → websocket_connections.id + Timestamp time.Time + Direction events.WebSocketDirection + Opcode events.WebSocketOpcode + Payload []byte +} + +// WebSocketConnection represents the initial HTTP upgrade that started a +// WebSocket session. It is tied to the history transaction of the upgrade +// request so the UI can link from history to the frame list. +type WebSocketConnection struct { + ID uint64 + HistoryID uint64 + Host string + URL string + TLS bool + InScope bool + Timestamp time.Time + FrameCount int // populated on read, not stored +} + +// LogWebSocketConnection inserts a new WebSocket connection record and returns +// its ID. Called once when the upgrade handshake completes. +func (s *Store) LogWebSocketConnection(conn WebSocketConnection) (uint64, error) { + var id uint64 + err := s.write(func() error { + res, err := s.db.Exec(` + INSERT INTO websocket_connections (history_id, host, url, tls, in_scope, timestamp) + VALUES (?, ?, ?, ?, ?, ?)`, + conn.HistoryID, + conn.Host, + conn.URL, + boolToInt(conn.TLS), + boolToInt(conn.InScope), + conn.Timestamp.UTC().Format(time.RFC3339), + ) + if err != nil { + return fmt.Errorf("store: log websocket connection: %w", err) + } + lastID, err := res.LastInsertId() + if err != nil { + return fmt.Errorf("store: websocket connection last insert id: %w", err) + } + id = uint64(lastID) + return nil + }) + if err != nil { + return 0, err + } + + // Notify UI listeners — non-blocking, same pattern as Updates. + conn.ID = id + select { + case s.WebSocketConnections <- conn: + default: + } + + return id, nil +} + +// LogWebSocketFrame appends a single frame to an existing connection. +func (s *Store) LogWebSocketFrame(frame WebSocketFrame) error { + err := s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO websocket_frames (connection_id, timestamp, direction, opcode, payload) + VALUES (?, ?, ?, ?, ?)`, + frame.ConnectionID, + frame.Timestamp.UTC().Format(time.RFC3339), + int(frame.Direction), + int(frame.Opcode), + frame.Payload, + ) + if err != nil { + return fmt.Errorf("store: log websocket frame: %w", err) + } + return nil + }) + if err != nil { + return err + } + + // Notify UI listeners — non-blocking. + select { + case s.WebSocketFrames <- frame: + default: + } + + return nil +} + +// AllWebSocketConnections returns all WebSocket connections ordered by id +// descending, with frame counts. +func (s *Store) AllWebSocketConnections() ([]WebSocketConnection, error) { + rows, err := s.db.Query(` + SELECT c.id, c.history_id, c.host, c.url, c.tls, c.in_scope, c.timestamp, + COUNT(f.id) as frame_count + FROM websocket_connections c + LEFT JOIN websocket_frames f ON f.connection_id = c.id + GROUP BY c.id + ORDER BY c.id DESC`) + if err != nil { + return nil, fmt.Errorf("store: query websocket connections: %w", err) + } + defer rows.Close() + + var conns []WebSocketConnection + for rows.Next() { + var c WebSocketConnection + var ts string + var tlsInt, scopeInt int + if err := rows.Scan(&c.ID, &c.HistoryID, &c.Host, &c.URL, &tlsInt, &scopeInt, &ts, &c.FrameCount); err != nil { + return nil, fmt.Errorf("store: scan websocket connection: %w", err) + } + c.Timestamp, _ = time.Parse(time.RFC3339, ts) + c.TLS = tlsInt == 1 + c.InScope = scopeInt == 1 + conns = append(conns, c) + } + return conns, rows.Err() +} + +// FramesForConnection returns all frames for a given connection ordered by id +// ascending (chronological). +func (s *Store) FramesForConnection(connectionID uint64) ([]WebSocketFrame, error) { + rows, err := s.db.Query(` + SELECT id, connection_id, timestamp, direction, opcode, payload + FROM websocket_frames + WHERE connection_id = ? + ORDER BY id ASC`, connectionID) + if err != nil { + return nil, fmt.Errorf("store: query websocket frames: %w", err) + } + defer rows.Close() + + var frames []WebSocketFrame + for rows.Next() { + var f WebSocketFrame + var ts string + if err := rows.Scan(&f.ID, &f.ConnectionID, &ts, &f.Direction, &f.Opcode, &f.Payload); err != nil { + return nil, fmt.Errorf("store: scan websocket frame: %w", err) + } + f.Timestamp, _ = time.Parse(time.RFC3339, ts) + frames = append(frames, f) + } + return frames, rows.Err() +} + +// ObserveWebSocketConnection implements events.WebSocketConnectionObserver. +func (s *Store) ObserveWebSocketConnection(e events.WebSocketConnectionEvent) uint64 { + id, err := s.LogWebSocketConnection(WebSocketConnection{ + Host: e.Host, + URL: e.URL, + TLS: e.TLS, + InScope: s.InScope(e.Host), + Timestamp: e.Timestamp, + }) + if err != nil { + logger.Error("store: observe ws connection: %v", err) + return 0 + } + return id +} + +// ObserveWebSocketFrame implements events.WebSocketFrameObserver. +func (s *Store) ObserveWebSocketFrame(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + if e.ConnectionID == 0 { + return events.WebSocketFrameResult{Payload: e.Payload} + } + payload := e.Payload + if len(payload) > wsFrameSizeLimit { + payload = payload[:wsFrameSizeLimit] + } + if err := s.LogWebSocketFrame(WebSocketFrame{ + ConnectionID: e.ConnectionID, + Timestamp: e.Timestamp, + Direction: e.Direction, + Opcode: e.Opcode, + Payload: payload, + }); err != nil { + logger.Error("store: observe ws frame: %v", err) + } + return events.WebSocketFrameResult{Payload: e.Payload} +} diff --git a/internal/store/websocket_test.go b/internal/store/websocket_test.go new file mode 100644 index 0000000..6340080 --- /dev/null +++ b/internal/store/websocket_test.go @@ -0,0 +1,200 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLogWebSocketConnection_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + conn := store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/chat", + TLS: true, + InScope: true, + Timestamp: time.Now().UTC().Truncate(time.Second), + } + id, err := st.LogWebSocketConnection(conn) + require.NoError(t, err) + assert.NotZero(t, id) + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 1) + assert.Equal(t, id, conns[0].ID) + assert.Equal(t, "example.com:443", conns[0].Host) + assert.Equal(t, "wss://example.com:443/chat", conns[0].URL) + assert.True(t, conns[0].TLS) + assert.True(t, conns[0].InScope) +} + +func TestLogWebSocketFrame_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + TLS: true, + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + frame := store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC().Truncate(time.Second), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: []byte("hello world"), + } + require.NoError(t, st.LogWebSocketFrame(frame)) + + frames, err := st.FramesForConnection(connID) + require.NoError(t, err) + require.Len(t, frames, 1) + assert.Equal(t, connID, frames[0].ConnectionID) + assert.Equal(t, events.WebSocketClient, frames[0].Direction) + assert.Equal(t, events.WebSocketText, frames[0].Opcode) + assert.Equal(t, []byte("hello world"), frames[0].Payload) +} + +func TestAllWebSocketConnections_OrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + for _, host := range []string{"first.com:443", "second.com:443", "third.com:443"} { + _, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: host, + URL: "wss://" + host + "/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + } + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 3) + assert.Equal(t, "third.com:443", conns[0].Host) + assert.Equal(t, "second.com:443", conns[1].Host) + assert.Equal(t, "first.com:443", conns[2].Host) +} + +func TestAllWebSocketConnections_FrameCountIsCorrect(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + for i := 0; i < 5; i++ { + require.NoError(t, st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: []byte("msg"), + })) + } + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 1) + assert.Equal(t, 5, conns[0].FrameCount) +} + +func TestFramesForConnection_OrderedByIDAsc(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + directions := []events.WebSocketDirection{ + events.WebSocketClient, + events.WebSocketServer, + events.WebSocketClient, + } + for _, dir := range directions { + require.NoError(t, st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: dir, + Opcode: events.WebSocketText, + Payload: []byte("x"), + })) + } + + frames, err := st.FramesForConnection(connID) + require.NoError(t, err) + require.Len(t, frames, 3) + assert.Equal(t, events.WebSocketClient, frames[0].Direction) + assert.Equal(t, events.WebSocketServer, frames[1].Direction) + assert.Equal(t, events.WebSocketClient, frames[2].Direction) +} + +func TestFramesForConnection_EmptyForUnknownConnection(t *testing.T) { + st := newTestStore(t) + frames, err := st.FramesForConnection(99999) + require.NoError(t, err) + assert.Empty(t, frames) +} + +func TestLogWebSocketConnection_EmitsToChannel(t *testing.T) { + st := newTestStore(t) + + conn := store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + } + _, err := st.LogWebSocketConnection(conn) + require.NoError(t, err) + + select { + case c := <-st.WebSocketConnections: + assert.Equal(t, "example.com:443", c.Host) + case <-time.After(time.Second): + t.Fatal("timed out waiting for WebSocketConnections channel emit") + } +} + +func TestLogWebSocketFrame_EmitsToChannel(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + // Drain the connection channel. + <-st.WebSocketConnections + + err = st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketText, + Payload: []byte("streamed"), + }) + require.NoError(t, err) + + select { + case f := <-st.WebSocketFrames: + assert.Equal(t, connID, f.ConnectionID) + assert.Equal(t, events.WebSocketServer, f.Direction) + assert.Equal(t, []byte("streamed"), f.Payload) + case <-time.After(time.Second): + t.Fatal("timed out waiting for WebSocketFrames channel emit") + } +} diff --git a/internal/ui/app.go b/internal/ui/app.go new file mode 100644 index 0000000..27d6f9a --- /dev/null +++ b/internal/ui/app.go @@ -0,0 +1,239 @@ +package ui + +import ( + "fmt" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/canvas" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/data/binding" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/storage" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/assets" + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" +) + +const ( + prefKeyDarkTheme = "dark_theme" + prefKeyProxyHost = "proxy_host" + prefKeyProxyPort = "proxy_port" + prefKeyProxyEnabled = "proxy_enabled" + prefKeyUserTheme = "user_theme" + + defaultDarkTheme = true + defaultProxyHost = "127.0.0.1" + defaultProxyPort = 9090 + defaultProxyEnabled = true +) + +// Package-level bindings shared between app.go and settings.go. +// Both files are in the same package so this is safe — there is only +// ever one app instance in a desktop process. +var proxyStatusBinding = binding.NewString() +var proxyRunningBinding = binding.NewBool() + +// loadActiveTheme reads the user theme from prefs and loads it. +// Returns nil if no user theme is set or loading fails (falls back to default). +func loadActiveTheme(fyneApp fyne.App) *LoadedTheme { + name := fyneApp.Preferences().String(prefKeyUserTheme) + if name == "" { + return nil + } + dir := ThemesDir(fyneApp.Storage().RootURI().Path()) + lt, err := loadThemeByName(name, dir) + if err != nil { + return nil + } + return lt +} + +// applyTheme reads current dark/light and user theme prefs and applies the theme. +// Call this from anywhere — toggle, settings, startup. No shared mutable state. +func applyTheme(fyneApp fyne.App) { + isDark := fyneApp.Preferences().BoolWithFallback(prefKeyDarkTheme, defaultDarkTheme) + fyneApp.Settings().SetTheme(NewVagueThemeWithLoaded(isDark, loadActiveTheme(fyneApp))) +} + +// setProxyStatus updates both package-level bindings and persists the running +// state to prefs. Must be called on the Fyne main goroutine. +func setProxyStatus(prefs fyne.Preferences, running bool) { + proxyHost := prefs.StringWithFallback(prefKeyProxyHost, defaultProxyHost) + proxyPort := prefs.IntWithFallback(prefKeyProxyPort, defaultProxyPort) + prefs.SetBool(prefKeyProxyEnabled, running) + _ = proxyRunningBinding.Set(running) + if running { + _ = proxyStatusBinding.Set(fmt.Sprintf("Proxy: running on %s:%d", proxyHost, proxyPort)) + } else { + _ = proxyStatusBinding.Set("Proxy: stopped") + } +} + +func ShowMainWindow(fyneApp fyne.App, projectStore *store.Store, bus *events.Bus, launchWin fyne.Window, onSwitchProject func(string)) { + prefs := fyneApp.Preferences() + + mainWin := fyneApp.NewWindow("Shiv") + mainWin.Resize(fyne.NewSize(1280, 800)) + mainWin.SetIcon(fyne.NewStaticResource("logo.png", assets.Logo)) + + mainWin.SetMaster() + applyTheme(fyneApp) + + browserBtn := widget.NewButtonWithIcon("", AppIcon("web"), func() { + launchDefaultBrowser(fyneApp, mainWin) + }) + + proxyToggleBtn := widget.NewButtonWithIcon("", AppIcon("off-button"), nil) + + proxyRunningBinding.AddListener(binding.NewDataListener(func() { + running, _ := proxyRunningBinding.Get() + if running { + proxyToggleBtn.SetIcon(AppIcon("on-button")) + } else { + proxyToggleBtn.SetIcon(AppIcon("off-button")) + } + proxyToggleBtn.Refresh() + })) + + proxyToggleBtn.OnTapped = func() { + isRunning, _ := proxyRunningBinding.Get() + isRunning = !isRunning + + if isRunning { + proxyHost := prefs.StringWithFallback(prefKeyProxyHost, defaultProxyHost) + proxyPort := prefs.IntWithFallback(prefKeyProxyPort, defaultProxyPort) + addr := fmt.Sprintf("%s:%d", proxyHost, proxyPort) + bus.EmitProxyRestart(events.ProxyRestartEvent{Addr: addr}) + } else { + bus.EmitProxyStop(events.ProxyStopEvent{}) + } + + setProxyStatus(prefs, isRunning) + } + + // Initialise status from persisted prefs. + setProxyStatus(prefs, prefs.BoolWithFallback(prefKeyProxyEnabled, defaultProxyEnabled)) + + proxyLabel := widget.NewLabelWithData(proxyStatusBinding) + proxyLabel.Importance = widget.LowImportance + + logo := canvas.NewImageFromResource(fyne.NewStaticResource("logo.png", assets.Logo)) + logo.FillMode = canvas.ImageFillContain + logo.SetMinSize(fyne.NewSize(24, 24)) + + appName := widget.NewLabel("Shiv") + appName.TextStyle = fyne.TextStyle{Bold: true} + + repeater := newRepeaterTab(projectStore, bus, mainWin) + loot := newLootTab(projectStore, mainWin, repeater) + repeater.loot = loot + intruder := newIntruderTab(mainWin, projectStore, repeater, loot) + history := newHistoryTab(projectStore, mainWin, repeater, loot, intruder) + intercept := newInterceptTab(projectStore) + plugins := newPluginsTab(projectStore, bus, mainWin) + decoder := newDecoderTab() + tabs := container.NewAppTabs( + container.NewTabItemWithIcon("History", AppIcon("history"), history.build()), + container.NewTabItemWithIcon("Intercept", AppIcon("intercept"), intercept.build()), + container.NewTabItemWithIcon("Repeater", AppIcon("repeater"), repeater.build()), + container.NewTabItemWithIcon("Intruder", AppIcon("intruder"), intruder.build()), + container.NewTabItemWithIcon("Decoder", AppIcon("toolbox"), decoder.build()), + container.NewTabItemWithIcon("Loot", AppIcon("loot"), loot.build()), + container.NewTabItemWithIcon("Plugins", AppIcon("plugin"), plugins.build()), + ) + tabs.SetTabLocation(container.TabLocationTop) + + keybinds := newKeybinds(mainWin, tabs, history, repeater, prefs) + + settingsMenuItem := fyne.NewMenuItem("Settings", func() { + showSettingsDialog(fyneApp, mainWin, bus, keybinds) + }) + settingsMenuItem.Icon = AppIcon("toolbox") + + toggleThemeMenuItem := fyne.NewMenuItem("Toggle Theme", func() { + isDark := !prefs.BoolWithFallback(prefKeyDarkTheme, defaultDarkTheme) + prefs.SetBool(prefKeyDarkTheme, isDark) + applyTheme(fyneApp) + }) + toggleThemeMenuItem.Icon = theme.ColorChromaticIcon() + launchBrowserMenuItem := fyne.NewMenuItem("Launch Browser", func() { + launchDefaultBrowser(fyneApp, mainWin) + }) + + newProjectItem := fyne.NewMenuItem("New Project", func() { + d := dialog.NewFileSave(func(wc fyne.URIWriteCloser, err error) { + if err != nil || wc == nil { + return + } + wc.Close() + path := wc.URI().Path() + saveRecent(path) + mainWin.Hide() + onSwitchProject(path) + }, mainWin) + d.SetFileName("untitled.shiv") + d.SetFilter(storage.NewExtensionFileFilter([]string{".shiv"})) + d.Show() + }) + newProjectItem.Icon = theme.DocumentCreateIcon() + + openProjectItem := fyne.NewMenuItem("Open Project", func() { + d := dialog.NewFileOpen(func(rc fyne.URIReadCloser, err error) { + if err != nil || rc == nil { + return + } + rc.Close() + path := rc.URI().Path() + saveRecent(path) + mainWin.Hide() + onSwitchProject(path) + }, mainWin) + d.SetFilter(storage.NewExtensionFileFilter([]string{".shiv"})) + d.Show() + }) + openProjectItem.Icon = theme.FolderOpenIcon() + + quitItem := fyne.NewMenuItem("Quit", func() { + fyneApp.Quit() + }) + + menuItems := []*fyne.MenuItem{ + newProjectItem, + openProjectItem, + fyne.NewMenuItemSeparator(), + launchBrowserMenuItem, + toggleThemeMenuItem, + settingsMenuItem, + fyne.NewMenuItemSeparator(), + quitItem, + } + if lt := loadActiveTheme(fyneApp); lt != nil && !lt.HasBoth { + toggleThemeMenuItem.Disabled = true + } + + menu := fyne.NewMenu("", menuItems...) + var menuBtn *widget.Button + + menuBtn = widget.NewButtonWithIcon("", theme.MenuIcon(), func() { + widget.ShowPopUpMenuAtRelativePosition(menu, mainWin.Canvas(), fyne.NewPos(0, menuBtn.Size().Height), menuBtn) + }) + + optionsBar := container.NewBorder( + container.NewBorder(nil, nil, nil, + container.NewHBox(proxyLabel, proxyToggleBtn, browserBtn, menuBtn), + ), + nil, nil, nil, + ) + + mainWin.SetContent(container.NewStack( + tabs, + optionsBar, + )) + mainWin.Show() + if launchWin != nil { + launchWin.Close() + } +} diff --git a/internal/ui/browser.go b/internal/ui/browser.go new file mode 100644 index 0000000..2dddb0f --- /dev/null +++ b/internal/ui/browser.go @@ -0,0 +1,283 @@ +package ui + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/dialog" + + "github.com/shiv/internal/logger" +) + +const prefKeyDefaultBrowser = "default_browser" + +// DetectedBrowser holds a browser's display name and executable path. +type DetectedBrowser struct { + Name string + Path string +} + +// isFirefox returns true if the executable looks like Firefox. +func isFirefox(path string) bool { + base := strings.ToLower(filepath.Base(path)) + return strings.Contains(base, "firefox") +} + +// DetectBrowsers scans the system for supported browsers and returns all found. +func DetectBrowsers() []DetectedBrowser { + switch runtime.GOOS { + case "linux": + return detectBrowsersLinux() + case "darwin": + return detectBrowsersDarwin() + case "windows": + return detectBrowsersWindows() + } + return nil +} + +func detectBrowsersLinux() []DetectedBrowser { + candidates := []struct{ name, bin string }{ + {"Google Chrome", "google-chrome"}, + {"Google Chrome (stable)", "google-chrome-stable"}, + {"Chromium", "chromium"}, + {"Chromium (browser)", "chromium-browser"}, + {"Brave (Browser)", "brave-browser"}, + {"Brave", "brave"}, + {"Microsoft Edge", "microsoft-edge"}, + {"Vivaldi", "vivaldi"}, + {"Vivaldi (stable)", "vivaldi-Stable"}, + {"Opera", "opera"}, + {"Opera (stable)", "opera-stable"}, + {"Opera (browser)", "opera-browser"}, + {"Firefox", "firefox"}, + {"Firefox ESR", "firefox-esr"}, + } + var found []DetectedBrowser + seen := map[string]bool{} + for _, c := range candidates { + path, err := exec.LookPath(c.bin) + if err != nil { + continue + } + if seen[path] { + continue + } + seen[path] = true + found = append(found, DetectedBrowser{Name: c.name, Path: path}) + } + return found +} + +func detectBrowsersDarwin() []DetectedBrowser { + candidates := []struct{ name, app string }{ + {"Google Chrome", "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"}, + {"Chromium", "/Applications/Chromium.app/Contents/MacOS/Chromium"}, + {"Brave", "/Applications/Brave Browser.app/Contents/MacOS/Brave Browser"}, + {"Microsoft Edge", "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"}, + {"Vivaldi", "/Applications/Vivaldi.app/Contents/MacOS/Vivaldi"}, + {"Opera", "/Applications/Opera.app/Contents/MacOS/Opera"}, + {"Firefox", "/Applications/Firefox.app/Contents/MacOS/firefox"}, + } + var found []DetectedBrowser + for _, c := range candidates { + if _, err := os.Stat(c.app); err == nil { + found = append(found, DetectedBrowser{Name: c.name, Path: c.app}) + } + } + return found +} + +func detectBrowsersWindows() []DetectedBrowser { + candidates := []struct{ name, path string }{ + {"Google Chrome", `C:\Program Files\Google\Chrome\Application\chrome.exe`}, + {"Google Chrome (x86)", `C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`}, + {"Chromium", `C:\Program Files\Chromium\Application\chrome.exe`}, + {"Brave", `C:\Program Files\BraveSoftware\Brave-Browser\Application\brave.exe`}, + {"Microsoft Edge", `C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`}, + {"Vivaldi", `C:\Program Files\Vivaldi\Application\vivaldi.exe`}, + {"Opera", `C:\Program Files\Opera\opera.exe`}, + {"Firefox", `C:\Program Files\Mozilla Firefox\firefox.exe`}, + {"Firefox (x86)", `C:\Program Files (x86)\Mozilla Firefox\firefox.exe`}, + } + var found []DetectedBrowser + seen := map[string]bool{} + for _, c := range candidates { + if _, err := os.Stat(c.path); err == nil { + if !seen[c.path] { + seen[c.path] = true + found = append(found, DetectedBrowser{Name: c.name, Path: c.path}) + } + } + } + return found +} + +// profileDir returns the persistent browser profile directory for the given browser executable. +func profileDir(fyneApp fyne.App, browserPath string) string { + root := fyneApp.Storage().RootURI().Path() + if isFirefox(browserPath) { + return filepath.Join(root, "browser-profile-firefox") + } + return filepath.Join(root, "browser-profile-chrome") +} + +// setupFirefoxProfile creates a Firefox profile dir with proxy settings written into user.js. +func setupFirefoxProfile(profilePath, proxyHost string, proxyPort int) error { + if err := os.MkdirAll(profilePath, 0700); err != nil { + return fmt.Errorf("create firefox profile dir: %w", err) + } + userJS := fmt.Sprintf(`// Shiv proxy profile — auto-generated, do not edit manually. +user_pref("network.proxy.type", 1); +user_pref("network.proxy.http", "%s"); +user_pref("network.proxy.http_port", %d); +user_pref("network.proxy.ssl", "%s"); +user_pref("network.proxy.ssl_port", %d); +user_pref("network.proxy.no_proxies_on", ""); +user_pref("browser.startup.homepage_override.mstone", "ignore"); +user_pref("browser.shell.checkDefaultBrowser", false); +user_pref("extensions.enabledScopes", 0); +user_pref("extensions.autoDisableScopes", 15); +user_pref("network.stricttransportsecurity.preloadlist", false); +user_pref("security.enterprise_roots.enabled", true); +`, proxyHost, proxyPort, proxyHost, proxyPort) + return os.WriteFile(filepath.Join(profilePath, "user.js"), []byte(userJS), 0600) +} + +// LaunchBrowser launches the given browser pointed at the proxy with an isolated profile. +func LaunchBrowser(fyneApp fyne.App, browser DetectedBrowser, proxyHost string, proxyPort int) error { + profile := profileDir(fyneApp, browser.Path) + proxyAddr := fmt.Sprintf("%s:%d", proxyHost, proxyPort) + + var args []string + + if isFirefox(browser.Path) { + if err := setupFirefoxProfile(profile, proxyHost, proxyPort); err != nil { + return err + } + args = []string{ + "-profile", profile, + "-no-remote", + } + } else { + // Chromium-based — profile dir is created automatically by the browser. + args = []string{ + "--proxy-server=http://" + proxyAddr, + "--user-data-dir=" + profile, + "--no-first-run", + "--no-default-browser-check", + "--disable-extensions", + "--disable-sync", + "--disable-background-networking", + "--disable-default-apps", + "--disable-component-update", + "--disable-breakpad", + "--metrics-recording-only", + "--disable-features=OptimizationHints,MediaRouter", + } + } + + cmd := exec.Command(browser.Path, args...) + cmd.Stderr = nil + cmd.Stdout = nil + if err := cmd.Start(); err != nil { + return fmt.Errorf("launch browser: %w", err) + } + + logger.Info("browser: launched %s (pid %d) via %s", browser.Name, cmd.Process.Pid, proxyAddr) + + // Detach — we don't want to wait for the browser to exit. + go func() { + err := cmd.Wait() + if err != nil { + logger.Error("browser exited early: %v", err) + } + }() + return nil +} + +// ClearBrowserProfile deletes the stored profile for the given browser executable. +func ClearBrowserProfile(fyneApp fyne.App, browserPath string) error { + profile := profileDir(fyneApp, browserPath) + if err := os.RemoveAll(profile); err != nil { + return fmt.Errorf("clear browser profile: %w", err) + } + return nil +} + +// launchDefaultBrowser reads the saved default browser from prefs and launches it. +// If no default is set it shows a picker dialog first. +func launchDefaultBrowser(fyneApp fyne.App, win fyne.Window) { + prefs := fyneApp.Preferences() + proxyHost := prefs.StringWithFallback(prefKeyProxyHost, defaultProxyHost) + proxyPort := prefs.IntWithFallback(prefKeyProxyPort, defaultProxyPort) + + savedPath := prefs.String(prefKeyDefaultBrowser) + if savedPath != "" { + // Verify it still exists. + if _, err := os.Stat(savedPath); err == nil { + browser := DetectedBrowser{Name: filepath.Base(savedPath), Path: savedPath} + if err := LaunchBrowser(fyneApp, browser, proxyHost, proxyPort); err != nil { + dialog.ShowError(err, win) + } + return + } + // Saved browser no longer exists — fall through to picker. + prefs.SetString(prefKeyDefaultBrowser, "") + } + + browsers := DetectBrowsers() + if len(browsers) == 0 { + dialog.ShowInformation("No Browser Found", + "No supported browser was detected on your system.\nInstall Chrome, Brave, Firefox, or another Chromium-based browser.", + win) + return + } + + // If only one browser found, use it directly and save as default. + if len(browsers) == 1 { + prefs.SetString(prefKeyDefaultBrowser, browsers[0].Path) + if err := LaunchBrowser(fyneApp, browsers[0], proxyHost, proxyPort); err != nil { + dialog.ShowError(err, win) + } + return + } + + // Multiple browsers — show picker. + showBrowserPickerDialog(fyneApp, win, browsers, proxyHost, proxyPort) +} + +// showBrowserPickerDialog shows a one-time dialog to pick a default browser. +func showBrowserPickerDialog(fyneApp fyne.App, win fyne.Window, browsers []DetectedBrowser, proxyHost string, proxyPort int) { + names := make([]string, len(browsers)) + for i, b := range browsers { + names[i] = b.Name + } + + selected := browsers[0] + picker := newBrowserPickerContent(names, func(name string) { + for _, b := range browsers { + if b.Name == name { + selected = b + break + } + } + }) + + d := dialog.NewCustomConfirm("Select Default Browser", "Launch", "Cancel", picker, func(confirmed bool) { + if !confirmed { + return + } + fyneApp.Preferences().SetString(prefKeyDefaultBrowser, selected.Path) + if err := LaunchBrowser(fyneApp, selected, proxyHost, proxyPort); err != nil { + dialog.ShowError(err, win) + } + }, win) + closeOnEscape(win, d.Dismiss) + d.Show() +} diff --git a/internal/ui/decoder.go b/internal/ui/decoder.go new file mode 100644 index 0000000..c5c0653 --- /dev/null +++ b/internal/ui/decoder.go @@ -0,0 +1,254 @@ +package ui + +import ( + "crypto/md5" + "crypto/sha1" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "html" + "net/url" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/widget" +) + +type decoderOperation struct { + name string + fwdLabel string // label for the forward button + bwdLabel string // label for the backward button; empty = no backward + forward func(string) string + backward func(string) (string, error) // nil = one-way +} + +var decoderOps = []decoderOperation{ + { + name: "URL Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: url.QueryEscape, + backward: func(s string) (string, error) { return url.QueryUnescape(s) }, + }, + { + name: "URL Path Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: url.PathEscape, + backward: func(s string) (string, error) { return url.PathUnescape(s) }, + }, + { + name: "Base64 Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: func(s string) string { return base64.StdEncoding.EncodeToString([]byte(s)) }, + backward: func(s string) (string, error) { + s = strings.TrimSpace(s) + b, err := base64.StdEncoding.DecodeString(s) + if err != nil { + b, err = base64.URLEncoding.DecodeString(s) + } + return string(b), err + }, + }, + { + name: "Base64 URL-safe Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: func(s string) string { return base64.URLEncoding.EncodeToString([]byte(s)) }, + backward: func(s string) (string, error) { + b, err := base64.URLEncoding.DecodeString(strings.TrimSpace(s)) + return string(b), err + }, + }, + { + name: "Hex Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: func(s string) string { return hex.EncodeToString([]byte(s)) }, + backward: func(s string) (string, error) { + b, err := hex.DecodeString(strings.ReplaceAll(strings.TrimSpace(s), " ", "")) + return string(b), err + }, + }, + { + name: "HTML Entity Encode / Decode", + fwdLabel: "Encode →", + bwdLabel: "← Decode", + forward: html.EscapeString, + backward: func(s string) (string, error) { return html.UnescapeString(s), nil }, + }, + { + name: "JSON Pretty / Minify", + fwdLabel: "Pretty →", + bwdLabel: "← Minify", + forward: func(s string) string { + var v any + if err := json.Unmarshal([]byte(s), &v); err != nil { + return "Error: " + err.Error() + } + var buf strings.Builder + enc := json.NewEncoder(&buf) + enc.SetIndent("", " ") + if err := enc.Encode(v); err != nil { + return "Error: " + err.Error() + } + return strings.TrimRight(buf.String(), "\n") + }, + backward: func(s string) (string, error) { + var v any + if err := json.Unmarshal([]byte(s), &v); err != nil { + return "", err + } + b, err := json.Marshal(v) + return string(b), err + }, + }, + { + name: "JWT Decode (no verify)", + fwdLabel: "Decode →", + forward: func(s string) string { + parts := strings.Split(strings.TrimSpace(s), ".") + if len(parts) < 2 { + return "Error: not a JWT (expected header.payload.signature)" + } + var result strings.Builder + for i, part := range parts[:min(2, len(parts))] { + label := []string{"Header", "Payload"}[i] + b, err := base64.RawURLEncoding.DecodeString(part) + if err != nil { + result.WriteString(fmt.Sprintf("=== %s ===\nError: %v\n\n", label, err)) + continue + } + result.WriteString(fmt.Sprintf("=== %s ===\n%s\n\n", label, string(PrettyJSON(b)))) + } + if len(parts) >= 3 { + result.WriteString(fmt.Sprintf("=== Signature ===\n%s", parts[2])) + } + return result.String() + }, + }, + { + name: "MD5 Hash", + fwdLabel: "Hash →", + forward: func(s string) string { + h := md5.Sum([]byte(s)) + return hex.EncodeToString(h[:]) + }, + }, + { + name: "SHA1 Hash", + fwdLabel: "Hash →", + forward: func(s string) string { + h := sha1.Sum([]byte(s)) + return hex.EncodeToString(h[:]) + }, + }, + { + name: "SHA256 Hash", + fwdLabel: "Hash →", + forward: func(s string) string { + h := sha256.Sum256([]byte(s)) + return hex.EncodeToString(h[:]) + }, + }, +} + +type decoderTab struct{} + +func newDecoderTab() *decoderTab { return &decoderTab{} } + +func (d *decoderTab) build() fyne.CanvasObject { + inputEntry := widget.NewMultiLineEntry() + inputEntry.SetPlaceHolder("Input...") + inputEntry.TextStyle = fyne.TextStyle{Monospace: true} + inputEntry.Wrapping = fyne.TextWrapOff + + outputEntry := widget.NewMultiLineEntry() + outputEntry.SetPlaceHolder("Output...") + outputEntry.TextStyle = fyne.TextStyle{Monospace: true} + outputEntry.Wrapping = fyne.TextWrapOff + + opNames := make([]string, len(decoderOps)) + for i, op := range decoderOps { + opNames[i] = op.name + } + + selectedIdx := 0 + + fwdBtn := widget.NewButton(decoderOps[0].fwdLabel, nil) + fwdBtn.Importance = widget.HighImportance + + bwdBtn := widget.NewButton("← Decode", nil) + + fwdBtn.OnTapped = func() { + op := decoderOps[selectedIdx] + outputEntry.SetText(op.forward(inputEntry.Text)) + } + bwdBtn.OnTapped = func() { + op := decoderOps[selectedIdx] + if op.backward == nil { + return + } + result, err := op.backward(inputEntry.Text) + if err != nil { + outputEntry.SetText("Error: " + err.Error()) + return + } + outputEntry.SetText(result) + } + + opSelect := widget.NewSelect(opNames, func(s string) { + for i, op := range decoderOps { + if op.name == s { + selectedIdx = i + fwdBtn.SetText(op.fwdLabel) + if op.backward != nil { + bwdBtn.SetText(op.bwdLabel) + bwdBtn.Show() + } else { + bwdBtn.Hide() + } + return + } + } + }) + opSelect.SetSelected(opNames[0]) + + // SHA1/MD5/SHA256/JWT have no backward — hide decode button initially if needed. + if decoderOps[0].backward == nil { + bwdBtn.Hide() + } + + swapBtn := widget.NewButton("⇅ Swap", func() { + in := inputEntry.Text + inputEntry.SetText(outputEntry.Text) + outputEntry.SetText(in) + }) + + copyBtn := widget.NewButton("Copy Output", func() { + fyne.CurrentApp().Clipboard().SetContent(outputEntry.Text) + }) + + toolbar := container.NewBorder(nil, nil, + container.NewHBox(fwdBtn, bwdBtn), + container.NewHBox(swapBtn, copyBtn), + opSelect, + ) + + inputPane := container.NewBorder(newBoldLabel("Input"), nil, nil, nil, inputEntry) + outputPane := container.NewBorder(newBoldLabel("Output"), nil, nil, nil, outputEntry) + + split := container.NewVSplit(inputPane, outputPane) + split.SetOffset(0.5) + + return container.NewBorder( + container.New(layout.NewCustomPaddedLayout(8, 8, 0, 0), toolbar), + nil, nil, nil, + split, + ) +} diff --git a/internal/ui/format.go b/internal/ui/format.go new file mode 100644 index 0000000..8038d42 --- /dev/null +++ b/internal/ui/format.go @@ -0,0 +1,120 @@ +package ui + +import ( + "bytes" + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" + + "github.com/shiv/internal/store" +) + +// PathOf extracts the path+query portion from a full URL string. +func PathOf(rawURL string) string { + for _, prefix := range []string{"https://", "http://"} { + if strings.HasPrefix(rawURL, prefix) { + rest := rawURL[len(prefix):] + if slash := strings.Index(rest, "/"); slash >= 0 { + return rest[slash:] + } + return "/" + } + } + return rawURL +} + +// PrettyJSON re-indents JSON, returning the original bytes on any error. +func PrettyJSON(body []byte) []byte { + if len(body) == 0 { + return body + } + var buf bytes.Buffer + if err := json.Indent(&buf, body, "", " "); err != nil { + return body + } + return buf.Bytes() +} + +// hostHeader returns the value to use for the Host header in a reconstructed +// raw request. Default ports (80 for HTTP, 443 for HTTPS) are omitted per +// RFC 7230 §5.4. Non-standard ports are always included. +func hostHeader(tx store.Transaction) string { + if tx.Port == 0 || (tx.TLS && tx.Port == 443) || (!tx.TLS && tx.Port == 80) { + return tx.Host + } + return tx.Host + ":" + strconv.Itoa(tx.Port) +} + +// FormatStoreRequest serialises a Transaction into a raw HTTP/1.1 request string +// suitable for display in the UI or sending via Repeater/Intruder. +// +// The protocol is always written as HTTP/1.1 regardless of what was captured — +func FormatStoreRequest(tx store.Transaction) string { + proto := tx.Proto + if proto == "" { + proto = "HTTP/1.1" + } + + var b strings.Builder + fmt.Fprintf(&b, "%s %s %s\r\n", tx.Method, PathOf(tx.URL), proto) + fmt.Fprintf(&b, "Host: %s\r\n", hostHeader(tx)) + + keys := sortedKeys(tx.ReqHeaders) + for _, k := range keys { + // Skip Host — we already wrote a canonical one above. + if strings.EqualFold(k, "host") { + continue + } + for _, v := range tx.ReqHeaders[k] { + fmt.Fprintf(&b, "%s: %s\r\n", k, v) + } + } + b.WriteString("\r\n") + if len(tx.ReqBody) > 0 { + b.Write(PrettyJSON(tx.ReqBody)) + } + return b.String() +} + +// FormatStoreResponse serialises a Transaction into a raw HTTP response string +// suitable for display in the UI. +// +// The store already holds decompressed body bytes — we must NOT decompress again. +func FormatStoreResponse(tx store.Transaction) string { + var b strings.Builder + // Always render as HTTP/1.1; H2 pseudo-headers are not meaningful in raw text. + const proto = "HTTP/1.1" + fmt.Fprintf(&b, "%s %d\r\n", proto, tx.StatusCode) + + keys := sortedKeys(tx.RespHeaders) + for _, k := range keys { + for _, v := range tx.RespHeaders[k] { + fmt.Fprintf(&b, "%s: %s\r\n", k, v) + } + } + b.WriteString("\r\n") + + if len(tx.RespBody) > 0 { + // tx.RespBody is already decompressed by the proxy before storage. + // Do NOT call Decompress here — the body is plain bytes. + body := tx.RespBody + ct := tx.RespHeaders.Get("Content-Type") + if strings.Contains(ct, "application/json") { + b.Write(PrettyJSON(body)) + } else { + b.Write(body) + } + } + return b.String() +} + +func sortedKeys(h map[string][]string) []string { + keys := make([]string, 0, len(h)) + for k := range h { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} diff --git a/internal/ui/har.go b/internal/ui/har.go new file mode 100644 index 0000000..ba3e085 --- /dev/null +++ b/internal/ui/har.go @@ -0,0 +1,88 @@ +package ui + +import ( + "fmt" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/dialog" + + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +// exportHAR exports all transactions matching the current filter to a HAR file. +// It paginates through the DB so all rows are exported, not just the visible page. +func (h *historyTab) exportHAR() { + if len(h.displayed) == 0 { + dialog.ShowInformation("Export HAR", "No transactions to export.", h.win) + return + } + + f := h.currentFilter() + + go func() { + // Paginate through all matching rows from DB. + var all []store.Transaction + var cursor uint64 + for { + page, err := h.projectStore.TransactionsPage(cursor, f) + if err != nil { + fyne.Do(func() { + dialog.ShowError(fmt.Errorf("HAR export failed: %w", err), h.win) + }) + return + } + if len(page) == 0 { + break + } + all = append(all, page...) + if len(page) < store.PageSize { + break + } + cursor = page[len(page)-1].ID + } + + if len(all) == 0 { + fyne.Do(func() { + dialog.ShowInformation("Export HAR", "No transactions to export.", h.win) + }) + return + } + + // Fetch full bodies — page rows omit bodies for performance. + full := make([]store.Transaction, 0, len(all)) + for _, tx := range all { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("har: get transaction %d: %v", tx.ID, err) + full = append(full, tx) + continue + } + full = append(full, *fullTx) + } + + data, err := store.ExportHAR(full) + if err != nil { + fyne.Do(func() { + dialog.ShowError(fmt.Errorf("HAR export failed: %w", err), h.win) + }) + return + } + + fyne.Do(func() { + saveDialog := dialog.NewFileSave(func(wc fyne.URIWriteCloser, err error) { + if err != nil || wc == nil { + return + } + defer wc.Close() + if _, err := wc.Write(data); err != nil { + logger.Error("har: write: %v", err) + dialog.ShowError(err, h.win) + } + }, h.win) + saveDialog.SetFileName(fmt.Sprintf("shiv-%s.har", time.Now().Format("2006-01-02-150405"))) + saveDialog.Show() + }) + }() +} diff --git a/internal/ui/helpers.go b/internal/ui/helpers.go new file mode 100644 index 0000000..6b706d4 --- /dev/null +++ b/internal/ui/helpers.go @@ -0,0 +1,91 @@ +package ui + +import ( + "fmt" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/widget" + "github.com/shiv/internal/logger" +) + +func closeOnEscape(win fyne.Window, closeFn func()) { + previous := win.Canvas().OnTypedKey() + win.Canvas().SetOnTypedKey(func(keyEvent *fyne.KeyEvent) { + if keyEvent.Name == fyne.KeyEscape { + closeFn() + win.Canvas().SetOnTypedKey(previous) + } + }) +} + +// newBoldLabel returns a label with bold text style. +func newBoldLabel(text string) *widget.Label { + label := widget.NewLabel(text) + label.TextStyle = fyne.TextStyle{Bold: true} + return label +} + +// recoverPanic logs a panic without crashing the process. +// Use as: defer recoverPanic("context description") +func recoverPanic(context string) { + if r := recover(); r != nil { + logger.Error("panic in %s: %v", context, r) + } +} + +func hasSuffix(s string, suffixes ...string) bool { + for _, suffix := range suffixes { + if strings.HasSuffix(s, suffix) { + return true + } + } + return false +} + +// saveToFile opens a file-save dialog and writes content to the chosen path. +func saveToFile(content []byte, defaultName string, win fyne.Window) { + d := dialog.NewFileSave(func(wc fyne.URIWriteCloser, err error) { + if err != nil || wc == nil { + return + } + defer wc.Close() + if _, err := wc.Write(content); err != nil { + dialog.ShowError(err, win) + } + }, win) + d.SetFileName(defaultName) + d.Show() +} + +// copyButton returns a button labelled "Copy" that copies text to clipboard. +func copyButton(label string, text func() string) *widget.Button { + return widget.NewButton(label, func() { + fyne.CurrentApp().Clipboard().SetContent(text()) + }) +} + +// paneHeader returns a border container with a bold label on the left and +// optional extra widgets on the right, used consistently across all detail panes. +func paneHeader(title string, right ...fyne.CanvasObject) fyne.CanvasObject { + var rightObj fyne.CanvasObject + if len(right) == 1 { + rightObj = right[0] + } else if len(right) > 1 { + rightObj = container.NewHBox(right...) + } + return container.NewBorder(nil, nil, newBoldLabel(title), rightObj) +} + +func formatSize(bytes int) string { + switch { + case bytes >= 1<<20: + return fmt.Sprintf("%.1fMB", float64(bytes)/float64(1<<20)) + case bytes >= 1<<10: + return fmt.Sprintf("%.1fKB", float64(bytes)/float64(1<<10)) + default: + return fmt.Sprintf("%dB", bytes) + } +} diff --git a/internal/ui/history.go b/internal/ui/history.go new file mode 100644 index 0000000..18ead77 --- /dev/null +++ b/internal/ui/history.go @@ -0,0 +1,1303 @@ +package ui + +import ( + "fmt" + "image/color" + "net" + "path/filepath" + "slices" + "strconv" + "strings" + "sync" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/widget" + + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +type siteMapNode struct { + children map[string]*siteMapNode +} + +func newSiteMapNode() *siteMapNode { + return &siteMapNode{children: make(map[string]*siteMapNode)} +} + +type siteMap struct { + mu sync.RWMutex + hosts map[string]*siteMapNode + scopeFilter func(host string) bool + textFilter string +} + +func newSiteMap() *siteMap { + return &siteMap{hosts: make(map[string]*siteMapNode)} +} + +func (sm *siteMap) add(tx store.Transaction) { + sm.mu.Lock() + defer sm.mu.Unlock() + key := hostPortKey(tx.Host, tx.Port) + if _, ok := sm.hosts[key]; !ok { + sm.hosts[key] = newSiteMapNode() + } + node := sm.hosts[key] + for _, seg := range splitPath(treePathOf(tx.URL)) { + if _, ok := node.children[seg]; !ok { + node.children[seg] = newSiteMapNode() + } + node = node.children[seg] + } +} + +func (sm *siteMap) clear() { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.hosts = make(map[string]*siteMapNode) +} + +func (sm *siteMap) childUIDs(uid widget.TreeNodeID) []widget.TreeNodeID { + sm.mu.RLock() + defer sm.mu.RUnlock() + if uid == "" { + ids := make([]widget.TreeNodeID, 0, len(sm.hosts)) + for key := range sm.hosts { + // key is "host:port"; extract bare host for scope and text filter checks. + h, portStr, err := net.SplitHostPort(key) + if err != nil || h == "" { + h = key + } + if sm.scopeFilter != nil && !sm.scopeFilter(h) { + continue + } + if sm.textFilter != "" { + port, _ := strconv.Atoi(portStr) + display := strings.ToLower(displayHost(h, port)) + if !strings.Contains(display, strings.ToLower(sm.textFilter)) { + continue + } + } + ids = append(ids, "h:"+key) + } + slices.Sort(ids) + return ids + } + host, port, path, ok := parseNodeID(uid) + if !ok { + return nil + } + node, ok := sm.hosts[hostPortKey(host, port)] + if !ok { + return nil + } + for _, seg := range splitPath(path) { + child, ok := node.children[seg] + if !ok { + return nil + } + node = child + } + ids := make([]widget.TreeNodeID, 0, len(node.children)) + for seg := range node.children { + var childPath string + if path == "" { + childPath = "/" + seg + } else { + childPath = path + "/" + seg + } + ids = append(ids, "p:"+hostPortKey(host, port)+"|"+childPath) + } + slices.Sort(ids) + return ids +} + +func (sm *siteMap) isBranch(uid widget.TreeNodeID) bool { + sm.mu.RLock() + defer sm.mu.RUnlock() + if uid == "" { + return true + } + host, port, path, ok := parseNodeID(uid) + if !ok { + return false + } + node, ok := sm.hosts[hostPortKey(host, port)] + if !ok { + return false + } + for _, seg := range splitPath(path) { + child, ok := node.children[seg] + if !ok { + return false + } + node = child + } + return len(node.children) > 0 +} + +func labelFor(uid widget.TreeNodeID) string { + host, port, path, ok := parseNodeID(uid) + if !ok { + return uid + } + if path == "" { + return displayHost(host, port) + } + segs := splitPath(path) + return "/" + segs[len(segs)-1] +} + +// parseNodeID decodes a tree node UID into its components. +// Host nodes: "h:example.com:1337" → host="example.com", port=1337, path="" +// Path nodes: "p:example.com:1337|/api/v1" → host="example.com", port=1337, path="/api/v1" +func parseNodeID(uid widget.TreeNodeID) (host string, port int, path string, ok bool) { + switch { + case strings.HasPrefix(uid, "h:"): + h, p, err := net.SplitHostPort(uid[2:]) + if err != nil { + // Backward-compat: bare hostname with no port (old DB entries). + return uid[2:], 0, "", true + } + port, _ = strconv.Atoi(p) + return h, port, "", true + case strings.HasPrefix(uid, "p:"): + hostPort, p, ok0 := strings.Cut(uid[2:], "|") + if !ok0 { + return "", 0, "", false + } + h, portStr, err := net.SplitHostPort(hostPort) + if err != nil { + return hostPort, 0, p, true + } + port, _ = strconv.Atoi(portStr) + return h, port, p, true + } + return "", 0, "", false +} + +// hostPortKey is the internal site-map map key for a host+port pair. +// It is always "host:port" regardless of whether the port is default. +func hostPortKey(host string, port int) string { + return net.JoinHostPort(host, strconv.Itoa(port)) +} + +// displayHost formats a host+port for display. Default ports (80, 443) are +// omitted since they add no information and clutter the site map. +func displayHost(host string, port int) string { + if port == 0 || port == 80 || port == 443 { + return host + } + return net.JoinHostPort(host, strconv.Itoa(port)) +} + +func splitPath(path string) []string { + var segs []string + for s := range strings.SplitSeq(path, "/") { + if s != "" { + segs = append(segs, s) + } + } + return segs +} + +// treePathOf returns the path portion of a URL with query string stripped. +// Used for site map tree nodes so that /api?page=1 and /api?page=2 map to +// the same /api node instead of creating separate leaves per unique query. +func treePathOf(rawURL string) string { + p := PathOf(rawURL) + if idx := strings.IndexByte(p, '?'); idx >= 0 { + return p[:idx] + } + return p +} + +var staticExtensions = []string{ + ".js", ".css", ".png", ".jpg", ".jpeg", ".gif", ".ico", ".svg", + ".woff", ".woff2", ".ttf", ".eot", ".map", ".webp", ".avif", +} + +func matchesFilter(tx store.Transaction, f store.TransactionFilter) bool { + if f.Host != "" && tx.Host != f.Host { + return false + } + if f.Port != 0 && tx.Port != f.Port { + return false + } + if f.PathPrefix != "" { + p := treePathOf(tx.URL) + if p != f.PathPrefix && !strings.HasPrefix(p, f.PathPrefix+"/") { + return false + } + } + if !f.ShowOutScope && !tx.InScope { + return false + } + if len(f.Methods) > 0 { + found := false + for _, m := range f.Methods { + if strings.EqualFold(tx.Method, m) { + found = true + break + } + } + if !found { + return false + } + } + if f.StatusMin > 0 && tx.StatusCode < f.StatusMin { + return false + } + if f.StatusMax > 0 && tx.StatusCode > f.StatusMax { + return false + } + if f.HideStatic { + p := strings.ToLower(treePathOf(tx.URL)) + for _, ext := range staticExtensions { + if strings.HasSuffix(p, ext) { + return false + } + } + } + return true +} + +type historyTab struct { + projectStore *store.Store + win fyne.Window + repeater *repeaterTab + loot *lootTab + intruder *intruderTab + + mu sync.RWMutex + + // displayed holds the rows currently shown in the table. + // New traffic is prepended; older rows are appended via infinite scroll. + displayed []store.Transaction + + // cursor is the lowest ID in displayed. The next page loads id < cursor. + cursor uint64 + + // loadingMore prevents concurrent next-page loads. + loadingMore bool + + // hasMore is false when the DB has no more rows below the cursor. + hasMore bool + + // queryID increments on every loadFirstPage call. Goroutines check it + // before writing results so stale responses from old queries are dropped. + queryID uint64 + + siteMap *siteMap + tree *widget.Tree + selectedUID widget.TreeNodeID + + table *widgets.DataTable + filterEntry *widget.Entry + treeFilterEntry *widget.Entry + showOutScope *widget.Check + methodSelect *widget.Select + statusSelect *widget.Select + hideStatic *widget.Check + reqLabel *widgets.TextView + respLabel *widgets.TextView + scopeBtn *widget.Button + deleteBtn *widget.Button + reqCopyBtn *widget.Button + respCopyBtn *widget.Button + respViewFullBtn *widget.Button + + selectedTx store.Transaction + hasSelected bool + + annotationsMu sync.RWMutex + annotations map[uint64]store.Annotation +} + +var tableColumns = []widgets.DataTableColumn{ + {Header: "Method", Width: 80}, + {Header: "Host", Width: 180}, + {Header: "Path", Width: 260}, + {Header: "Status", Width: 70}, + {Header: "Proto", Width: 80}, + {Header: "Size", Width: 80}, + {Header: "Duration", Width: 80}, + {Header: "Time", Width: 160}, + {Header: "Comment", Width: 160}, +} + +func newHistoryTab(projectStore *store.Store, win fyne.Window, repeater *repeaterTab, loot *lootTab, intruder *intruderTab) *historyTab { + return &historyTab{ + projectStore: projectStore, + win: win, + repeater: repeater, + loot: loot, + intruder: intruder, + siteMap: newSiteMap(), + hasMore: true, + annotations: make(map[uint64]store.Annotation), + } +} + +func (h *historyTab) currentFilter() store.TransactionFilter { + var f store.TransactionFilter + if h.filterEntry != nil { + f.Search = h.filterEntry.Text + } + if h.showOutScope != nil { + f.ShowOutScope = h.showOutScope.Checked + } + if h.hideStatic != nil { + f.HideStatic = h.hideStatic.Checked + } + if h.methodSelect != nil { + if sel := h.methodSelect.Selected; sel != "" && sel != "All" { + f.Methods = []string{sel} + } + } + if h.statusSelect != nil { + switch h.statusSelect.Selected { + case "2xx": + f.StatusMin, f.StatusMax = 200, 299 + case "3xx": + f.StatusMin, f.StatusMax = 300, 399 + case "4xx": + f.StatusMin, f.StatusMax = 400, 499 + case "5xx": + f.StatusMin, f.StatusMax = 500, 599 + } + } + + if h.selectedUID != "" { + host, port, path, ok := parseNodeID(h.selectedUID) + if ok { + f.Host = host + f.Port = port + f.PathPrefix = path + } + } + return f +} + +// fetchPage fetches a page of transactions from the DB and applies +// path-prefix filtering. Returns filtered results and unfiltered count. +func (h *historyTab) fetchPage(cursor uint64, f store.TransactionFilter) ([]store.Transaction, int, error) { + txs, err := h.projectStore.TransactionsPage(cursor, f) + if err != nil { + return nil, 0, err + } + fullLen := len(txs) + if f.PathPrefix != "" || f.HideStatic { + filtered := make([]store.Transaction, 0, len(txs)) + for _, tx := range txs { + if matchesFilter(tx, f) { + filtered = append(filtered, tx) + } + } + txs = filtered + } + return txs, fullLen, nil +} + +func (h *historyTab) loadFirstPage() { + f := h.currentFilter() + h.mu.Lock() + h.queryID++ + localID := h.queryID + h.mu.Unlock() + + go func() { + txs, fullLen, err := h.fetchPage(0, f) + if err != nil { + logger.Error("history: load first page: %v", err) + return + } + fyne.Do(func() { + h.mu.Lock() + if localID != h.queryID { + h.mu.Unlock() + return + } + h.loadingMore = false + h.displayed = txs + h.hasMore = fullLen == store.PageSize + if len(txs) > 0 { + h.cursor = txs[len(txs)-1].ID + } else { + h.cursor = 0 + h.hasMore = false + } + h.mu.Unlock() + h.table.ScrollToRow(0) + }) + }() +} + +func (h *historyTab) loadNextPage() { + h.mu.Lock() + if h.loadingMore || !h.hasMore || h.cursor == 0 { + h.mu.Unlock() + return + } + cursor := h.cursor + localID := h.queryID + h.loadingMore = true + h.mu.Unlock() + + f := h.currentFilter() + go func() { + txs, fullLen, err := h.fetchPage(cursor, f) + if err != nil { + logger.Error("history: load next page: %v", err) + h.mu.Lock() + h.loadingMore = false + h.mu.Unlock() + return + } + fyne.Do(func() { + h.mu.Lock() + if localID != h.queryID { + h.loadingMore = false + h.mu.Unlock() + return + } + h.displayed = append(h.displayed, txs...) + h.hasMore = fullLen == store.PageSize + if len(txs) > 0 { + h.cursor = txs[len(txs)-1].ID + } + h.loadingMore = false + h.mu.Unlock() + h.table.Refresh() + }) + }() +} + +func (h *historyTab) buildFilterBar() fyne.CanvasObject { + h.filterEntry = widget.NewEntry() + h.filterEntry.SetPlaceHolder("Filter — host, path, method, status...") + h.filterEntry.OnChanged = func(_ string) { h.loadFirstPage() } + + h.methodSelect = widget.NewSelect( + []string{"All", "GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"}, + func(_ string) { h.loadFirstPage() }, + ) + h.methodSelect.SetSelected("All") + + h.statusSelect = widget.NewSelect( + []string{"All", "2xx", "3xx", "4xx", "5xx"}, + func(_ string) { h.loadFirstPage() }, + ) + h.statusSelect.SetSelected("All") + + h.hideStatic = widget.NewCheck("Hide static", func(_ bool) { h.loadFirstPage() }) + + h.showOutScope = widget.NewCheck("Show out-of-scope", func(checked bool) { + if checked { + h.siteMap.scopeFilter = nil + } else { + h.siteMap.scopeFilter = func(host string) bool { + return h.projectStore.InScope(host) + } + } + h.tree.ScrollToTop() + h.tree.Refresh() + h.loadFirstPage() + }) + h.showOutScope.Checked = true + + exportHARBtn := widget.NewButtonWithIcon("Export HAR", AppIcon("document"), func() { + h.exportHAR() + }) + wsBtn := widget.NewButtonWithIcon("WebSockets", AppIcon("web"), func() { + showWebSocketWindow(fyne.CurrentApp(), h.projectStore, h.win, h.repeater) + }) + clearBtn := widget.NewButtonWithIcon("Clear", AppIcon("delete"), func() { + h.clearHistory() + }) + + filterRow2 := container.NewBorder(nil, nil, nil, + container.NewHBox(h.showOutScope, h.hideStatic, exportHARBtn, wsBtn, clearBtn), + container.NewHBox( + widget.NewLabel("Method:"), h.methodSelect, + widget.NewLabel("Status:"), h.statusSelect, + ), + ) + return container.NewVBox( + container.NewBorder(nil, nil, nil, nil, h.filterEntry), + filterRow2, + ) +} + +func (h *historyTab) deleteSelectedHost() { + if h.selectedUID == "" || !strings.HasPrefix(h.selectedUID, "h:") { + return + } + host, port, _, ok := parseNodeID(h.selectedUID) + if !ok { + return + } + if err := h.projectStore.DeleteTransactionsByHostPort(host, port); err != nil { + logger.Error("history: delete %s:%d: %v", host, port, err) + return + } + h.siteMap.mu.Lock() + delete(h.siteMap.hosts, hostPortKey(host, port)) + h.siteMap.mu.Unlock() + h.selectedUID = "" + h.mu.Lock() + h.displayed = nil + h.cursor = 0 + h.hasMore = true + h.mu.Unlock() + h.tree.UnselectAll() + h.tree.Refresh() + h.deleteBtn.Disable() + h.loadFirstPage() +} + +func (h *historyTab) updateTreeNode(uid widget.TreeNodeID, branch bool, obj fyne.CanvasObject) { + box := obj.(*fyne.Container) + icon := box.Objects[0].(*widget.Icon) + label := box.Objects[1].(*widget.Label) + if strings.HasPrefix(uid, "h:") { + icon.SetResource(AppIcon("web")) + } else if branch { + icon.SetResource(AppIcon("folder")) + } else if hasSuffix(labelFor(uid), ".jpg", ".jpeg", ".png", ".svg", ".gif") { + icon.SetResource(AppIcon("media")) + } else { + icon.SetResource(AppIcon("document")) + } + label.SetText(labelFor(uid)) +} + +func (h *historyTab) buildSiteMap() fyne.CanvasObject { + h.scopeBtn = widget.NewButtonWithIcon("Scope", AppIcon("scope"), func() { + showScopeDialog(h.projectStore, h.win, func() { + h.tree.Refresh() + h.loadFirstPage() + }) + }) + + h.deleteBtn = widget.NewButtonWithIcon("", AppIcon("delete"), func() { + h.deleteSelectedHost() + }) + h.deleteBtn.Disable() + + h.tree = widget.NewTree( + h.siteMap.childUIDs, + h.siteMap.isBranch, + func(_ bool) fyne.CanvasObject { + return container.NewHBox(widget.NewIcon(nil), widget.NewLabel("template")) + }, + h.updateTreeNode, + ) + h.tree.OnSelected = func(uid widget.TreeNodeID) { + h.selectedUID = uid + if strings.HasPrefix(uid, "h:") { + h.deleteBtn.Enable() + } else { + h.deleteBtn.Disable() + } + h.loadFirstPage() + } + + h.treeFilterEntry = widget.NewEntry() + h.treeFilterEntry.SetPlaceHolder("Filter hosts...") + h.treeFilterEntry.OnChanged = func(text string) { + h.siteMap.mu.Lock() + h.siteMap.textFilter = text + h.siteMap.mu.Unlock() + h.tree.Refresh() + } + + allBtn := widget.NewButton("All", func() { + h.selectedUID = "" + h.tree.UnselectAll() + h.deleteBtn.Disable() + h.loadFirstPage() + }) + + siteMapHeader := container.NewHBox(newBoldLabel("Site Map"), layout.NewSpacer(), allBtn, h.deleteBtn, h.scopeBtn) + return container.NewBorder( + container.NewVBox(siteMapHeader, h.treeFilterEntry), + nil, nil, nil, + h.tree, + ) +} + +func (h *historyTab) tableCellStyle(row, col int) widget.Importance { + h.mu.RLock() + defer h.mu.RUnlock() + if row >= len(h.displayed) { + return widget.MediumImportance + } + tx := h.displayed[row] + switch col { + case 0: + switch tx.Method { + case "POST", "PUT", "PATCH", "DELETE": + return widget.WarningImportance + } + case 3: + switch { + case tx.StatusCode >= 500: + return widget.DangerImportance + case tx.StatusCode >= 400: + return widget.WarningImportance + case tx.StatusCode >= 300: + return widget.LowImportance + case tx.StatusCode >= 200: + return widget.SuccessImportance + } + } + return widget.MediumImportance +} + +func (h *historyTab) tableOnSelect(row int) { + h.mu.RLock() + if row >= len(h.displayed) { + h.mu.RUnlock() + return + } + tx := h.displayed[row] + h.mu.RUnlock() + h.selectedTx = tx + h.hasSelected = true + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: get transaction: %v", err) + return + } + fyne.Do(func() { + h.selectedTx = *fullTx + h.showDetail(*fullTx) + }) + }() +} + +func (h *historyTab) tableRowBgColor(row int) color.Color { + h.mu.RLock() + if row >= len(h.displayed) { + h.mu.RUnlock() + return nil + } + id := h.displayed[row].ID + h.mu.RUnlock() + h.annotationsMu.RLock() + a, ok := h.annotations[id] + h.annotationsMu.RUnlock() + if !ok || a.Color == "" { + return nil + } + return annotationColor(a.Color) +} + +func (h *historyTab) buildTable() fyne.CanvasObject { + h.table = widgets.NewDataTable() + h.table.SetWindow(h.win) + h.table.Columns = tableColumns + h.table.RowCount = func() int { + h.mu.RLock() + defer h.mu.RUnlock() + return len(h.displayed) + } + h.table.CellValue = func(row, col int) string { + h.mu.RLock() + defer h.mu.RUnlock() + if row >= len(h.displayed) { + return "" + } + return h.cellText(h.displayed[row], col) + } + h.table.CellStyle = h.tableCellStyle + h.table.RowID = func(row int) int64 { + h.mu.RLock() + defer h.mu.RUnlock() + if row >= len(h.displayed) { + return 0 + } + return int64(h.displayed[row].ID) + } + h.table.OnSelect = h.tableOnSelect + h.table.MenuItems = func(row int) []widgets.ContextMenuItem { + h.mu.RLock() + if row >= len(h.displayed) { + h.mu.RUnlock() + return nil + } + tx := h.displayed[row] + h.mu.RUnlock() + return h.contextMenuItems(tx) + } + h.table.OnNearBottom = func() { + h.loadNextPage() + } + h.table.RowBgColor = h.tableRowBgColor + return h.table.Build() +} + +func (h *historyTab) buildDetailPane() fyne.CanvasObject { + h.reqLabel = widgets.NewTextView() + h.reqLabel.SetWindow(h.win) + h.respLabel = widgets.NewTextView() + h.respLabel.SetWindow(h.win) + + h.reqCopyBtn = widget.NewButton("Copy", func() { + fyne.CurrentApp().Clipboard().SetContent(FormatStoreRequest(h.selectedTx)) + }) + h.reqCopyBtn.Disable() + + h.respCopyBtn = widget.NewButton("Copy", func() { + fyne.CurrentApp().Clipboard().SetContent(FormatStoreResponse(h.selectedTx)) + }) + h.respCopyBtn.Disable() + + h.respViewFullBtn = widget.NewButton("View Full", func() { + showFullBodyDialog(h.selectedTx.RespBody, h.win) + }) + h.respViewFullBtn.Disable() + h.respViewFullBtn.Hide() + + inspectBtn := widget.NewButtonWithIcon("Inspector", AppIcon("inspector"), func() { + if !h.hasSelected { + return + } + showInspectorDialog(h.selectedTx, h.win) + }) + + reqPane := container.NewBorder( + paneHeader("Request", h.reqCopyBtn), + nil, nil, nil, h.reqLabel.Build(), + ) + respPane := container.NewBorder( + paneHeader("Response", h.respViewFullBtn, h.respCopyBtn, inspectBtn), + nil, nil, nil, + h.respLabel.Build(), + ) + + detailSplit := container.NewHSplit(reqPane, respPane) + detailSplit.SetOffset(0.5) + return detailSplit +} + +func (h *historyTab) loadSiteMapEntries() { + entries, err := h.projectStore.SiteMapEntries() + if err != nil { + logger.Error("history: build site map: %v", err) + return + } + fyne.Do(func() { + for _, e := range entries { + h.siteMap.add(store.Transaction{Host: e.Host, Port: e.Port, URL: e.URL}) + } + h.tree.Refresh() + }) +} + +func (h *historyTab) loadAnnotations() { + all, err := h.projectStore.AllAnnotations() + if err != nil { + logger.Error("history: load annotations: %v", err) + return + } + fyne.Do(func() { + h.annotationsMu.Lock() + for _, a := range all { + h.annotations[a.HistoryID] = a + } + h.annotationsMu.Unlock() + h.table.Refresh() + }) +} + +func (h *historyTab) clearHistory() { + if err := h.projectStore.ClearHistory(); err != nil { + logger.Error("clear history: %v", err) + return + } + h.mu.Lock() + h.displayed = nil + h.cursor = 0 + h.hasMore = false + h.mu.Unlock() + h.siteMap.clear() + h.selectedUID = "" + h.tree.Refresh() + h.table.Refresh() + h.selectedTx = store.Transaction{} + h.hasSelected = false + h.reqLabel.SetText("") + h.respLabel.SetText("") +} + +func (h *historyTab) build() fyne.CanvasObject { + filterBar := h.buildFilterBar() + siteMapPane := h.buildSiteMap() + tableObj := h.buildTable() + detailSplit := h.buildDetailPane() + + topSplit := container.NewHSplit( + siteMapPane, + container.NewBorder(filterBar, nil, nil, nil, tableObj), + ) + topSplit.SetOffset(0.25) + + mainSplit := container.NewVSplit(topSplit, detailSplit) + mainSplit.SetOffset(0.55) + + go h.loadSiteMapEntries() + go h.loadAnnotations() + h.loadFirstPage() + go h.watchUpdates() + + return mainSplit +} + +func (h *historyTab) watchUpdates() { + defer recoverPanic("watchUpdates") + for tx := range h.projectStore.Updates { + fyne.Do(func() { + h.mu.Lock() + for _, r := range h.displayed { + if r.ID == tx.ID { + h.mu.Unlock() + return + } + } + f := h.currentFilter() + show := matchesFilter(tx, f) + if show { + h.displayed = append([]store.Transaction{tx}, h.displayed...) + h.cursor = h.displayed[len(h.displayed)-1].ID + } + isSelected := h.hasSelected && h.selectedTx.ID == tx.ID + if isSelected { + h.selectedTx = tx + } + h.mu.Unlock() + h.siteMap.add(tx) + h.tree.Refresh() + if show { + h.table.Refresh() + } + if isSelected { + h.showDetail(tx) + } + }) + } +} + +func (h *historyTab) showDetail(tx store.Transaction) { + truncated := len(tx.RespBody) > internalhttp.MaxDisplayBytes + display := tx + display.RespBody = internalhttp.TruncateBody(tx.RespBody) + display.ReqBody = internalhttp.TruncateBody(tx.ReqBody) + h.reqLabel.SetText(FormatStoreRequest(display)) + h.respLabel.SetText(FormatStoreResponse(display)) + + h.reqCopyBtn.Enable() + h.respCopyBtn.Enable() + if truncated { + h.respViewFullBtn.Show() + h.respViewFullBtn.Enable() + } else { + h.respViewFullBtn.Hide() + } +} + +func showFullBodyDialog(body []byte, win fyne.Window) { + tv := widgets.NewTextView() + tv.SetWindow(win) + tv.SetText(string(body)) + d := dialog.NewCustom("Full Response Body", "Close", tv.Build(), win) + d.Resize(fyne.NewSize(800, 600)) + closeOnEscape(win, d.Dismiss) + d.Show() +} + +func (h *historyTab) annotationMenuItems(tx store.Transaction) []widgets.ContextMenuItem { + return []widgets.ContextMenuItem{ + { + Label: "Add Comment...", + Action: func() { h.showCommentDialog(tx.ID) }, + }, + { + Label: "Highlight", + Action: func() { h.showHighlightMenu(tx.ID) }, + }, + { + Label: "Clear Annotation", + Action: func() { + if err := h.projectStore.DeleteAnnotation(tx.ID); err != nil { + logger.Error("history: clear annotation: %v", err) + return + } + h.annotationsMu.Lock() + delete(h.annotations, tx.ID) + h.annotationsMu.Unlock() + h.table.Refresh() + }, + }, + } +} + +func (h *historyTab) contextMenuItems(tx store.Transaction) []widgets.ContextMenuItem { + items := []widgets.ContextMenuItem{ + { + Label: "Send to Repeater", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("send to repeater: %v", err) + return + } + fyne.Do(func() { h.sendToRepeater(*fullTx) }) + }() + }, + }, + {Label: "Send to Intruder", Action: func() { h.sendToIntruder(tx) }}, + { + Label: "Send to Loot", + Action: func() { + id := tx.ID + h.loot.showAddDialog(&id, "", "") + }, + }, + {Label: "Copy URL", Action: func() { fyne.CurrentApp().Clipboard().SetContent(tx.URL) }}, + { + Label: "Copy Request", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: copy request: %v", err) + return + } + fyne.Do(func() { + fyne.CurrentApp().Clipboard().SetContent(FormatStoreRequest(*fullTx)) + }) + }() + }, + }, + { + Label: "Copy Response", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: copy response: %v", err) + return + } + fyne.Do(func() { + fyne.CurrentApp().Clipboard().SetContent(FormatStoreResponse(*fullTx)) + }) + }() + }, + }, + { + Label: "Save Request...", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: save request: %v", err) + return + } + fyne.Do(func() { + saveToFile([]byte(FormatStoreRequest(*fullTx)), "request.txt", h.win) + }) + }() + }, + }, + { + Label: "Save Response...", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: save response: %v", err) + return + } + fyne.Do(func() { + saveResponseToFile(*fullTx, h.win) + }) + }() + }, + }, + { + Label: "Generate CSRF PoC", + Action: func() { + go func() { + fullTx, err := h.projectStore.GetTransaction(tx.ID) + if err != nil { + logger.Error("history: csrf poc: %v", err) + return + } + fyne.Do(func() { + showCSRFPoCDialog(*fullTx, h.win) + }) + }() + }, + }, + } + return append(items, h.annotationMenuItems(tx)...) +} + +func (h *historyTab) sendToRepeater(tx store.Transaction) { + path := PathOf(tx.URL) + if len(path) > 20 { + path = path[:20] + "..." + } + h.repeater.AddTab(fmt.Sprintf("%s %s", tx.Method, path), tx.Host, tx.Port, tx.TLS, FormatStoreRequest(tx)) +} + +func (h *historyTab) sendToIntruder(tx store.Transaction) { + h.intruder.reqEditor.SetText(FormatStoreRequest(tx)) +} + +func protoLabel(tx store.Transaction) string { + if tx.TLS { + if tx.Proto == "HTTP/2" { + return "H2 TLS" + } + return "H1 TLS" + } + if tx.Proto == "HTTP/2" { + return "H2" + } + return "H1.1" +} + +func (h *historyTab) cellText(tx store.Transaction, col int) string { + switch col { + case 0: + return tx.Method + case 1: + return displayHost(tx.Host, tx.Port) + case 2: + return PathOf(tx.URL) + case 3: + return fmt.Sprintf("%d", tx.StatusCode) + case 4: + return protoLabel(tx) + case 5: + return formatSize(tx.RespSize) + case 6: + return fmt.Sprintf("%dms", tx.DurationMs) + case 7: + return tx.Timestamp.Local().Format("2006-01-02 15:04:05") + case 8: + h.annotationsMu.RLock() + a, ok := h.annotations[tx.ID] + h.annotationsMu.RUnlock() + if ok { + return a.Comment + } + return "" + } + return "" +} + +func saveResponseToFile(tx store.Transaction, win fyne.Window) { + ext := ".txt" + ct := tx.RespHeaders.Get("Content-Type") + switch { + case strings.Contains(ct, "application/json"): + ext = ".json" + case strings.Contains(ct, "text/html"): + ext = ".html" + case strings.Contains(ct, "image/png"): + ext = ".png" + case strings.Contains(ct, "image/jpeg"): + ext = ".jpg" + case strings.Contains(ct, "application/pdf"): + ext = ".pdf" + } + defaultName := "response" + ext + + // Try to derive filename from URL path. + if base := filepath.Base(strings.TrimRight(treePathOf(tx.URL), "/")); base != "" && base != "." && base != "/" { + if filepath.Ext(base) != "" { + defaultName = base + } + } + + // Write raw bytes (not formatted) so binary content is preserved. + saveToFile(tx.RespBody, defaultName, win) +} + +func showCSRFPoCDialog(tx store.Transaction, win fyne.Window) { + poc := generateCSRFPoC(tx) + + tv := widgets.NewTextView() + tv.SetWindow(win) + tv.SetText(poc) + + copyBtn := widget.NewButton("Copy HTML", func() { + fyne.CurrentApp().Clipboard().SetContent(poc) + }) + saveBtn := widget.NewButton("Save as .html...", func() { + saveToFile([]byte(poc), "csrf-poc.html", win) + }) + + content := container.NewBorder( + container.NewHBox(copyBtn, saveBtn), nil, nil, nil, + tv.Build(), + ) + + d := dialog.NewCustom("CSRF PoC", "Close", content, win) + d.Resize(fyne.NewSize(700, 500)) + closeOnEscape(win, d.Dismiss) + d.Show() +} + +func generateCSRFPoC(tx store.Transaction) string { + action := tx.URL + method := strings.ToUpper(tx.Method) + if method == "" { + method = "POST" + } + + var fields strings.Builder + ct := tx.ReqHeaders.Get("Content-Type") + if strings.Contains(ct, "application/x-www-form-urlencoded") && len(tx.ReqBody) > 0 { + // Parse form body and emit hidden inputs. + pairs := strings.Split(string(tx.ReqBody), "&") + for _, pair := range pairs { + kv := strings.SplitN(pair, "=", 2) + name := htmlAttrEscape(kv[0]) + val := "" + if len(kv) == 2 { + val = htmlAttrEscape(kv[1]) + } + fields.WriteString(fmt.Sprintf(" \n", name, val)) + } + } + + var b strings.Builder + b.WriteString("\n\n\n") + b.WriteString(fmt.Sprintf("
\n", htmlAttrEscape(action), method)) + b.WriteString(fields.String()) + b.WriteString(" \n") + b.WriteString("
\n") + b.WriteString("\n") + b.WriteString("\n\n") + return b.String() +} + +func htmlAttrEscape(s string) string { + s = strings.ReplaceAll(s, "&", "&") + s = strings.ReplaceAll(s, "\"", """) + s = strings.ReplaceAll(s, "<", "<") + s = strings.ReplaceAll(s, ">", ">") + return s +} + +func (h *historyTab) showCommentDialog(historyID uint64) { + h.annotationsMu.RLock() + existing := h.annotations[historyID].Comment + h.annotationsMu.RUnlock() + + entry := widget.NewMultiLineEntry() + entry.SetText(existing) + entry.SetPlaceHolder("Enter comment...") + entry.Resize(fyne.NewSize(400, 120)) + + d := dialog.NewCustomConfirm("Add Comment", "Save", "Cancel", + container.NewStack(entry), + func(ok bool) { + if !ok { + return + } + h.annotationsMu.RLock() + a := h.annotations[historyID] + h.annotationsMu.RUnlock() + a.HistoryID = historyID + a.Comment = entry.Text + if err := h.projectStore.SetAnnotation(a); err != nil { + logger.Error("history: set annotation comment: %v", err) + return + } + h.annotationsMu.Lock() + h.annotations[historyID] = a + h.annotationsMu.Unlock() + h.table.Refresh() + }, h.win) + d.Resize(fyne.NewSize(440, 220)) + closeOnEscape(h.win, d.Dismiss) + d.Show() +} + +var highlightColors = []struct { + name string + color string +}{ + {"Red", "red"}, + {"Orange", "orange"}, + {"Yellow", "yellow"}, + {"Green", "green"}, + {"Blue", "blue"}, + {"Purple", "purple"}, +} + +func (h *historyTab) showHighlightMenu(historyID uint64) { + items := make([]*widget.Button, 0, len(highlightColors)+1) + var d *dialog.CustomDialog + for _, hc := range highlightColors { + hc := hc + btn := widget.NewButton(hc.name, func() { + h.annotationsMu.RLock() + a := h.annotations[historyID] + h.annotationsMu.RUnlock() + a.HistoryID = historyID + a.Color = hc.color + if err := h.projectStore.SetAnnotation(a); err != nil { + logger.Error("history: set annotation color: %v", err) + return + } + h.annotationsMu.Lock() + h.annotations[historyID] = a + h.annotationsMu.Unlock() + h.table.Refresh() + d.Hide() + }) + items = append(items, btn) + } + var objs []fyne.CanvasObject + for _, btn := range items { + objs = append(objs, btn) + } + content := container.NewVBox(objs...) + d = dialog.NewCustom("Highlight Color", "Cancel", content, h.win) + closeOnEscape(h.win, d.Dismiss) + d.Show() +} + +func annotationColor(name string) color.Color { + switch name { + case "red": + return color.NRGBA{R: 220, G: 80, B: 80, A: 100} + case "orange": + return color.NRGBA{R: 230, G: 140, B: 50, A: 100} + case "yellow": + return color.NRGBA{R: 220, G: 200, B: 50, A: 100} + case "green": + return color.NRGBA{R: 60, G: 180, B: 80, A: 100} + case "blue": + return color.NRGBA{R: 60, G: 130, B: 220, A: 100} + case "purple": + return color.NRGBA{R: 160, G: 80, B: 200, A: 100} + } + return nil +} + diff --git a/internal/ui/history_test.go b/internal/ui/history_test.go new file mode 100644 index 0000000..b251460 --- /dev/null +++ b/internal/ui/history_test.go @@ -0,0 +1,124 @@ +package ui + +import ( + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" +) + +func TestTreePathOf(t *testing.T) { + tests := []struct { + name string + url string + expected string + }{ + {"absolute https no query", "https://example.com/api/v1", "/api/v1"}, + {"absolute https with query", "https://example.com/api/v1?page=1&limit=10", "/api/v1"}, + {"absolute http with query", "http://example.com/search?q=foo&lang=en", "/search"}, + {"root with query", "https://example.com/?utm_source=x", "/"}, + {"no path", "https://example.com", "/"}, + {"path only no query", "/api/users", "/api/users"}, + {"path only with query", "/api/users?id=123", "/api/users"}, + {"root path with query", "/?q=foo", "/"}, + {"empty string", "", ""}, + {"query only after root", "https://example.com?x=1", "/"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.expected, treePathOf(tt.url)) + }) + } +} + +func TestSiteMap_Add_DeduplicatesQueryStrings(t *testing.T) { + sm := newSiteMap() + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api?page=1"}) + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api?page=2"}) + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api?page=3"}) + + hostUIDs := sm.childUIDs("") + assert.Len(t, hostUIDs, 1, "should have exactly one host") + + pathUIDs := sm.childUIDs(hostUIDs[0]) + assert.Len(t, pathUIDs, 1, "query strings must be deduplicated to a single /api node") +} + +func TestSiteMap_Add_SeparatesByPort(t *testing.T) { + sm := newSiteMap() + sm.add(store.Transaction{Host: "example.com", Port: 80, URL: "http://example.com/"}) + sm.add(store.Transaction{Host: "example.com", Port: 8080, URL: "http://example.com:8080/"}) + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/"}) + + assert.Len(t, sm.childUIDs(""), 3, "each port must be a separate tree entry") +} + +func TestSiteMap_TextFilter(t *testing.T) { + sm := newSiteMap() + sm.add(store.Transaction{Host: "example.com", Port: 80, URL: "http://example.com/"}) + sm.add(store.Transaction{Host: "other.com", Port: 80, URL: "http://other.com/"}) + sm.add(store.Transaction{Host: "test.example.com", Port: 80, URL: "http://test.example.com/"}) + + assert.Len(t, sm.childUIDs(""), 3, "no filter: all hosts visible") + + sm.mu.Lock() + sm.textFilter = "example" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 2, "filter 'example': matches example.com and test.example.com") + + sm.mu.Lock() + sm.textFilter = "other" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 1, "filter 'other': matches only other.com") + + sm.mu.Lock() + sm.textFilter = "EXAMPLE" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 2, "filter must be case-insensitive") + + sm.mu.Lock() + sm.textFilter = "xyz" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 0, "filter 'xyz': no matches") + + sm.mu.Lock() + sm.textFilter = "" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 3, "cleared filter: all hosts visible again") +} + +func TestSiteMap_TextFilter_NonStandardPort(t *testing.T) { + sm := newSiteMap() + sm.add(store.Transaction{Host: "portquiz.net", Port: 8080, URL: "http://portquiz.net:8080/"}) + sm.add(store.Transaction{Host: "portquiz.net", Port: 80, URL: "http://portquiz.net/"}) + + // Port 8080 displays as "portquiz.net:8080" — filter on port should work + sm.mu.Lock() + sm.textFilter = "8080" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 1, "filter on port number must match display host") + + // Filter on hostname alone must match both ports + sm.mu.Lock() + sm.textFilter = "portquiz" + sm.mu.Unlock() + assert.Len(t, sm.childUIDs(""), 2, "filter on hostname must match all ports") +} + +func TestSiteMap_Add_DeepPaths(t *testing.T) { + sm := newSiteMap() + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api/v1/users"}) + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api/v1/posts"}) + sm.add(store.Transaction{Host: "example.com", Port: 443, URL: "https://example.com/api/v2/users"}) + + hostUIDs := sm.childUIDs("") + assert.Len(t, hostUIDs, 1) + + // /api + apiUIDs := sm.childUIDs(hostUIDs[0]) + assert.Len(t, apiUIDs, 1, "should have one /api child") + + // /api/v1 and /api/v2 + versionUIDs := sm.childUIDs(apiUIDs[0]) + assert.Len(t, versionUIDs, 2, "should have v1 and v2 children") +} diff --git a/internal/ui/inspector.go b/internal/ui/inspector.go new file mode 100644 index 0000000..600cb17 --- /dev/null +++ b/internal/ui/inspector.go @@ -0,0 +1,125 @@ +package ui + +import ( + "net/http" + "sort" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +func showInspectorDialog(tx store.Transaction, win fyne.Window) { + tabs := container.NewAppTabs() + + if len(tx.ReqHeaders) > 0 { + tabs.Append(container.NewTabItem("Req Headers", buildHeadersList(tx.ReqHeaders))) + } + if len(tx.ReqBody) > 0 { + body := widgets.NewTextView() + body.SetWindow(win) + body.SetText(string(PrettyJSON(tx.ReqBody))) + tabs.Append(container.NewTabItem("Req Body", body.Build())) + } + + tabs.Append(container.NewTabItem("Resp Headers", buildHeadersList(tx.RespHeaders))) + + cookies := parseCookies(tx.RespHeaders) + tabs.Append(container.NewTabItem("Cookies", buildKVTable(cookies))) + + if len(tx.RespBody) > 0 { + body := widgets.NewTextView() + body.SetWindow(win) + ct := tx.RespHeaders.Get("Content-Type") + if strings.Contains(ct, "application/json") { + body.SetText(string(PrettyJSON(tx.RespBody))) + tabs.Append(container.NewTabItem("JSON", body.Build())) + } else { + body.SetText(string(tx.RespBody)) + tabs.Append(container.NewTabItem("Resp Body", body.Build())) + } + } + + d := dialog.NewCustom("Inspector", "Close", tabs, win) + d.Resize(fyne.NewSize(700, 500)) + closeOnEscape(win, d.Dismiss) + d.Show() +} + +func parseCookies(headers http.Header) [][2]string { + var result [][2]string + for _, line := range headers["Set-Cookie"] { + parts := strings.SplitN(line, ";", 2) + cookieParts := strings.SplitN(strings.TrimSpace(parts[0]), "=", 2) + if len(cookieParts) == 2 { + result = append(result, [2]string{cookieParts[0], cookieParts[1]}) + } else if len(cookieParts) == 1 { + result = append(result, [2]string{cookieParts[0], ""}) + } + if len(parts) > 1 { + for attr := range strings.SplitSeq(parts[1], ";") { + attr = strings.TrimSpace(attr) + if attr == "" { + continue + } + attrParts := strings.SplitN(attr, "=", 2) + if len(attrParts) == 2 { + result = append(result, [2]string{" " + attrParts[0], attrParts[1]}) + } else { + result = append(result, [2]string{" " + attrParts[0], ""}) + } + } + } + } + return result +} + +func buildHeadersList(headers http.Header) fyne.CanvasObject { + keys := make([]string, 0, len(headers)) + for k := range headers { + keys = append(keys, k) + } + sort.Strings(keys) + + var pairs [][2]string + for _, k := range keys { + for _, v := range headers[k] { + pairs = append(pairs, [2]string{k, v}) + } + } + return buildKVTable(pairs) +} + +func buildKVTable(pairs [][2]string) fyne.CanvasObject { + if len(pairs) == 0 { + l := widget.NewLabel("None") + l.Importance = widget.LowImportance + return container.NewCenter(l) + } + t := widget.NewTable( + func() (int, int) { return len(pairs), 2 }, + func() fyne.CanvasObject { + l := widget.NewLabel("") + l.Truncation = fyne.TextTruncateEllipsis + return l + }, + func(id widget.TableCellID, obj fyne.CanvasObject) { + l := obj.(*widget.Label) + if id.Col == 0 { + l.TextStyle = fyne.TextStyle{Bold: true} + l.SetText(pairs[id.Row][0]) + } else { + l.TextStyle = fyne.TextStyle{} + l.SetText(pairs[id.Row][1]) + } + }, + ) + t.SetColumnWidth(0, 200) + t.SetColumnWidth(1, 450) + return t +} diff --git a/internal/ui/intercept.go b/internal/ui/intercept.go new file mode 100644 index 0000000..4ce8e74 --- /dev/null +++ b/internal/ui/intercept.go @@ -0,0 +1,216 @@ +package ui + +import ( + "bufio" + "bytes" + "fmt" + "io" + "net/http" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +type interceptTab struct { + projectStore *store.Store + pending *store.PendingRequest + + toggle *widget.Check + editor *widget.Entry + forward *widget.Button + drop *widget.Button + forwardAll *widget.Button +} + +func newInterceptTab(projectStore *store.Store) *interceptTab { + return &interceptTab{projectStore: projectStore} +} + +func (t *interceptTab) build() fyne.CanvasObject { + t.toggle = widget.NewCheck("Intercept ON", func(on bool) { + t.projectStore.Intercept.SetEnabled(on) + if !on { + // SetEnabled(false) triggers bypass internally — all goroutines + // blocked in Hold are released. We only need to handle whatever + // the UI is currently displaying. + if t.pending != nil { + old := t.pending + t.pending = nil + old.Reply <- store.Decision{Forward: true, Request: old.Request, Body: old.Body} + } + t.clearEditor() + } + }) + + t.editor = widget.NewMultiLineEntry() + t.editor.SetPlaceHolder("No request intercepted yet.") + t.editor.TextStyle = fyne.TextStyle{Monospace: true} + t.editor.Wrapping = fyne.TextWrapOff + t.editor.Disable() + + t.forward = widget.NewButtonWithIcon("Forward", theme.ConfirmIcon(), func() { + t.decide(true) + }) + t.drop = widget.NewButtonWithIcon("Drop", theme.DeleteIcon(), func() { + t.decide(false) + }) + t.forwardAll = widget.NewButtonWithIcon("Forward All", theme.MediaFastForwardIcon(), func() { + // Forward whatever is currently displayed, then signal the gate to + // release all goroutines currently blocked in Hold. Future requests + // will be intercepted normally — ForwardAll is a one-shot release. + if t.pending != nil { + old := t.pending + t.pending = nil + old.Reply <- store.Decision{Forward: true, Request: old.Request, Body: old.Body} + } + t.clearEditor() + t.projectStore.Intercept.ForwardAll() + }) + + t.forward.Importance = widget.HighImportance + t.forward.Disable() + t.drop.Disable() + t.forwardAll.Disable() + + scopeOnly := widget.NewCheck("Scope only", func(on bool) { + if on { + t.projectStore.Intercept.SetScopeFilter(t.projectStore.InScope) + } else { + t.projectStore.Intercept.SetScopeFilter(nil) + } + }) + + copyBtn := widget.NewButton("Copy", func() { + fyne.CurrentApp().Clipboard().SetContent(t.editor.Text) + }) + + buttons := container.NewHBox(t.forward, t.drop, t.forwardAll) + + content := container.NewBorder( + container.NewVBox( + container.NewBorder(nil, nil, + container.NewHBox(t.toggle, scopeOnly), + copyBtn, nil), + widget.NewSeparator(), + ), + buttons, + nil, nil, + t.editor, + ) + + go t.watchQueue() + + return content +} + +// watchQueue reads from the gate queue and presents requests to the UI one at +// a time. Cancelled pending (released by bypass before the UI processed them) +// are skipped via the double IsDone check — once on the watchQueue goroutine +// and once inside fyne.Do on the main goroutine. +func (t *interceptTab) watchQueue() { + for pending := range t.projectStore.Intercept.Queue() { + p := pending + // Fast path: already cancelled before we even schedule fyne.Do. + if p.IsDone() { + continue + } + fyne.Do(func() { + // Second check: may have been cancelled between the fast path above + // and this callback executing on the main goroutine. + if p.IsDone() { + return + } + t.showRequest(p) + }) + } +} + +func formatLiveRequest(req *http.Request, body []byte) string { + var builder bytes.Buffer + path := req.URL.RequestURI() + if path == "" { + path = "/" + } + fmt.Fprintf(&builder, "%s %s HTTP/1.1\r\n", req.Method, path) + fmt.Fprintf(&builder, "Host: %s\r\n", req.Host) + for headerKey, headerValues := range req.Header { + for _, headerValue := range headerValues { + fmt.Fprintf(&builder, "%s: %s\r\n", headerKey, headerValue) + } + } + builder.WriteString("\r\n") + if len(body) > 0 { + builder.Write(body) + } + return builder.String() +} + +func (t *interceptTab) showRequest(p *store.PendingRequest) { + t.pending = p + t.editor.Enable() + t.editor.SetText(formatLiveRequest(p.Request, p.Body)) + t.forward.Enable() + t.drop.Enable() + t.forwardAll.Enable() +} + +func (t *interceptTab) decide(forward bool) { + if t.pending == nil { + return + } + pending := t.pending + rawText := t.editor.Text + t.pending = nil + t.clearEditor() + + if forward { + req, body, err := parseRawRequest(rawText, pending.Request) + if err != nil { + logger.Error("intercept: parse edited request: %v — forwarding original", err) + req = pending.Request + body = pending.Body + } + pending.Reply <- store.Decision{Forward: true, Request: req, Body: body} + } else { + pending.Reply <- store.Decision{Forward: false, Request: pending.Request, Body: pending.Body} + } +} + +func (t *interceptTab) clearEditor() { + t.editor.SetText("") + t.editor.Disable() + t.forward.Disable() + t.drop.Disable() + t.forwardAll.Disable() +} + +func parseRawRequest(raw string, original *http.Request) (*http.Request, []byte, error) { + req, err := http.ReadRequest(bufio.NewReader(bytes.NewBufferString(raw))) + if err != nil { + return nil, nil, err + } + + req.URL.Scheme = original.URL.Scheme + req.URL.Host = original.URL.Host + if req.Host == "" { + req.Host = original.Host + } + + var body []byte + if req.Body != nil { + var err error + body, err = io.ReadAll(req.Body) + req.Body.Close() + if err != nil { + return nil, nil, fmt.Errorf("read request body: %w", err) + } + } + req.Body = http.NoBody + + return req, body, nil +} diff --git a/internal/ui/intruder.go b/internal/ui/intruder.go new file mode 100644 index 0000000..7527b82 --- /dev/null +++ b/internal/ui/intruder.go @@ -0,0 +1,742 @@ +package ui + +import ( + "encoding/base64" + "encoding/hex" + "fmt" + "io" + "net" + "net/url" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/storage" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +var intruderMarkerRegex = regexp.MustCompile(`\$<[^>]+>`) + +func applyPayloadEncoding(payload, encoding string) string { + switch encoding { + case "URL encode": + return url.QueryEscape(payload) + case "Double URL encode": + return url.QueryEscape(url.QueryEscape(payload)) + case "Base64": + return base64.StdEncoding.EncodeToString([]byte(payload)) + case "Hex": + return hex.EncodeToString([]byte(payload)) + default: + return payload + } +} + +type intruderResult struct { + payload string + statusCode int + size int + durationMs int64 + err string + rawResp string + rawReq string +} + +type intruderTab struct { + win fyne.Window + projectStore *store.Store + repeater *repeaterTab + loot *lootTab + + reqEditor *widgets.TextViewEntry + payloadEntry *widget.Entry + filterEntry *widget.Entry + startBtn *widget.Button + stopBtn *widget.Button + progressLabel *widget.Label + responsePane *widgets.TextView + requestPane *widgets.TextView + + sendToRepeaterBtn *widget.Button + sendToLootBtn *widget.Button + viewFullBtn *widget.Button + loadPayloadsBtn *widget.Button + + payloadEncoding string // "none", "url", "double-url", "base64", "hex" + + selectedResult *intruderResult + + config store.IntruderConfig + + mu sync.RWMutex + results []intruderResult + filtered []intruderResult + + table *widgets.DataTable + + running atomic.Bool + stopChan chan struct{} +} + +func newIntruderTab(win fyne.Window, projectStore *store.Store, repeater *repeaterTab, loot *lootTab) *intruderTab { + return &intruderTab{ + win: win, + projectStore: projectStore, + repeater: repeater, + loot: loot, + config: projectStore.LoadIntruderConfig(), + } +} + +var intruderColumns = []widgets.DataTableColumn{ + {Header: "Payload", Width: 200}, + {Header: "Status", Width: 80}, + {Header: "Size", Width: 100}, + {Header: "Duration", Width: 100}, +} + +func (t *intruderTab) showConfigDialog() { + delayEntry := widget.NewEntry() + delayEntry.SetText(strconv.Itoa(t.config.DelayMs)) + delayEntry.SetPlaceHolder("0") + + stopOnStatusEntry := widget.NewEntry() + if t.config.StopOnStatus != 0 { + stopOnStatusEntry.SetText(strconv.Itoa(t.config.StopOnStatus)) + } + stopOnStatusEntry.SetPlaceHolder("0 = disabled") + + maxRedirectsEntry := widget.NewEntry() + maxRedirectsEntry.SetText(strconv.Itoa(t.config.MaxRedirects)) + maxRedirectsEntry.SetPlaceHolder("10") + + followRedirectsSelect := widget.NewSelect([]string{"never", "always", "in-scope"}, nil) + followRedirectsSelect.SetSelected(t.config.FollowRedirects) + + timeoutEntry := widget.NewEntry() + timeoutEntry.SetText(strconv.Itoa(t.config.TimeoutMs)) + timeoutEntry.SetPlaceHolder("30000") + + encodingSelect := widget.NewSelect( + []string{"None", "URL encode", "Double URL encode", "Base64", "Hex"}, + nil, + ) + enc := t.payloadEncoding + if enc == "" { + enc = "None" + } + encodingSelect.SetSelected(enc) + + form := widget.NewForm( + widget.NewFormItem("Delay between requests (ms)", delayEntry), + widget.NewFormItem("Stop on status code", stopOnStatusEntry), + widget.NewFormItem("Follow redirects", followRedirectsSelect), + widget.NewFormItem("Max redirects", maxRedirectsEntry), + widget.NewFormItem("Timeout (ms)", timeoutEntry), + widget.NewFormItem("Payload encoding", encodingSelect), + ) + + sized := container.NewGridWrap(fyne.NewSize(420, 310), form) + configDialog := dialog.NewCustomConfirm("Attack Configuration", "Save", "Cancel", sized, func(confirmed bool) { + if !confirmed { + return + } + if delay, err := strconv.Atoi(strings.TrimSpace(delayEntry.Text)); err == nil && delay >= 0 { + t.config.DelayMs = delay + } + if statusText := strings.TrimSpace(stopOnStatusEntry.Text); statusText == "" || statusText == "0" { + t.config.StopOnStatus = 0 + } else if status, err := strconv.Atoi(statusText); err == nil && status > 0 { + t.config.StopOnStatus = status + } + if maxRedirects, err := strconv.Atoi(strings.TrimSpace(maxRedirectsEntry.Text)); err == nil && maxRedirects >= 0 { + t.config.MaxRedirects = maxRedirects + } + t.config.FollowRedirects = followRedirectsSelect.Selected + if timeout, err := strconv.Atoi(strings.TrimSpace(timeoutEntry.Text)); err == nil && timeout > 0 { + t.config.TimeoutMs = timeout + } + t.payloadEncoding = encodingSelect.Selected + t.config.RawRequest = t.reqEditor.GetText() + t.config.Payloads = t.payloadEntry.Text + t.projectStore.SaveIntruderConfig(t.config) + }, t.win) + + closeOnEscape(t.win, configDialog.Dismiss) + configDialog.Show() +} + +func (t *intruderTab) buildRequestPane() fyne.CanvasObject { + t.reqEditor = widgets.NewTextViewEntry() + t.reqEditor.SetWindow(t.win) + t.reqEditor.SetPlaceHolder("Paste raw HTTP request here, mark injection points with $\n\nExample:\nGET /search?q=$ HTTP/1.1\nHost: example.com") + t.reqEditor.SetText(t.config.RawRequest) + + t.payloadEntry = widget.NewMultiLineEntry() + t.payloadEntry.SetPlaceHolder("One payload per line:\nadmin\npassword\n' OR 1=1--") + t.payloadEntry.SetMinRowsVisible(8) + t.payloadEntry.SetText(t.config.Payloads) + + t.loadPayloadsBtn = widget.NewButtonWithIcon("Load from file", theme.FolderOpenIcon(), func() { + fileDialog := dialog.NewFileOpen(func(readCloser fyne.URIReadCloser, err error) { + if err != nil || readCloser == nil { + return + } + defer readCloser.Close() + data, err := io.ReadAll(readCloser) + if err != nil { + return + } + t.payloadEntry.SetText(string(data)) + }, t.win) + fileDialog.SetFilter(storage.NewExtensionFileFilter([]string{".txt"})) + fileDialog.Show() + }) + + payloadPane := container.NewBorder( + paneHeader("Payloads"), + nil, nil, nil, + t.payloadEntry, + ) + + configSplit := container.NewHSplit( + container.NewBorder(paneHeader("Request"), nil, nil, nil, t.reqEditor.Build()), + payloadPane, + ) + configSplit.SetOffset(0.6) + return configSplit +} + +func (t *intruderTab) tableCellValue(row, col int) string { + t.mu.RLock() + defer t.mu.RUnlock() + if row >= len(t.filtered) { + return "" + } + result := t.filtered[row] + switch col { + case 0: + return result.payload + case 1: + if result.err != "" { + return "ERR" + } + return fmt.Sprintf("%d", result.statusCode) + case 2: + if result.err != "" { + return "-" + } + return fmt.Sprintf("%db", result.size) + case 3: + if result.err != "" { + return "-" + } + return fmt.Sprintf("%dms", result.durationMs) + } + return "" +} + +func (t *intruderTab) tableCellStyle(row, col int) widget.Importance { + if col != 1 { + return widget.MediumImportance + } + t.mu.RLock() + defer t.mu.RUnlock() + if row >= len(t.filtered) { + return widget.MediumImportance + } + result := t.filtered[row] + if result.err != "" { + return widget.DangerImportance + } + switch { + case result.statusCode >= 500: + return widget.DangerImportance + case result.statusCode >= 400: + return widget.WarningImportance + case result.statusCode >= 200 && result.statusCode < 300: + return widget.SuccessImportance + } + return widget.MediumImportance +} + +func (t *intruderTab) tableOnSelect(row int) { + t.mu.RLock() + if row >= len(t.filtered) { + t.mu.RUnlock() + return + } + result := t.filtered[row] + t.mu.RUnlock() + t.selectedResult = &result + t.responsePane.SetText(result.rawResp) + t.requestPane.SetText(result.rawReq) + t.sendToRepeaterBtn.Enable() + t.sendToLootBtn.Enable() + if len(result.rawResp) > internalhttp.MaxDisplayBytes { + t.viewFullBtn.Show() + } else { + t.viewFullBtn.Hide() + } +} + +func (t *intruderTab) buildResultsTable() fyne.CanvasObject { + t.filterEntry = widget.NewEntry() + t.filterEntry.SetPlaceHolder("Filter results — payload, status...") + t.filterEntry.OnChanged = func(_ string) { t.applyFilter() } + + t.table = widgets.NewDataTable() + t.table.SetWindow(t.win) + t.table.Columns = intruderColumns + t.table.RowCount = func() int { + t.mu.RLock() + defer t.mu.RUnlock() + return len(t.filtered) + } + t.table.CellValue = t.tableCellValue + t.table.CellStyle = t.tableCellStyle + t.table.RowID = func(row int) int64 { return int64(row) } + t.table.OnSelect = t.tableOnSelect + t.table.MenuItems = func(row int) []widgets.ContextMenuItem { + t.mu.RLock() + if row >= len(t.filtered) { + t.mu.RUnlock() + return nil + } + result := t.filtered[row] + t.mu.RUnlock() + return t.contextMenuItems(result) + } + + return container.NewBorder( + container.NewBorder(nil, nil, newBoldLabel("Results"), nil, t.filterEntry), + nil, nil, nil, + t.table.Build(), + ) +} + +func (t *intruderTab) buildResultDetailPane() fyne.CanvasObject { + t.responsePane = widgets.NewTextView() + t.responsePane.SetWindow(t.win) + t.responsePane.SetPlaceHolder("Select a result to view response...") + + t.requestPane = widgets.NewTextView() + t.requestPane.SetWindow(t.win) + t.requestPane.SetPlaceHolder("Select a result to view request...") + + copyReqBtn := widget.NewButton("Copy", func() { + if t.selectedResult != nil { + fyne.CurrentApp().Clipboard().SetContent(t.selectedResult.rawReq) + } + }) + copyRespBtn := widget.NewButton("Copy", func() { + if t.selectedResult != nil { + fyne.CurrentApp().Clipboard().SetContent(t.selectedResult.rawResp) + } + }) + t.viewFullBtn = widget.NewButton("View Full", func() { + if t.selectedResult != nil { + showFullBodyDialog([]byte(t.selectedResult.rawResp), t.win) + } + }) + t.viewFullBtn.Hide() + + detailPane := container.NewHSplit( + container.NewBorder( + paneHeader("Request", copyReqBtn), + nil, nil, nil, t.requestPane.Build(), + ), + container.NewBorder( + paneHeader("Response", t.viewFullBtn, copyRespBtn), + nil, nil, nil, t.responsePane.Build(), + ), + ) + return detailPane +} + +func (t *intruderTab) clearResults() { + t.mu.Lock() + t.results = nil + t.filtered = nil + t.mu.Unlock() + t.table.Refresh() + t.progressLabel.SetText("") + t.responsePane.SetText("") + t.requestPane.SetText("") + t.selectedResult = nil + t.sendToRepeaterBtn.Disable() + t.sendToLootBtn.Disable() +} + +func (t *intruderTab) sendSelectedToRepeater() { + if t.selectedResult == nil { + return + } + host, port, useTLS := internalhttp.ParseHostFromRaw(t.selectedResult.rawReq) + path := PathOf(internalhttp.ExtractURL(t.selectedResult.rawReq)) + if len(path) > 20 { + path = path[:20] + "..." + } + t.repeater.AddTab(fmt.Sprintf("Intruder %s", path), host, port, useTLS, t.selectedResult.rawReq) +} + +func (t *intruderTab) buildToolbar() fyne.CanvasObject { + t.progressLabel = widget.NewLabel("") + t.progressLabel.Importance = widget.LowImportance + + t.startBtn = widget.NewButtonWithIcon("Start Attack", theme.MediaPlayIcon(), func() { t.startAttack() }) + t.startBtn.Importance = widget.HighImportance + t.stopBtn = widget.NewButtonWithIcon("Stop", theme.MediaStopIcon(), func() { t.stopAttack() }) + t.stopBtn.Disable() + + configBtn := widget.NewButtonWithIcon("Config", AppIcon("settings"), func() { t.showConfigDialog() }) + + t.sendToRepeaterBtn = widget.NewButtonWithIcon("Repeater", theme.MailForwardIcon(), func() { t.sendSelectedToRepeater() }) + t.sendToRepeaterBtn.Disable() + t.sendToLootBtn = widget.NewButtonWithIcon("Loot", AppIcon("loot"), func() { + if t.selectedResult != nil { + t.loot.showAddDialog(nil, t.selectedResult.rawReq, t.selectedResult.rawResp) + } + }) + t.sendToLootBtn.Disable() + + clearBtn := widget.NewButtonWithIcon("Clear Results", AppIcon("delete"), func() { t.clearResults() }) + + return container.NewHBox(t.startBtn, t.stopBtn, configBtn, clearBtn, t.progressLabel, layout.NewSpacer(), t.loadPayloadsBtn) +} + +func (t *intruderTab) build() fyne.CanvasObject { + requestPane := t.buildRequestPane() + detailPane := t.buildResultDetailPane() + tablePane := t.buildResultsTable() + toolbar := t.buildToolbar() + + resultsSplit := container.NewHSplit(tablePane, detailPane) + resultsSplit.SetOffset(0.4) + + mainSplit := container.NewVSplit(requestPane, resultsSplit) + mainSplit.SetOffset(0.45) + + return container.NewBorder(toolbar, nil, nil, nil, mainSplit) +} + +func (t *intruderTab) followRedirect(rawResp string, originalHost string) (string, bool) { + firstLine := strings.SplitN(rawResp, "\r\n", 2)[0] + parts := strings.Fields(firstLine) + if len(parts) < 2 { + return "", false + } + + var statusCode int + fmt.Sscanf(parts[1], "%d", &statusCode) + if statusCode < 300 || statusCode >= 400 { + return "", false + } + + location := "" + for _, line := range strings.Split(rawResp, "\r\n") { + if strings.HasPrefix(strings.ToLower(line), "location:") { + location = strings.TrimSpace(line[9:]) + break + } + } + if location == "" { + return "", false + } + + switch t.config.FollowRedirects { + case "never": + return "", false + case "in-scope": + // location may be an absolute URL — extract the hostname before the scope check. + scopeHost := location + if strings.HasPrefix(location, "http://") || strings.HasPrefix(location, "https://") { + rest := location[strings.Index(location, "//")+2:] + if slash := strings.Index(rest, "/"); slash >= 0 { + scopeHost = rest[:slash] + } else { + scopeHost = rest + } + } + if !t.projectStore.InScope(scopeHost) { + return "", false + } + } + + var redirectHost, redirectPath string + var useTLS bool + + if strings.HasPrefix(location, "https://") { + rest := location[8:] + useTLS = true + if slash := strings.Index(rest, "/"); slash >= 0 { + redirectHost = rest[:slash] + redirectPath = rest[slash:] + } else { + redirectHost = rest + redirectPath = "/" + } + } else if strings.HasPrefix(location, "http://") { + rest := location[7:] + useTLS = false + if slash := strings.Index(rest, "/"); slash >= 0 { + redirectHost = rest[:slash] + redirectPath = rest[slash:] + } else { + redirectHost = rest + redirectPath = "/" + } + } else { + redirectHost = originalHost + redirectPath = location + useTLS = strings.HasSuffix(originalHost, ":443") + } + + port := 443 + if !useTLS { + port = 80 + } + if host, portStr, err := net.SplitHostPort(redirectHost); err == nil { + redirectHost = host + fmt.Sscanf(portStr, "%d", &port) + } + + redirectReq := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\n\r\n", redirectPath, redirectHost) + result, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: redirectHost, + Port: port, + TLS: useTLS, + RawReq: redirectReq, + TimeoutMs: t.config.TimeoutMs, + }) + if err != nil { + return "", false + } + return result.Raw, true +} + +func (t *intruderTab) startAttack() { + rawReq := t.reqEditor.GetText() + if rawReq == "" { + dialog.ShowInformation("Error", "Please enter a request.", t.win) + return + } + + markers := intruderMarkerRegex.FindAllString(rawReq, -1) + if len(markers) == 0 { + dialog.ShowInformation("Error", "No injection points found. Mark them with $.", t.win) + return + } + + payloadsText := strings.TrimSpace(t.payloadEntry.Text) + if payloadsText == "" { + dialog.ShowInformation("Error", "Please enter at least one payload.", t.win) + return + } + + payloads := strings.Split(payloadsText, "\n") + var cleanPayloads []string + for _, payload := range payloads { + payload = strings.TrimSpace(payload) + if payload != "" { + cleanPayloads = append(cleanPayloads, payload) + } + } + + if len(cleanPayloads) == 0 { + dialog.ShowInformation("Error", "No valid payloads found.", t.win) + return + } + + t.config.RawRequest = rawReq + t.config.Payloads = t.payloadEntry.Text + t.projectStore.SaveIntruderConfig(t.config) + + t.mu.Lock() + t.results = nil + t.filtered = nil + t.mu.Unlock() + t.table.ClearSelection() + t.table.Refresh() + t.responsePane.SetText("") + t.requestPane.SetText("") + t.selectedResult = nil + t.sendToRepeaterBtn.Disable() + t.sendToLootBtn.Disable() + + t.running.Store(true) + t.stopChan = make(chan struct{}) + t.startBtn.Disable() + t.stopBtn.Enable() + + config := t.config + encoding := t.payloadEncoding + total := len(cleanPayloads) * len(markers) + done := 0 + + go func() { + for _, marker := range markers { + for _, payload := range cleanPayloads { + select { + case <-t.stopChan: + fyne.Do(func() { + t.progressLabel.SetText(fmt.Sprintf("Stopped after %d/%d requests", done, total)) + t.startBtn.Enable() + t.stopBtn.Disable() + }) + return + default: + } + + if config.DelayMs > 0 { + time.Sleep(time.Duration(config.DelayMs) * time.Millisecond) + } + + encodedPayload := applyPayloadEncoding(payload, encoding) + injected := strings.ReplaceAll(rawReq, marker, encodedPayload) + host, port, useTLS := internalhttp.ParseHostFromRaw(injected) + + var result intruderResult + result.payload = payload + result.rawReq = injected + + if host == "" { + result.err = "no Host header found" + } else { + start := time.Now() + resp, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, + Port: port, + TLS: useTLS, + RawReq: injected, + TimeoutMs: config.TimeoutMs, + }) + elapsed := time.Since(start).Milliseconds() + if err != nil { + result.err = err.Error() + } else { + rawResp := resp.Raw + if config.FollowRedirects != "never" && config.MaxRedirects > 0 { + for redirectCount := 0; redirectCount < config.MaxRedirects; redirectCount++ { + redirectResp, followed := t.followRedirect(rawResp, host) + if !followed { + break + } + rawResp = redirectResp + } + } + result.durationMs = elapsed + result.size = len(rawResp) + result.rawResp = rawResp + firstLine := strings.SplitN(rawResp, "\r\n", 2)[0] + parts := strings.Fields(firstLine) + if len(parts) >= 2 { + fmt.Sscanf(parts[1], "%d", &result.statusCode) + } + } + } + + done++ + shouldStop := config.StopOnStatus != 0 && result.statusCode == config.StopOnStatus + resultCopy := result + + fyne.Do(func() { + t.mu.Lock() + t.results = append(t.results, resultCopy) + t.mu.Unlock() + t.applyFilter() + t.progressLabel.SetText(fmt.Sprintf("%d/%d", done, total)) + }) + + if shouldStop { + fyne.Do(func() { + t.progressLabel.SetText(fmt.Sprintf("Stopped — status %d matched after %d requests", config.StopOnStatus, done)) + t.startBtn.Enable() + t.stopBtn.Disable() + }) + return + } + } + } + + fyne.Do(func() { + t.running.Store(false) + t.startBtn.Enable() + t.stopBtn.Disable() + t.progressLabel.SetText(fmt.Sprintf("Done — %d requests", done)) + }) + }() +} + +func (t *intruderTab) stopAttack() { + if t.stopChan != nil { + close(t.stopChan) + } + t.running.Store(false) +} + +func (t *intruderTab) applyFilter() { + query := strings.ToLower(t.filterEntry.Text) + + t.mu.Lock() + var filtered []intruderResult + for _, result := range t.results { + if query != "" { + searchable := strings.ToLower(fmt.Sprintf("%s %d %s", result.payload, result.statusCode, result.err)) + if !strings.Contains(searchable, query) { + continue + } + } + filtered = append(filtered, result) + } + t.filtered = filtered + t.mu.Unlock() + + t.table.Refresh() +} + +func (t *intruderTab) contextMenuItems(result intruderResult) []widgets.ContextMenuItem { + return []widgets.ContextMenuItem{ + { + Label: "Send to Repeater", + Action: func() { + host, port, useTLS := internalhttp.ParseHostFromRaw(result.rawReq) + path := PathOf(internalhttp.ExtractURL(result.rawReq)) + if len(path) > 20 { + path = path[:20] + "..." + } + t.repeater.AddTab(fmt.Sprintf("Intruder %s", path), host, port, useTLS, result.rawReq) + }, + }, + { + Label: "Send to Loot", + Action: func() { + t.loot.showAddDialog(nil, result.rawReq, result.rawResp) + }, + }, + { + Label: "Copy Request", + Action: func() { + fyne.CurrentApp().Clipboard().SetContent(result.rawReq) + }, + }, + { + Label: "Copy Response", + Action: func() { + fyne.CurrentApp().Clipboard().SetContent(result.rawResp) + }, + }, + } +} diff --git a/internal/ui/keybinds.go b/internal/ui/keybinds.go new file mode 100644 index 0000000..b493bf7 --- /dev/null +++ b/internal/ui/keybinds.go @@ -0,0 +1,92 @@ +package ui + +import ( + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/driver/desktop" +) + +const ( + prefKeySendRequest = "shortcut_send_request" + prefKeyCloseTab = "shortcut_close_tab" + prefKeyToRepeater = "shortcut_to_repeater" + + defaultKeySendRequest fyne.KeyName = fyne.KeyS + defaultKeyCloseTab fyne.KeyName = fyne.KeyW + defaultKeyToRepeater fyne.KeyName = fyne.KeyR +) + +type Keybinds struct { + win fyne.Window + tabs *container.AppTabs + history *historyTab + repeater *repeaterTab + active []fyne.Shortcut + + KeySendRequest fyne.KeyName + KeyCloseTab fyne.KeyName + KeyToRepeater fyne.KeyName +} + +func newKeybinds(win fyne.Window, tabs *container.AppTabs, history *historyTab, repeater *repeaterTab, prefs fyne.Preferences) *Keybinds { + keybinds := &Keybinds{ + win: win, + tabs: tabs, + history: history, + repeater: repeater, + KeySendRequest: fyne.KeyName(prefs.StringWithFallback(prefKeySendRequest, string(defaultKeySendRequest))), + KeyCloseTab: fyne.KeyName(prefs.StringWithFallback(prefKeyCloseTab, string(defaultKeyCloseTab))), + KeyToRepeater: fyne.KeyName(prefs.StringWithFallback(prefKeyToRepeater, string(defaultKeyToRepeater))), + } + keybinds.Update() + return keybinds +} + +func (keybinds *Keybinds) Update() { + for _, registeredShortcut := range keybinds.active { + keybinds.win.Canvas().RemoveShortcut(registeredShortcut) + } + keybinds.active = nil + + isActiveTab := func(tabName string) bool { + selectedTab := keybinds.tabs.Selected() + return selectedTab != nil && selectedTab.Text == tabName + } + + sendRequestShortcut := &desktop.CustomShortcut{KeyName: keybinds.KeySendRequest, Modifier: fyne.KeyModifierControl} + keybinds.win.Canvas().AddShortcut(sendRequestShortcut, func(_ fyne.Shortcut) { + if !isActiveTab("Repeater") { + return + } + selectedTab := keybinds.repeater.tabs.Selected() + if selectedTab == nil { + return + } + if sendFn, ok := keybinds.repeater.sendFns[selectedTab]; ok { + sendFn() + } + }) + keybinds.active = append(keybinds.active, sendRequestShortcut) + + closeTabShortcut := &desktop.CustomShortcut{KeyName: keybinds.KeyCloseTab, Modifier: fyne.KeyModifierControl} + keybinds.win.Canvas().AddShortcut(closeTabShortcut, func(_ fyne.Shortcut) { + if !isActiveTab("Repeater") { + return + } + selectedTab := keybinds.repeater.tabs.Selected() + if selectedTab == nil { + return + } + keybinds.repeater.closeTab(selectedTab) + keybinds.repeater.tabs.Remove(selectedTab) + }) + keybinds.active = append(keybinds.active, closeTabShortcut) + + toRepeaterShortcut := &desktop.CustomShortcut{KeyName: keybinds.KeyToRepeater, Modifier: fyne.KeyModifierControl} + keybinds.win.Canvas().AddShortcut(toRepeaterShortcut, func(_ fyne.Shortcut) { + if isActiveTab("History") && keybinds.history.hasSelected { + keybinds.history.sendToRepeater(keybinds.history.selectedTx) + } + }) + keybinds.active = append(keybinds.active, toRepeaterShortcut) +} diff --git a/internal/ui/launch.go b/internal/ui/launch.go new file mode 100644 index 0000000..75dca54 --- /dev/null +++ b/internal/ui/launch.go @@ -0,0 +1,219 @@ +package ui + +import ( + "encoding/json" + "os" + "path/filepath" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/canvas" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/storage" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/assets" +) + +const maxRecentProjects = 10 + +type RecentProject struct { + Path string `json:"path"` + LastOpen time.Time `json:"last_open"` +} + +func ShowLaunchScreen(app fyne.App, onSelect func(projectPath string, w fyne.Window)) { + launchWin := app.NewWindow("Shiv") + applyTheme(app) + launchWin.Resize(fyne.NewSize(560, 420)) + launchWin.CenterOnScreen() + launchWin.SetFixedSize(true) + + var list *widget.List + var onDeleteBtn func(widget.ListItemID) + + recents, _ := loadRecents() + + list = widget.NewList( + func() int { return len(recents) }, + func() fyne.CanvasObject { + return container.NewHBox( + widget.NewIcon(AppIcon("project")), + container.NewHBox( + widget.NewLabel("name"), + widget.NewLabel("date"), + ), + layout.NewSpacer(), + widget.NewButtonWithIcon("", AppIcon("delete"), nil), + ) + }, + func(i widget.ListItemID, obj fyne.CanvasObject) { + hbox := obj.(*fyne.Container) + box := hbox.Objects[1].(*fyne.Container) + box.Objects[0].(*widget.Label).SetText(filepath.Base(recents[i].Path)) + box.Objects[1].(*widget.Label).SetText(recents[i].LastOpen.Format("2006-01-02 15:04")) + box.Objects[1].(*widget.Label).Importance = widget.LowImportance + hbox.Objects[3].(*widget.Button).OnTapped = func() { onDeleteBtn(i) } + }, + ) + + onDeleteBtn = func(i widget.ListItemID) { + deleteRecent(recents[i].Path) + recents, _ = loadRecents() + list.Refresh() + } + + list.OnSelected = func(i widget.ListItemID) { + path := recents[i].Path + if _, err := os.Stat(path); err != nil { + dialog.ShowError(err, launchWin) + list.UnselectAll() + return + } + saveRecent(path) + onSelect(path, launchWin) + } + + newBtn := widget.NewButtonWithIcon("New Project", theme.DocumentCreateIcon(), func() { + saveDialog := dialog.NewFileSave(func(writeCloser fyne.URIWriteCloser, err error) { + if err != nil || writeCloser == nil { + return + } + writeCloser.Close() + path := writeCloser.URI().Path() + saveRecent(path) + onSelect(path, launchWin) + }, launchWin) + saveDialog.SetFileName("untitled.shiv") + saveDialog.SetFilter(storage.NewExtensionFileFilter([]string{".shiv"})) + saveDialog.Show() + }) + + openBtn := widget.NewButtonWithIcon("Open Project", theme.FolderOpenIcon(), func() { + openDialog := dialog.NewFileOpen(func(readCloser fyne.URIReadCloser, err error) { + if err != nil || readCloser == nil { + return + } + readCloser.Close() + path := readCloser.URI().Path() + saveRecent(path) + onSelect(path, launchWin) + }, launchWin) + openDialog.SetFilter(storage.NewExtensionFileFilter([]string{".shiv"})) + openDialog.Show() + }) + newBtn.Importance = widget.HighImportance + + logo := canvas.NewImageFromResource(fyne.NewStaticResource("logo.png", assets.Logo)) + logo.FillMode = canvas.ImageFillContain + logo.SetMinSize(fyne.NewSize(128, 128)) + + subtitle := widget.NewLabel("HTTP/HTTPS Proxy Ninja") + subtitle.Importance = widget.LowImportance + + recentLabel := widget.NewLabel("Recent Projects") + recentLabel.TextStyle = fyne.TextStyle{Bold: true} + + header := container.NewVBox( + container.NewCenter(logo), + container.NewCenter(subtitle), + widget.NewSeparator(), + ) + + buttons := container.NewGridWithColumns(2, newBtn, openBtn) + + var content fyne.CanvasObject + if len(recents) == 0 { + empty := widget.NewLabel("No recent projects") + empty.Importance = widget.LowImportance + content = container.NewBorder( + container.NewVBox(header, recentLabel), + buttons, + nil, nil, + container.NewCenter(empty), + ) + } else { + content = container.NewBorder( + container.NewVBox(header, recentLabel), + buttons, + nil, nil, + list, + ) + } + + launchWin.SetContent(container.NewPadded(content)) + launchWin.Show() +} + +func recentsPath() (string, error) { + configDir, err := os.UserConfigDir() + if err != nil { + return "", err + } + return filepath.Join(configDir, "shiv", "recent_projects.json"), nil +} + +func loadRecents() ([]RecentProject, error) { + path, err := recentsPath() + if err != nil { + return nil, err + } + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var recents []RecentProject + if err := json.Unmarshal(data, &recents); err != nil { + return nil, err + } + return recents, nil +} + +func saveRecent(projectPath string) { + recents, _ := loadRecents() + + filtered := recents[:0] + for _, recent := range recents { + if recent.Path != projectPath { + filtered = append(filtered, recent) + } + } + + updated := append([]RecentProject{{Path: projectPath, LastOpen: time.Now()}}, filtered...) + if len(updated) > maxRecentProjects { + updated = updated[:maxRecentProjects] + } + + path, err := recentsPath() + if err != nil { + return + } + data, err := json.Marshal(updated) + if err != nil { + return + } + os.WriteFile(path, data, 0o644) +} + +func deleteRecent(projectPath string) { + os.Remove(projectPath) + recents, _ := loadRecents() + filtered := recents[:0] + for _, r := range recents { + if r.Path != projectPath { + filtered = append(filtered, r) + } + } + path, err := recentsPath() + if err != nil { + return + } + data, err := json.Marshal(filtered) + if err != nil { + return + } + os.WriteFile(path, data, 0o644) +} diff --git a/internal/ui/loot.go b/internal/ui/loot.go new file mode 100644 index 0000000..df98ee0 --- /dev/null +++ b/internal/ui/loot.go @@ -0,0 +1,394 @@ +package ui + +import ( + "fmt" + "strings" + "text/template" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/assets" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +var severities = []string{"Critical", "High", "Medium", "Low", "Info"} + +var severityImportance = map[string]widget.Importance{ + "Critical": widget.DangerImportance, + "High": widget.DangerImportance, + "Medium": widget.WarningImportance, + "Low": widget.MediumImportance, + "Info": widget.LowImportance, +} + +type lootTab struct { + projectStore *store.Store + win fyne.Window + repeater *repeaterTab + entries []store.LootEntry + + table *widgets.DataTable + notesArea *widgets.TextView + viewReqBtn *widget.Button + deleteBtn *widget.Button + exportBtn *widget.Button + + selectedIdx int +} + +var lootColumns = []widgets.DataTableColumn{ + {Header: "Severity", Width: 300}, + {Header: "Title", Width: 500}, + {Header: "Created", Width: 360}, +} + +func newLootTab(projectStore *store.Store, win fyne.Window, repeater *repeaterTab) *lootTab { + return &lootTab{ + projectStore: projectStore, + win: win, + repeater: repeater, + selectedIdx: -1, + } +} + +func (l *lootTab) build() fyne.CanvasObject { + l.table = widgets.NewDataTable() + l.table.SetWindow(l.win) + l.table.Columns = lootColumns + l.table.RowCount = func() int { return len(l.entries) } + l.table.CellValue = func(row, col int) string { + if row >= len(l.entries) { + return "" + } + entry := l.entries[row] + switch col { + case 0: + return entry.Severity + case 1: + return entry.Title + case 2: + return entry.CreatedAt.Format("2006-01-02 15:04") + } + return "" + } + l.table.CellStyle = func(row, col int) widget.Importance { + if col != 0 || row >= len(l.entries) { + return widget.MediumImportance + } + if imp, ok := severityImportance[l.entries[row].Severity]; ok { + return imp + } + return widget.MediumImportance + } + l.table.RowID = func(row int) int64 { + if row >= len(l.entries) { + return 0 + } + return int64(l.entries[row].ID) + } + l.table.OnSelect = func(row int) { + if row >= len(l.entries) { + return + } + l.selectedIdx = row + l.deleteBtn.Enable() + entry := l.entries[row] + l.notesArea.SetText(fmt.Sprintf("Severity: %s\nCreated: %s\n\n%s", + entry.Severity, entry.CreatedAt.Format("2006-01-02 15:04"), entry.Notes)) + if entry.HistoryID != nil || entry.RawRequest != "" { + l.viewReqBtn.Enable() + } else { + l.viewReqBtn.Disable() + } + } + l.table.MenuItems = func(row int) []widgets.ContextMenuItem { + if row >= len(l.entries) { + return nil + } + entry := l.entries[row] + return []widgets.ContextMenuItem{ + { + Label: "Delete", + Action: func() { + dialog.ShowConfirm("Delete", fmt.Sprintf("Delete '%s'?", entry.Title), func(ok bool) { + if !ok { + return + } + if err := l.projectStore.DeleteLoot(entry.ID); err != nil { + logger.Error("loot: delete: %v", err) + return + } + l.selectedIdx = -1 + l.deleteBtn.Disable() + l.notesArea.SetText("") + }, l.win) + }, + }, + { + Label: "View Request", + Action: func() { + if entry.HistoryID != nil || entry.RawRequest != "" { + l.showLinkedRequest(entry) + } + }, + }, + } + } + + tableObj := l.table.Build() + + l.notesArea = widgets.NewTextView() + l.notesArea.SetWindow(l.win) + l.notesArea.SetPlaceHolder("Select an entry to view details...") + + l.deleteBtn = widget.NewButtonWithIcon("Delete", theme.DeleteIcon(), func() { + if l.selectedIdx < 0 || l.selectedIdx >= len(l.entries) { + return + } + entry := l.entries[l.selectedIdx] + dialog.ShowConfirm("Delete", fmt.Sprintf("Delete '%s'?", entry.Title), func(ok bool) { + if !ok { + return + } + if err := l.projectStore.DeleteLoot(entry.ID); err != nil { + logger.Error("loot: delete: %v", err) + return + } + l.selectedIdx = -1 + l.deleteBtn.Disable() + l.notesArea.SetText("") + }, l.win) + }) + l.deleteBtn.Disable() + + addBtn := widget.NewButtonWithIcon("Add", theme.ContentAddIcon(), func() { + l.showAddDialog(nil, "", "") + }) + addBtn.Importance = widget.HighImportance + + l.exportBtn = widget.NewButtonWithIcon("Export Markdown", theme.DocumentSaveIcon(), func() { + l.exportMarkdown() + }) + + l.viewReqBtn = widget.NewButtonWithIcon("View Request", theme.SearchIcon(), func() { + if l.selectedIdx < 0 || l.selectedIdx >= len(l.entries) { + return + } + l.showLinkedRequest(l.entries[l.selectedIdx]) + }) + l.viewReqBtn.Disable() + + split := container.NewVSplit( + container.NewBorder(container.NewHBox(addBtn, l.deleteBtn, l.viewReqBtn, l.exportBtn), nil, nil, nil, tableObj), + container.NewBorder(paneHeader("Notes"), nil, nil, nil, l.notesArea.Build()), + ) + split.SetOffset(0.6) + + l.reload() + + // Watch for new loot entries added by plugins or other goroutines. + go func() { + for range l.projectStore.LootEntries { + fyne.Do(func() { + l.reload() + }) + } + }() + + return split +} + +func (l *lootTab) reload() { + entries, err := l.projectStore.AllLoot() + if err != nil { + logger.Error("loot: load: %v", err) + return + } + l.entries = entries + l.table.Refresh() +} + +func (l *lootTab) showAddDialog(historyID *uint64, rawRequest string, rawResponse string) { + titleEntry := widget.NewEntry() + titleEntry.SetPlaceHolder("e.g. Admin credentials found") + + severitySelect := widget.NewSelect(severities, nil) + severitySelect.SetSelected("High") + + notesEntry := widget.NewMultiLineEntry() + notesEntry.SetPlaceHolder("Describe the finding, include evidence, impact...") + notesEntry.SetMinRowsVisible(6) + + form := widget.NewForm( + widget.NewFormItem("Title", titleEntry), + widget.NewFormItem("Severity", severitySelect), + widget.NewFormItem("Notes", notesEntry), + ) + + d := dialog.NewCustomConfirm("Add Loot", "Save", "Cancel", + container.NewGridWrap(fyne.NewSize(500, 350), form), + func(confirmed bool) { + if !confirmed { + return + } + title := strings.TrimSpace(titleEntry.Text) + if title == "" { + return + } + if _, err := l.projectStore.AddLoot(store.LootEntry{ + Title: title, + Severity: severitySelect.Selected, + Notes: notesEntry.Text, + HistoryID: historyID, + RawRequest: rawRequest, + RawResponse: rawResponse, + }); err != nil { + logger.Error("loot: add: %v", err) + dialog.ShowError(err, l.win) + return + } + }, l.win) + closeOnEscape(l.win, d.Dismiss) + d.Show() +} + +func (l *lootTab) exportMarkdown() { + entries, err := l.projectStore.AllLoot() + if err != nil { + logger.Error("loot: export: %v", err) + return + } + if len(entries) == 0 { + dialog.ShowInformation("Export", "No loot entries to export.", l.win) + return + } + + type entryData struct { + Severity string + Title string + Date string + Notes string + Request string + Response string + } + type templateData struct { + Generated string + Entries []entryData + } + + data := templateData{Generated: time.Now().Format("2006-01-02 15:04")} + for _, entry := range entries { + ed := entryData{ + Severity: entry.Severity, + Title: entry.Title, + Date: entry.CreatedAt.Format("2006-01-02 15:04"), + Notes: entry.Notes, + } + if entry.HistoryID != nil { + if tx, err := l.projectStore.GetTransaction(*entry.HistoryID); err == nil { + ed.Request = FormatStoreRequest(*tx) + ed.Response = FormatStoreResponse(*tx) + } + } else { + ed.Request = entry.RawRequest + ed.Response = entry.RawResponse + } + data.Entries = append(data.Entries, ed) + } + + tmpl, err := template.New("findings").Parse(assets.FindingsTemplate) + if err != nil { + logger.Error("loot: parse template: %v", err) + dialog.ShowError(err, l.win) + return + } + + var builder strings.Builder + if err := tmpl.Execute(&builder, data); err != nil { + logger.Error("loot: execute template: %v", err) + dialog.ShowError(err, l.win) + return + } + + saveDialog := dialog.NewFileSave(func(wc fyne.URIWriteCloser, err error) { + if err != nil || wc == nil { + return + } + defer wc.Close() + if _, err := wc.Write([]byte(builder.String())); err != nil { + logger.Error("loot: write export: %v", err) + dialog.ShowError(err, l.win) + } + }, l.win) + saveDialog.SetFileName(fmt.Sprintf("findings-%s.md", time.Now().Format("2006-01-02"))) + saveDialog.Show() +} + +func (l *lootTab) showLinkedRequest(entry store.LootEntry) { + if entry.HistoryID != nil { + go func() { + tx, err := l.projectStore.GetTransaction(*entry.HistoryID) + if err != nil { + logger.Error("loot: get linked request: %v", err) + return + } + fyne.Do(func() { + l.showRequestResponseDialog( + fmt.Sprintf("Linked Request — %s %s", tx.Method, tx.URL), + FormatStoreRequest(*tx), + FormatStoreResponse(*tx), + tx.Host, tx.Port, tx.TLS, + ) + }) + }() + } else { + host, port, useTLS := internalhttp.ParseHostFromRaw(entry.RawRequest) + l.showRequestResponseDialog("Linked Request", entry.RawRequest, entry.RawResponse, host, port, useTLS) + } +} + +func (l *lootTab) showRequestResponseDialog(title, rawRequest, rawResponse, host string, port int, useTLS bool) { + reqEntry := widgets.NewTextView() + reqEntry.SetWindow(l.win) + reqEntry.SetText(rawRequest) + + respEntry := widgets.NewTextView() + respEntry.SetWindow(l.win) + respEntry.SetText(rawResponse) + + sendBtn := widget.NewButtonWithIcon("Send to Repeater", theme.MailForwardIcon(), nil) + + split := container.NewHSplit( + container.NewBorder(paneHeader("Request"), nil, nil, nil, reqEntry.Build()), + container.NewBorder(paneHeader("Response"), nil, nil, nil, respEntry.Build()), + ) + split.SetOffset(0.5) + + d := dialog.NewCustom(title, "Close", + container.NewBorder(nil, sendBtn, nil, nil, split), + l.win) + + sendBtn.OnTapped = func() { + path := PathOf(internalhttp.ExtractURL(rawRequest)) + if len(path) > 20 { + path = path[:20] + "..." + } + l.repeater.AddTab( + fmt.Sprintf("%s %s", internalhttp.ExtractMethod(rawRequest), path), + host, port, useTLS, rawRequest, + ) + d.Hide() + } + + closeOnEscape(l.win, d.Dismiss) + d.Show() + d.Resize(fyne.NewSize(900, 600)) +} diff --git a/internal/ui/plugins.go b/internal/ui/plugins.go new file mode 100644 index 0000000..cf29f37 --- /dev/null +++ b/internal/ui/plugins.go @@ -0,0 +1,179 @@ +package ui + +import ( + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/storage" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +type pluginsTab struct { + projectStore *store.Store + bus *events.Bus + win fyne.Window +} + +func newPluginsTab(projectStore *store.Store, bus *events.Bus, win fyne.Window) *pluginsTab { + return &pluginsTab{ + projectStore: projectStore, + bus: bus, + win: win, + } +} + +func (t *pluginsTab) build() fyne.CanvasObject { + var plugins []store.PluginEntry + selectedIdx := -1 + + logView := widgets.NewTextView() + logView.SetWindow(t.win) + + toggleBtn := widget.NewButton("Disable", nil) + toggleBtn.Disable() + + var pluginList *widget.List + + refreshLogs := func() { + if selectedIdx < 0 || selectedIdx >= len(plugins) { + logView.SetText("") + return + } + lines := t.projectStore.PluginLogs(plugins[selectedIdx].Name) + logView.SetText(strings.Join(lines, "\n")) + } + + refreshAll := func() { + entries, err := t.projectStore.AllPlugins() + if err != nil { + return + } + plugins = entries + pluginList.Refresh() + + // Keep toggle button label in sync after external state change. + if selectedIdx >= 0 && selectedIdx < len(plugins) { + if plugins[selectedIdx].Enabled { + toggleBtn.SetText("Disable") + } else { + toggleBtn.SetText("Enable") + } + toggleBtn.Enable() + } + + refreshLogs() + } + + pluginList = widget.NewList( + func() int { return len(plugins) }, + func() fyne.CanvasObject { + return container.NewHBox( + widget.NewLabel(""), + widget.NewLabel(""), + ) + }, + func(id widget.ListItemID, item fyne.CanvasObject) { + if id >= len(plugins) { + return + } + p := plugins[id] + box := item.(*fyne.Container) + status := box.Objects[0].(*widget.Label) + name := box.Objects[1].(*widget.Label) + if p.Enabled { + status.SetText("✓") + status.Importance = widget.SuccessImportance + } else { + status.SetText("✗") + status.Importance = widget.DangerImportance + } + name.SetText(p.Name) + }, + ) + + pluginList.OnSelected = func(id widget.ListItemID) { + selectedIdx = id + if id >= len(plugins) { + toggleBtn.Disable() + return + } + if plugins[id].Enabled { + toggleBtn.SetText("Disable") + } else { + toggleBtn.SetText("Enable") + } + toggleBtn.Enable() + refreshLogs() + } + + pluginList.OnUnselected = func(_ widget.ListItemID) { + selectedIdx = -1 + toggleBtn.Disable() + logView.SetText("") + } + + toggleBtn.OnTapped = func() { + if selectedIdx < 0 || selectedIdx >= len(plugins) { + return + } + p := plugins[selectedIdx] + t.bus.EmitSetPluginEnabled(events.SetPluginEnabledEvent{ + Name: p.Name, + Enabled: !p.Enabled, + }) + } + + addBtn := widget.NewButton("Add", func() { + fd := dialog.NewFileOpen(func(uc fyne.URIReadCloser, err error) { + if err != nil || uc == nil { + return + } + uc.Close() + t.bus.EmitLoadPlugin(events.LoadPluginEvent{ + SourcePath: uc.URI().Path(), + }) + }, t.win) + fd.SetFilter(storage.NewExtensionFileFilter([]string{".lua"})) + fd.Show() + }) + + // Listen for plugin state and log changes from the store. + go func() { + for range t.projectStore.PluginEntries { + fyne.Do(refreshAll) + } + }() + + // Initial population. + refreshAll() + + toolbar := container.New( + layout.NewCustomPaddedLayout(8, 8, 0, 0), + container.NewHBox(addBtn, toggleBtn), + ) + + leftPane := container.NewBorder( + newBoldLabel("Plugins"), + toolbar, + nil, nil, + pluginList, + ) + + rightPane := container.NewBorder( + newBoldLabel("Log"), + nil, nil, nil, + logView.Build(), + ) + + split := container.NewHSplit(leftPane, rightPane) + split.SetOffset(0.3) + + return split +} diff --git a/internal/ui/repeater.go b/internal/ui/repeater.go new file mode 100644 index 0000000..72c61ae --- /dev/null +++ b/internal/ui/repeater.go @@ -0,0 +1,224 @@ +package ui + +import ( + "bytes" + "fmt" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/widget" + + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +type repeaterTab struct { + projectStore *store.Store + bus *events.Bus + loot *lootTab + tabs *container.DocTabs + win fyne.Window + tabIDs map[*container.TabItem]int64 + sendFns map[*container.TabItem]func() + loadFns map[*container.TabItem]func() +} + +func newRepeaterTab(projectStore *store.Store, bus *events.Bus, win fyne.Window) *repeaterTab { + return &repeaterTab{ + projectStore: projectStore, + bus: bus, + win: win, + tabIDs: make(map[*container.TabItem]int64), + sendFns: make(map[*container.TabItem]func()), + loadFns: make(map[*container.TabItem]func()), + } +} + +func (r *repeaterTab) closeTab(closed *container.TabItem) { + if tabID, ok := r.tabIDs[closed]; ok { + if err := r.projectStore.DeleteRepeaterTab(tabID); err != nil { + logger.Error("repeater: delete tab: %v", err) + } + delete(r.tabIDs, closed) + } + delete(r.sendFns, closed) + delete(r.loadFns, closed) +} + +func (r *repeaterTab) build() fyne.CanvasObject { + r.tabs = container.NewDocTabs() + r.tabs.SetTabLocation(container.TabLocationTop) + + r.tabs.OnSelected = func(item *container.TabItem) { + if fn, ok := r.loadFns[item]; ok { + fn() + delete(r.loadFns, item) + } + } + + r.tabs.OnClosed = r.closeTab + + r.tabs.CreateTab = func() *container.TabItem { + saved := store.RepeaterTab{ + Name: "New Tab", + Port: 443, + TabType: "http", + } + id, err := r.projectStore.SaveRepeaterTab(saved) + if err != nil { + logger.Error("repeater: save new tab: %v", err) + return container.NewTabItem("New Tab", widget.NewLabel("")) + } + saved.ID = id + return r.buildTabItem(saved) + } + + saved, err := r.projectStore.AllRepeaterTabs() + if err != nil { + logger.Error("repeater: load tabs: %v", err) + } + for _, t := range saved { + r.tabs.Append(r.buildTabItem(t)) + } + if selected := r.tabs.Selected(); selected != nil { + if fn, ok := r.loadFns[selected]; ok { + fn() + delete(r.loadFns, selected) + } + } + + return r.tabs +} + +func (r *repeaterTab) AddTab(name, host string, port int, useTLS bool, rawRequest string) { + saved := store.RepeaterTab{ + Name: name, + Host: host, + Port: port, + TLS: useTLS, + TabType: "http", + RawRequest: rawRequest, + } + id, err := r.projectStore.SaveRepeaterTab(saved) + if err != nil { + logger.Error("repeater: save tab: %v", err) + return + } + saved.ID = id + item := r.buildTabItem(saved) + r.tabs.Append(item) + r.tabs.Select(item) +} + +// AddWSTab opens a new WebSocket repeater tab for the given connection. +func (r *repeaterTab) AddWSTab(name, wsURL string, useTLS bool) { + saved := store.RepeaterTab{ + Name: name, + Host: wsURL, + TLS: useTLS, + TabType: "websocket", + RawRequest: wsURL, + } + id, err := r.projectStore.SaveRepeaterTab(saved) + if err != nil { + logger.Error("repeater: save ws tab: %v", err) + return + } + saved.ID = id + item := r.buildTabItem(saved) + r.tabs.Append(item) + r.tabs.Select(item) +} + +func (r *repeaterTab) buildTabItem(tab store.RepeaterTab) *container.TabItem { + if tab.TabType == "websocket" { + return r.buildWSTabItem(tab) + } + return r.buildHTTPTabItem(tab) +} + +// updateTabName updates the DocTab label after an HTTP send. +func (r *repeaterTab) updateTabName(rawReq string, tabID int64) { + firstLine := strings.SplitN(rawReq, "\n", 2)[0] + parts := strings.Fields(firstLine) + if len(parts) < 2 { + return + } + path := parts[1] + if len(path) > 24 { + path = path[:24] + "..." + } + name := fmt.Sprintf("%s %s", parts[0], path) + + for item, id := range r.tabIDs { + if id == tabID { + item.Text = name + r.tabs.Refresh() + break + } + } + + if err := r.projectStore.RenameRepeaterTab(tabID, name); err != nil { + logger.Error("repeater: rename tab: %v", err) + } +} + +// patchRaw applies plugin modifications surgically to the normalized raw +// request string. Only fields the plugin actually changed are rewritten. +// SendRaw owns Content-Length, Connection, and Accept-Encoding — we leave +// those untouched so its internal rewriteHeaders handles them correctly. +func patchRaw(raw, headerSection string, origBody []byte, origMethod, origPath string, result events.RequestResult) string { + lines := strings.Split(headerSection, "\r\n") + if len(lines) == 0 { + return raw + } + + // Rewrite request line only if method or path actually changed. + newMethod := result.Request.Method + newPath := result.Request.URL.RequestURI() + if newMethod == "" { + newMethod = origMethod + } + if newPath == "" || newPath == "?" { + newPath = origPath + } + if newMethod != origMethod || newPath != origPath { + lines[0] = fmt.Sprintf("%s %s HTTP/1.1", newMethod, newPath) + } + + // Apply header changes: update existing headers in-place (preserving + // original casing and order), append new ones at the end. + for pluginKey, pluginVals := range result.Request.Header { + if len(pluginVals) == 0 { + continue + } + pluginVal := strings.Join(pluginVals, ", ") + found := false + for i := 1; i < len(lines); i++ { + if lines[i] == "" { + break + } + parts := strings.SplitN(lines[i], ": ", 2) + if len(parts) == 2 && strings.EqualFold(parts[0], pluginKey) { + lines[i] = parts[0] + ": " + pluginVal + found = true + break + } + } + if !found { + lines = append(lines, pluginKey+": "+pluginVal) + } + } + + // Replace body only if the plugin actually changed it. + _, body := internalhttp.SplitRaw(raw) + finalBody := body + if !bytes.Equal(result.Body, origBody) { + finalBody = result.Body + } + + return strings.Join(lines, "\r\n") + "\r\n\r\n" + string(finalBody) +} diff --git a/internal/ui/repeater_http.go b/internal/ui/repeater_http.go new file mode 100644 index 0000000..8204a6e --- /dev/null +++ b/internal/ui/repeater_http.go @@ -0,0 +1,310 @@ +package ui + +import ( + "fmt" + "net/http" + "strings" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/events" + internalhttp "github.com/shiv/internal/http" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +type httpHistEntry struct{ req, resp string } + +type httpTabState struct { + r *repeaterTab + tab store.RepeaterTab + tabID int64 + reqEditor *widgets.TextViewEntry + respLabel *widgets.TextView + sendBtn *widget.Button + inspectBtn *widget.Button + sendToLootBtn *widget.Button + prevBtn *widget.Button + nextBtn *widget.Button + cookieJar map[string]string + lastRawRequest string + lastRawResponse string + lastRespBody []byte + sendHistory []httpHistEntry + histIdx int +} + +func (s *httpTabState) loadHistEntry(idx int) { + if idx < 0 || idx >= len(s.sendHistory) { + return + } + s.histIdx = idx + s.reqEditor.SetText(s.sendHistory[idx].req) + s.respLabel.SetText(s.sendHistory[idx].resp) + if s.histIdx > 0 { + s.prevBtn.Enable() + } else { + s.prevBtn.Disable() + } + if s.histIdx < len(s.sendHistory)-1 { + s.nextBtn.Enable() + } else { + s.nextBtn.Disable() + } +} + +// applyBusRequest runs the bus plugin pipeline on rawReq. Returns dropped=true +// when a plugin rejects the request (UI already updated); caller must return. +// busOK=false (with dropped=false) means no bus or NewRequest failed — caller +// should still send, just skip EmitResponse. +func (s *httpTabState) applyBusRequest(normalized, headerSection string, rawBody []byte, rawMethod, rawPath, scheme, addr string) (finalRaw string, emitReq *http.Request, emitBody []byte, busOK, dropped bool) { + finalRaw = normalized + if s.r.bus == nil { + return + } + busReq, err := http.NewRequest(rawMethod, fmt.Sprintf("%s://%s%s", scheme, addr, rawPath), nil) + if err != nil { + return + } + busReq.Header = internalhttp.ParseRawHeaders(headerSection) + result := s.r.bus.EmitRequest(events.RequestEvent{Request: busReq, Body: rawBody}) + if result.Drop { + fyne.Do(func() { + s.sendBtn.Enable() + s.respLabel.SetText("HTTP/1.1 403 Forbidden\r\n\r\nrequest dropped by plugin") + }) + dropped = true + return + } + // Apply only what the plugin actually changed back onto + // the normalized raw string. SendRaw owns Content-Length, + // Connection, and Accept-Encoding — we never touch those. + finalRaw = patchRaw(normalized, headerSection, rawBody, rawMethod, rawPath, result) + return finalRaw, result.Request, result.Body, true, false +} + +// applySendResult runs inside fyne.Do after a successful or failed send. +func (s *httpTabState) applySendResult(rawReq string, sendResult *internalhttp.RawResponse, err error) { + s.sendBtn.Enable() + if err != nil { + s.respLabel.SetText("Error: " + err.Error()) + logger.Error("repeater: send: %v", err) + return + } + s.respLabel.SetText(sendResult.Raw) + for _, c := range sendResult.Cookies { + s.cookieJar[c.Name] = c.Value + } + s.lastRawRequest = rawReq + s.lastRawResponse = sendResult.Raw + s.lastRespBody = sendResult.Body + s.r.updateTabName(rawReq, s.tabID) + s.inspectBtn.Enable() + s.sendToLootBtn.Enable() + if saveErr := s.r.projectStore.UpdateRepeaterTab(s.tabID, rawReq, sendResult.Raw); saveErr != nil { + logger.Error("repeater: update tab: %v", saveErr) + } + + // Record in per-tab send history. + const maxHistory = 50 + s.sendHistory = append(s.sendHistory, httpHistEntry{req: rawReq, resp: sendResult.Raw}) + if len(s.sendHistory) > maxHistory { + s.sendHistory = s.sendHistory[len(s.sendHistory)-maxHistory:] + } + s.histIdx = len(s.sendHistory) - 1 + if s.prevBtn != nil { + if s.histIdx > 0 { + s.prevBtn.Enable() + } + s.nextBtn.Disable() + } +} + +func (s *httpTabState) doSend() { + if s.sendBtn.Disabled() { + return + } + rawReq := s.reqEditor.GetText() + host, port, useTLS := internalhttp.ParseHostFromRaw(rawReq) + if host == "" { + s.respLabel.SetText("Error: no Host header found in request") + return + } + // ParseHostFromRaw defaults to port 80 / plain HTTP when the Host header + // has no explicit port and the request line has no scheme prefix. + // Use the tab's saved port/TLS as the authoritative values in that case. + if port == 80 && !useTLS && s.tab.Port != 0 { + port = s.tab.Port + useTLS = s.tab.TLS + } + s.sendBtn.Disable() + go func() { + start := time.Now() + scheme := "http" + if useTLS { + scheme = "https" + } + addr := fmt.Sprintf("%s:%d", host, port) + + // normalizeRaw ensures CRLF in headers before any parsing. + // SendRaw does the same internally — we must match it here so + // that our header/body split is consistent with what it sends. + normalized := internalhttp.NormalizeLineEndings(rawReq) + headerSection, rawBody := internalhttp.SplitRaw(normalized) + rawMethod := internalhttp.ExtractMethod(normalized) + rawPath := internalhttp.ExtractURL(normalized) + + finalRawReq, emitReq, emitBody, busOK, dropped := s.applyBusRequest(normalized, headerSection, rawBody, rawMethod, rawPath, scheme, addr) + if dropped { + return + } + + sendResult, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, + Port: port, + TLS: useTLS, + RawReq: finalRawReq, + CookieJar: s.cookieJar, + }) + + elapsed := time.Since(start).Milliseconds() + + if err == nil && busOK { + s.r.bus.EmitResponse(events.ResponseEvent{ + Timestamp: start, + Host: addr, + Proto: "HTTP/1.1", + Method: emitReq.Method, + URL: emitReq.URL.String(), + ReqHeaders: emitReq.Header, + ReqBody: emitBody, + StatusCode: sendResult.StatusCode, + RespHeaders: internalhttp.ParseRawHeaders(sendResult.Raw), + RespBody: sendResult.Body, + DurationMs: elapsed, + TLS: useTLS, + }) + } + + fyne.Do(func() { s.applySendResult(rawReq, sendResult, err) }) + }() +} + +func (s *httpTabState) buildEditors() { + s.reqEditor = widgets.NewTextViewEntry() + s.reqEditor.SetWindow(s.r.win) + s.reqEditor.SetPlaceHolder("Paste or edit raw HTTP request here...") + + s.respLabel = widgets.NewTextView() + s.respLabel.SetWindow(s.r.win) +} + +func (s *httpTabState) buildActionButtons() { + s.inspectBtn = widget.NewButtonWithIcon("Inspector", AppIcon("inspector"), func() { + _, reqBody := internalhttp.SplitRaw(s.lastRawRequest) + showInspectorDialog(store.Transaction{ + ReqHeaders: internalhttp.ParseRawHeaders(s.lastRawRequest), + ReqBody: reqBody, + RespHeaders: internalhttp.ParseRawHeaders(s.lastRawResponse), + RespBody: s.lastRespBody, + }, s.r.win) + }) + s.inspectBtn.Disable() + + s.sendToLootBtn = widget.NewButtonWithIcon("Loot", AppIcon("loot"), func() { + if s.r.loot != nil { + s.r.loot.showAddDialog(nil, s.lastRawRequest, s.lastRawResponse) + } + }) + s.sendToLootBtn.Disable() +} + +func (s *httpTabState) buildSendHistory() { + s.prevBtn = widget.NewButton("◀", func() { s.loadHistEntry(s.histIdx - 1) }) + s.nextBtn = widget.NewButton("▶", func() { s.loadHistEntry(s.histIdx + 1) }) + s.prevBtn.Disable() + s.nextBtn.Disable() +} + +func (s *httpTabState) buildToolbar() fyne.CanvasObject { + cloneBtn := widget.NewButtonWithIcon("Clone", theme.ContentCopyIcon(), func() { + raw := s.reqEditor.GetText() + firstLine := strings.SplitN(raw, "\n", 2)[0] + parts := strings.Fields(firstLine) + name := s.tab.Name + if len(parts) >= 2 { + path := parts[1] + if len(path) > 20 { + path = path[:20] + "..." + } + name = fmt.Sprintf("%s %s", parts[0], path) + } + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + s.r.AddTab(name, host, port, useTLS, raw) + }) + + return container.NewBorder(nil, nil, + container.NewHBox(s.sendBtn, cloneBtn, s.prevBtn, s.nextBtn), + s.sendToLootBtn, + widget.NewLabel(""), + ) +} + +func (s *httpTabState) buildLayout() fyne.CanvasObject { + toolbar := s.buildToolbar() + + copyReqBtn := widget.NewButton("Copy", func() { + fyne.CurrentApp().Clipboard().SetContent(s.reqEditor.GetText()) + }) + copyRespBtn := widget.NewButton("Copy", func() { + fyne.CurrentApp().Clipboard().SetContent(s.lastRawResponse) + }) + + reqPane := container.NewBorder( + paneHeader("Request", copyReqBtn), + nil, nil, nil, s.reqEditor.Build(), + ) + respPane := container.NewBorder( + paneHeader("Response", copyRespBtn, s.inspectBtn), + nil, nil, nil, s.respLabel.Build(), + ) + + split := container.NewHSplit(reqPane, respPane) + split.SetOffset(0.5) + + return container.NewBorder(container.New(layout.NewCustomPaddedLayout(8, 0, 0, 0), toolbar), nil, nil, nil, split) +} + +func (r *repeaterTab) buildHTTPTabItem(tab store.RepeaterTab) *container.TabItem { + s := &httpTabState{ + r: r, + tab: tab, + tabID: tab.ID, + cookieJar: make(map[string]string), + histIdx: -1, + } + + s.buildEditors() + s.buildActionButtons() + s.buildSendHistory() + + s.sendBtn = widget.NewButtonWithIcon("Send", theme.MailSendIcon(), nil) + s.sendBtn.Importance = widget.HighImportance + s.sendBtn.OnTapped = s.doSend + + content := s.buildLayout() + tabItem := container.NewTabItem(tab.Name, content) + + r.sendFns[tabItem] = s.doSend + r.tabIDs[tabItem] = s.tabID + s.reqEditor.SetText(tab.RawRequest) + r.loadFns[tabItem] = func() {} + + return tabItem +} diff --git a/internal/ui/repeater_test.go b/internal/ui/repeater_test.go new file mode 100644 index 0000000..c2c54af --- /dev/null +++ b/internal/ui/repeater_test.go @@ -0,0 +1,378 @@ +package ui + +import ( + "bytes" + "compress/gzip" + "fmt" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gorilla/websocket" + internalhttp "github.com/shiv/internal/http" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── ParseHostFromRaw ────────────────────────────────────────────────────────── + +func TestParseHostFromRaw_HostOnly(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n" + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_HostWithPort80(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:80\r\n\r\n" + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_HostWithPort443(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:443\r\n\r\n" + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) + assert.Equal(t, 443, port) + assert.True(t, useTLS) +} + +func TestParseHostFromRaw_CustomPort(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHost: example.com:8080\r\n\r\n" + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) + assert.Equal(t, 8080, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_NoHostHeader(t *testing.T) { + raw := "GET / HTTP/1.1\r\nContent-Type: text/plain\r\n\r\n" + host, port, useTLS := internalhttp.ParseHostFromRaw(raw) + assert.Empty(t, host) + assert.Equal(t, 0, port) + assert.False(t, useTLS) +} + +func TestParseHostFromRaw_CaseInsensitive(t *testing.T) { + raw := "GET / HTTP/1.1\r\nHOST: example.com\r\n\r\n" + host, _, _ := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) +} + +func TestParseHostFromRaw_LFOnly(t *testing.T) { + raw := "GET / HTTP/1.1\nHost: example.com\n\n" + host, port, _ := internalhttp.ParseHostFromRaw(raw) + assert.Equal(t, "example.com", host) + assert.Equal(t, 80, port) +} + +// ── Decompress ──────────────────────────────────────────────────────────────── + +func TestDecompressRepeaterBody_Gzip(t *testing.T) { + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write([]byte("hello gzip")) + require.NoError(t, err) + w.Close() + + hdr := http.Header{"Content-Encoding": []string{"gzip"}} + out := internalhttp.Decompress(hdr, buf.Bytes()) + assert.Equal(t, "hello gzip", string(out)) +} + +func TestDecompressRepeaterBody_NoEncoding(t *testing.T) { + body := []byte("plain body") + out := internalhttp.Decompress(http.Header{}, body) + assert.Equal(t, body, out) +} + +func TestDecompressRepeaterBody_BadGzip(t *testing.T) { + body := []byte("not gzip data") + hdr := http.Header{"Content-Encoding": []string{"gzip"}} + out := internalhttp.Decompress(hdr, body) + assert.Equal(t, body, out) +} + +// ── SendRaw ─────────────────────────────────────────────────────────────────── + +func startMockHTTPServer(t *testing.T, response string) (host string, port int) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + buf := make([]byte, 4096) + conn.Read(buf) + fmt.Fprint(conn, response) + }() + + t.Cleanup(func() { ln.Close() }) + + addr := ln.Addr().(*net.TCPAddr) + return "127.0.0.1", addr.Port +} + +func TestSendRawRequest_200OK(t *testing.T) { + body := "Hello, World!" + mockResp := fmt.Sprintf( + "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", + len(body), body, + ) + host, port := startMockHTTPServer(t, mockResp) + + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host, port) + result, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, Port: port, TLS: false, RawReq: raw, + }) + require.NoError(t, err) + assert.Contains(t, result.Raw, "200 OK") + assert.Contains(t, result.Raw, "Hello, World!") +} + +func TestSendRawRequest_404(t *testing.T) { + mockResp := "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n" + host, port := startMockHTTPServer(t, mockResp) + + raw := fmt.Sprintf("GET /missing HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host, port) + result, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, Port: port, TLS: false, RawReq: raw, + }) + require.NoError(t, err) + assert.Contains(t, result.Raw, "404") +} + +func TestSendRawRequest_WithBody(t *testing.T) { + mockResp := "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK" + host, port := startMockHTTPServer(t, mockResp) + + raw := fmt.Sprintf( + "POST /api HTTP/1.1\r\nHost: %s:%d\r\nContent-Type: application/json\r\n\r\n{\"key\":\"val\"}", + host, port, + ) + result, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, Port: port, TLS: false, RawReq: raw, + }) + require.NoError(t, err) + assert.Contains(t, result.Raw, "200 OK") +} + +func TestSendRawRequest_ConnectionRefused(t *testing.T) { + _, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: "127.0.0.1", Port: 1, TLS: false, + RawReq: "GET / HTTP/1.1\r\nHost: 127.0.0.1:1\r\n\r\n", + }) + assert.Error(t, err) +} + +func TestSendRawRequest_StripAcceptEncoding(t *testing.T) { + received := make(chan string, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { ln.Close() }) + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + received <- string(buf[:n]) + fmt.Fprint(conn, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + }() + + addr := ln.Addr().(*net.TCPAddr) + raw := fmt.Sprintf( + "GET / HTTP/1.1\r\nHost: 127.0.0.1:%d\r\nAccept-Encoding: gzip, deflate\r\n\r\n", + addr.Port, + ) + internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: "127.0.0.1", Port: addr.Port, TLS: false, RawReq: raw, + }) + + req := <-received + assert.NotContains(t, req, "Accept-Encoding") +} + +func TestSendRawRequest_MultipleRequests_IndependentConnections(t *testing.T) { + // Verify that two sequential SendRaw calls use independent connections + // and both get correct responses. + responses := []string{ + "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nfirst", + "HTTP/1.1 201 Created\r\nContent-Length: 6\r\n\r\nsecond", + } + + for i, expected := range responses { + host, port := startMockHTTPServer(t, expected) + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host, port) + result, err := internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: host, Port: port, TLS: false, RawReq: raw, + }) + require.NoError(t, err, "request %d failed", i) + if i == 0 { + assert.Contains(t, result.Raw, "200") + assert.Contains(t, result.Raw, "first") + } else { + assert.Contains(t, result.Raw, "201") + assert.Contains(t, result.Raw, "second") + } + } +} + +func TestSendRawRequest_CookieJarApplied(t *testing.T) { + received := make(chan string, 1) + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { ln.Close() }) + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + received <- string(buf[:n]) + fmt.Fprint(conn, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + }() + + addr := ln.Addr().(*net.TCPAddr) + raw := fmt.Sprintf("GET / HTTP/1.1\r\nHost: 127.0.0.1:%d\r\n\r\n", addr.Port) + internalhttp.SendRaw(internalhttp.RawRequestOptions{ + Host: "127.0.0.1", + Port: addr.Port, + TLS: false, + RawReq: raw, + CookieJar: map[string]string{"session": "abc123", "user": "test"}, + }) + + req := <-received + assert.Contains(t, req, "Cookie:") + assert.Contains(t, req, "session=abc123") +} + +// ── WebSocket send via repeater ─────────────────────────────────────────────── + +// wsEchoServer starts a plain (non-TLS) WebSocket echo server for repeater tests. +func wsEchoServer(t *testing.T) string { + t.Helper() + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + mt, msg, err := conn.ReadMessage() + if err != nil { + return + } + if err := conn.WriteMessage(mt, msg); err != nil { + return + } + } + })) + t.Cleanup(srv.Close) + // Convert http:// to ws:// + return strings.Replace(srv.URL, "http://", "ws://", 1) +} + +func TestRepeaterWS_SendRaw_EchoesMessage(t *testing.T) { + wsURL := wsEchoServer(t) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(wsURL+"/", nil) + require.NoError(t, err) + defer conn.Close() + + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("repeater test"))) + _, payload, err := conn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "repeater test", string(payload)) +} + +func TestRepeaterWS_SendRaw_MultipleMessages(t *testing.T) { + wsURL := wsEchoServer(t) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(wsURL+"/", nil) + require.NoError(t, err) + defer conn.Close() + + messages := []string{"first", "second", "third"} + for _, msg := range messages { + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte(msg))) + _, payload, err := conn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, msg, string(payload)) + } +} + +func TestRepeaterWS_SendRaw_EmptyMessage(t *testing.T) { + wsURL := wsEchoServer(t) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(wsURL+"/", nil) + require.NoError(t, err) + defer conn.Close() + + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte(""))) + _, payload, err := conn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "", string(payload)) +} + +func TestRepeaterWS_SendRaw_BinaryMessage(t *testing.T) { + wsURL := wsEchoServer(t) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(wsURL+"/", nil) + require.NoError(t, err) + defer conn.Close() + + payload := []byte{0x00, 0x01, 0x02, 0xFF} + require.NoError(t, conn.WriteMessage(websocket.BinaryMessage, payload)) + msgType, received, err := conn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, websocket.BinaryMessage, msgType) + assert.Equal(t, payload, received) +} + +func TestRepeaterWS_ConnectionRefused(t *testing.T) { + dialer := websocket.Dialer{} + _, _, err := dialer.Dial("ws://127.0.0.1:1/", nil) + assert.Error(t, err) +} + +func TestRepeaterWS_CleanClose(t *testing.T) { + wsURL := wsEchoServer(t) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(wsURL+"/", nil) + require.NoError(t, err) + + // Send a message, receive echo, then close cleanly. + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("bye"))) + _, payload, err := conn.ReadMessage() + require.NoError(t, err) + assert.Equal(t, "bye", string(payload)) + + // Clean close. + err = conn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done")) + require.NoError(t, err) + conn.Close() +} diff --git a/internal/ui/repeater_ws.go b/internal/ui/repeater_ws.go new file mode 100644 index 0000000..4a61263 --- /dev/null +++ b/internal/ui/repeater_ws.go @@ -0,0 +1,418 @@ +package ui + +import ( + "crypto/tls" + "net" + "net/http" + "strings" + "time" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/layout" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/gorilla/websocket" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +type wsConnState int + +const ( + wsDisconnected wsConnState = iota + wsConnecting + wsConnected +) + +type wsFrame struct { + direction string + payload string + ts time.Time +} + +type wsHeaderRow struct { + key string + value string +} + +type wsTabState struct { + r *repeaterTab + tab store.RepeaterTab + tabID int64 + wsConn *websocket.Conn + connState wsConnState + connID uint64 + customHeaders []wsHeaderRow + frames []wsFrame + frameTable *widgets.DataTable + payloadView *widgets.TextView + connectBtn *widget.Button + sendFrameBtn *widget.Button + urlEntry *widget.Entry +} + +func (s *wsTabState) disconnect() { + if s.wsConn != nil { + s.wsConn.Close() + s.wsConn = nil + } + s.connState = wsDisconnected + s.updateConnectBtn() +} + +func (s *wsTabState) updateConnectBtn() { + switch s.connState { + case wsDisconnected: + s.connectBtn.SetText("Connect") + s.connectBtn.Importance = widget.HighImportance + s.connectBtn.Enable() + s.sendFrameBtn.Disable() + case wsConnecting: + s.connectBtn.SetText("Connecting...") + s.connectBtn.Importance = widget.MediumImportance + s.connectBtn.Disable() + s.sendFrameBtn.Disable() + case wsConnected: + s.connectBtn.SetText("Disconnect") + s.connectBtn.Importance = widget.DangerImportance + s.connectBtn.Enable() + s.sendFrameBtn.Enable() + } + s.connectBtn.Refresh() +} + +func (s *wsTabState) startReadLoop(conn *websocket.Conn) { + go func() { + for { + msgType, payload, err := conn.ReadMessage() + if err != nil { + fyne.Do(func() { + if s.wsConn == conn { + s.wsConn = nil + s.connState = wsDisconnected + s.updateConnectBtn() + s.frames = append(s.frames, wsFrame{ + direction: "⚠ Closed", + payload: err.Error(), + ts: time.Now(), + }) + s.frameTable.Refresh() + s.frameTable.OnSelect(len(s.frames) - 1) + } + }) + return + } + finalPayload := payload + if s.r.bus != nil { + result := s.r.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: s.connID, + Timestamp: time.Now(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketOpcode(msgType), + Payload: payload, + }) + finalPayload = result.Payload + } + f := wsFrame{direction: "← Recv", payload: string(finalPayload), ts: time.Now()} + fyne.Do(func() { + s.frames = append(s.frames, f) + s.frameTable.Refresh() + s.frameTable.ScrollToRow(len(s.frames) - 1) + s.frameTable.OnSelect(len(s.frames) - 1) + }) + } + }() +} + +func (s *wsTabState) showHeadersDialog() { + editHeaders := make([]wsHeaderRow, len(s.customHeaders)) + copy(editHeaders, s.customHeaders) + + var rows *fyne.Container + var rebuildRows func() + + rebuildRows = func() { + rows.RemoveAll() + for i := range editHeaders { + idx := i + keyEntry := widget.NewEntry() + keyEntry.SetPlaceHolder("Header-Name") + keyEntry.SetText(editHeaders[idx].key) + keyEntry.OnChanged = func(v string) { editHeaders[idx].key = v } + + valEntry := widget.NewEntry() + valEntry.SetPlaceHolder("value") + valEntry.SetText(editHeaders[idx].value) + valEntry.OnChanged = func(v string) { editHeaders[idx].value = v } + + deleteBtn := widget.NewButtonWithIcon("", theme.DeleteIcon(), func() { + editHeaders = append(editHeaders[:idx], editHeaders[idx+1:]...) + rebuildRows() + }) + + rows.Add(container.NewBorder(nil, nil, nil, deleteBtn, + container.NewGridWithColumns(2, keyEntry, valEntry), + )) + } + rows.Refresh() + } + + rows = container.NewVBox() + rebuildRows() + + addBtn := widget.NewButtonWithIcon("Add Header", theme.ContentAddIcon(), func() { + editHeaders = append(editHeaders, wsHeaderRow{}) + rebuildRows() + }) + + scrolled := container.NewVScroll(rows) + scrolled.SetMinSize(fyne.NewSize(500, 240)) + + body := container.NewBorder(nil, addBtn, nil, nil, scrolled) + + dialog.ShowCustomConfirm("Edit Request Headers", "Save", "Cancel", body, func(save bool) { + if !save { + return + } + s.customHeaders = s.customHeaders[:0] + for _, h := range editHeaders { + if strings.TrimSpace(h.key) != "" { + s.customHeaders = append(s.customHeaders, h) + } + } + }, s.r.win) +} + +func (s *wsTabState) buildFrameTable() { + frameColumns := []widgets.DataTableColumn{ + {Header: "Dir", Width: 80}, + {Header: "Time", Width: 120}, + {Header: "Payload", Width: 300}, + } + + s.frameTable = widgets.NewDataTable() + s.frameTable.SetWindow(s.r.win) + s.frameTable.Columns = frameColumns + s.frameTable.RowCount = func() int { return len(s.frames) } + s.frameTable.RowID = func(row int) int64 { return int64(row) } + s.frameTable.CellValue = func(row, col int) string { + if row >= len(s.frames) { + return "" + } + f := s.frames[row] + switch col { + case 0: + return f.direction + case 1: + return f.ts.Format("15:04:05.000") + case 2: + return f.payload + } + return "" + } + s.frameTable.CellStyle = func(row, col int) widget.Importance { + if col != 0 || row >= len(s.frames) { + return widget.MediumImportance + } + if s.frames[row].direction == "→ Sent" { + return widget.HighImportance + } + return widget.SuccessImportance + } +} + +func (s *wsTabState) buildPayloadView() { + s.payloadView = widgets.NewTextView() + s.payloadView.SetWindow(s.r.win) + s.payloadView.SetPlaceHolder("Select a frame to view full payload...") + + s.frameTable.OnSelect = func(row int) { + if row >= len(s.frames) { + return + } + p := s.frames[row].payload + if p == "" { + s.payloadView.SetText("(empty)") + } else { + s.payloadView.SetText(p) + } + } +} + +func (s *wsTabState) buildConnectHandler() { + s.connectBtn.OnTapped = func() { + if s.connState == wsConnected { + s.disconnect() + return + } + rawURL := strings.TrimSpace(s.urlEntry.Text) + if rawURL == "" { + return + } + s.connState = wsConnecting + s.updateConnectBtn() + if err := s.r.projectStore.UpdateRepeaterTab(s.tabID, rawURL, ""); err != nil { + logger.Error("repeater: update ws tab: %v", err) + } + + reqHeader := http.Header{} + for _, h := range s.customHeaders { + if strings.TrimSpace(h.key) != "" { + cleaned := strings.NewReplacer("\n", "", "\r", "").Replace(h.value) + reqHeader.Set(h.key, cleaned) + } + } + + go func() { + dialer := websocket.Dialer{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec + HandshakeTimeout: 10 * time.Second, + NetDial: func(network, addr string) (net.Conn, error) { + return net.DialTimeout(network, addr, 10*time.Second) + }, + } + conn, _, err := dialer.Dial(rawURL, reqHeader) + fyne.Do(func() { + if err != nil { + s.connState = wsDisconnected + s.updateConnectBtn() + s.frames = append(s.frames, wsFrame{direction: "⚠ Error", payload: err.Error(), ts: time.Now()}) + s.frameTable.Refresh() + s.frameTable.OnSelect(len(s.frames) - 1) + return + } + s.wsConn = conn + s.connState = wsConnected + s.updateConnectBtn() + if s.r.bus != nil { + s.connID = s.r.bus.EmitWebSocketConnection(events.WebSocketConnectionEvent{ + Host: rawURL, + URL: rawURL, + TLS: strings.HasPrefix(rawURL, "wss://"), + Timestamp: time.Now(), + }) + } + s.frames = append(s.frames, wsFrame{direction: "✓ Connected", payload: rawURL, ts: time.Now()}) + s.frameTable.Refresh() + s.frameTable.OnSelect(len(s.frames) - 1) + s.startReadLoop(conn) + }) + }() + } +} + +func (s *wsTabState) buildSendHandler(payloadEntry *widget.Entry) { + s.sendFrameBtn.OnTapped = func() { + if s.wsConn == nil || s.connState != wsConnected { + return + } + payload := strings.ReplaceAll(strings.ReplaceAll(payloadEntry.Text, "\r\n", ""), "\n", "") + if payload == "" { + return + } + conn := s.wsConn + go func() { + finalPayload := []byte(payload) + if s.r.bus != nil { + result := s.r.bus.EmitWebSocketFrame(events.WebSocketFrameEvent{ + ConnectionID: s.connID, + Timestamp: time.Now(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: finalPayload, + }) + finalPayload = result.Payload + } + err := conn.WriteMessage(websocket.TextMessage, finalPayload) + logger.Info("ws: sending payload bytes: %q", finalPayload) + fyne.Do(func() { + if err != nil { + s.frames = append(s.frames, wsFrame{direction: "⚠ Error", payload: "send failed: " + err.Error(), ts: time.Now()}) + s.frameTable.Refresh() + s.frameTable.OnSelect(len(s.frames) - 1) + return + } + s.frames = append(s.frames, wsFrame{direction: "→ Sent", payload: string(finalPayload), ts: time.Now()}) + s.frameTable.Refresh() + s.frameTable.ScrollToRow(len(s.frames) - 1) + s.frameTable.OnSelect(len(s.frames) - 1) + }) + }() + } +} + +func (s *wsTabState) buildLayout(payloadEntry *widget.Entry, clearBtn, headersBtn *widget.Button) fyne.CanvasObject { + toolbar := container.NewBorder(nil, nil, + container.NewHBox(s.connectBtn, clearBtn, headersBtn), + nil, + s.urlEntry, + ) + + sendBar := container.NewBorder(nil, nil, nil, s.sendFrameBtn, payloadEntry) + + framePane := container.NewBorder(paneHeader("Frames"), nil, nil, nil, s.frameTable.Build()) + payloadPane := container.NewBorder(paneHeader("Payload"), nil, nil, nil, s.payloadView.Build()) + + frameSplit := container.NewVSplit(framePane, payloadPane) + frameSplit.SetOffset(0.6) + + return container.NewBorder( + container.New(layout.NewCustomPaddedLayout(8, 0, 0, 0), toolbar), + container.NewBorder(nil, nil, nil, nil, sendBar), + nil, nil, + frameSplit, + ) +} + +func (r *repeaterTab) buildWSTabItem(tab store.RepeaterTab) *container.TabItem { + s := &wsTabState{ + r: r, + tab: tab, + tabID: tab.ID, + } + + s.urlEntry = widget.NewEntry() + s.urlEntry.SetPlaceHolder("wss://host:port/path") + if tab.RawRequest != "" { + s.urlEntry.SetText(tab.RawRequest) + } + + s.connectBtn = widget.NewButtonWithIcon("Connect", AppIcon("web"), nil) + s.connectBtn.Importance = widget.HighImportance + + s.sendFrameBtn = widget.NewButtonWithIcon("Send", theme.MailSendIcon(), nil) + s.sendFrameBtn.Disable() + + s.buildFrameTable() + s.buildPayloadView() + s.buildConnectHandler() + + payloadEntry := widget.NewMultiLineEntry() + payloadEntry.SetPlaceHolder("Message to send...") + payloadEntry.SetMinRowsVisible(3) + + s.buildSendHandler(payloadEntry) + + clearBtn := widget.NewButtonWithIcon("Clear", AppIcon("delete"), func() { + s.frames = nil + s.frameTable.Refresh() + s.payloadView.SetText("") + }) + + headersBtn := widget.NewButtonWithIcon("Headers", theme.ListIcon(), func() { + s.showHeadersDialog() + }) + + content := s.buildLayout(payloadEntry, clearBtn, headersBtn) + tabItem := container.NewTabItem(tab.Name, content) + r.tabIDs[tabItem] = s.tabID + r.loadFns[tabItem] = func() {} + + return tabItem +} diff --git a/internal/ui/scope.go b/internal/ui/scope.go new file mode 100644 index 0000000..037472c --- /dev/null +++ b/internal/ui/scope.go @@ -0,0 +1,126 @@ +package ui + +import ( + "net" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +func showScopeDialog(projectStore *store.Store, win fyne.Window, onChanged func()) { + entries, err := projectStore.AllScopeEntries() + if err != nil { + logger.Error("scope: load entries: %v", err) + entries = nil + } + + // we keep a local copy to drive the list + type row struct { + id int64 + host string + } + rows := make([]row, len(entries)) + for i, entry := range entries { + rows[i] = row{entry.ID, entry.Host} + } + + var list *widget.List + var deleteBtn *widget.Button + selectedIdx := -1 + + list = widget.NewList( + func() int { return len(rows) }, + func() fyne.CanvasObject { + return widget.NewLabel("") + }, + func(i widget.ListItemID, obj fyne.CanvasObject) { + obj.(*widget.Label).SetText(rows[i].host) + }, + ) + list.OnSelected = func(i widget.ListItemID) { + selectedIdx = i + deleteBtn.Enable() + } + list.OnUnselected = func(_ widget.ListItemID) { + selectedIdx = -1 + deleteBtn.Disable() + } + + hostEntry := widget.NewEntry() + hostEntry.SetPlaceHolder("example.com") + + addBtn := widget.NewButtonWithIcon("Add", AppIcon("save"), func() { + host := strings.TrimSpace(hostEntry.Text) + if host == "" { + return + } + // strip protocol if someone pastes a URL + host = strings.TrimPrefix(host, "https://") + host = strings.TrimPrefix(host, "http://") + // strip path + if idx := strings.Index(host, "/"); idx >= 0 { + host = host[:idx] + } + // strip port, avoid breaking ipv6 + if hostname, _, err := net.SplitHostPort(host); err == nil { + host = hostname + } + if err := projectStore.AddScopeEntry(host); err != nil { + logger.Error("scope: add entry: %v", err) + return + } + newEntries, _ := projectStore.AllScopeEntries() + rows = make([]row, len(newEntries)) + for i, entry := range newEntries { + rows[i] = row{entry.ID, entry.Host} + } + hostEntry.SetText("") + list.Refresh() + onChanged() + }) + addBtn.Importance = widget.HighImportance + + deleteBtn = widget.NewButtonWithIcon("Remove", theme.DeleteIcon(), func() { + if selectedIdx < 0 || selectedIdx >= len(rows) { + return + } + entryId := rows[selectedIdx].id + if err := projectStore.DeleteScopeEntry(entryId); err != nil { + logger.Error("scope: delete entry: %v", err) + return + } + newEntries, _ := projectStore.AllScopeEntries() + rows = make([]row, len(newEntries)) + for i, e := range newEntries { + rows[i] = row{e.ID, e.Host} + } + selectedIdx = -1 + deleteBtn.Disable() + list.Refresh() + onChanged() + }) + deleteBtn.Disable() + + inputRow := container.NewBorder(nil, nil, nil, addBtn, hostEntry) + + content := container.NewBorder( + inputRow, + deleteBtn, + nil, nil, + list, + ) + content.Resize(fyne.NewSize(400, 300)) + + scopeDialog := dialog.NewCustom("Scope", "Close", content, win) + scopeDialog.Resize(fyne.NewSize(400, 350)) + + closeOnEscape(win, scopeDialog.Dismiss) + scopeDialog.Show() +} diff --git a/internal/ui/settings.go b/internal/ui/settings.go new file mode 100644 index 0000000..1565066 --- /dev/null +++ b/internal/ui/settings.go @@ -0,0 +1,401 @@ +package ui + +import ( + "fmt" + "io" + "os" + "path/filepath" + "strconv" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/dialog" + "fyne.io/fyne/v2/storage" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/cert" + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +var settingsWin fyne.Window + +func showSettingsDialog(fyneApp fyne.App, parentWin fyne.Window, bus *events.Bus, keybinds *Keybinds) { + if settingsWin != nil { + settingsWin.RequestFocus() + return + } + + tabs := container.NewAppTabs( + container.NewTabItem("Certificate", buildCertificateTab(parentWin)), + container.NewTabItem("Proxy", buildProxyTab(fyneApp, bus)), + container.NewTabItem("Browser", buildBrowserTab(fyneApp, parentWin)), + container.NewTabItem("Appearance & Shortcuts", buildAppearanceTab(fyneApp, keybinds)), + ) + tabs.SetTabLocation(container.TabLocationTop) + + settingsWin = fyneApp.NewWindow("Settings") + settingsWin.SetContent(container.NewPadded(tabs)) + settingsWin.Resize(fyne.NewSize(480, 420)) + settingsWin.SetOnClosed(func() { settingsWin = nil }) + closeOnEscape(settingsWin, settingsWin.Close) + settingsWin.Show() +} + +func buildCertificateTab(parentWin fyne.Window) fyne.CanvasObject { + statusLabel := widget.NewLabel("") + statusLabel.Wrapping = fyne.TextWrapBreak + + installBtn := widget.NewButton("Install CA Certificate", func() { + certAuthority, err := cert.Load() + if err != nil { + statusLabel.SetText("Error loading CA: " + err.Error()) + logger.Error("settings: load CA: %v", err) + return + } + msg, err := certAuthority.InstallCA() + if err != nil { + statusLabel.SetText("Installation failed:\n" + err.Error()) + logger.Error("settings: install CA: %v", err) + return + } + statusLabel.SetText(msg) + }) + installBtn.Importance = widget.HighImportance + + exportBtn := widget.NewButton("Export Certificate (PEM)", func() { + certAuthority, err := cert.Load() + if err != nil { + statusLabel.SetText("Error loading CA: " + err.Error()) + logger.Error("settings: load CA for export: %v", err) + return + } + pemBytes, err := certAuthority.CertPEM() + if err != nil { + statusLabel.SetText("Error reading certificate: " + err.Error()) + logger.Error("settings: cert PEM: %v", err) + return + } + saveDialog := dialog.NewFileSave(func(writeCloser fyne.URIWriteCloser, saveErr error) { + if saveErr != nil || writeCloser == nil { + return + } + defer writeCloser.Close() + if _, writeErr := writeCloser.Write(pemBytes); writeErr != nil { + statusLabel.SetText("Export failed: " + writeErr.Error()) + logger.Error("settings: export cert: %v", writeErr) + return + } + statusLabel.SetText("Certificate exported successfully.") + }, parentWin) + saveDialog.SetFileName("shiv-ca.pem") + saveDialog.Show() + }) + + return container.NewVBox( + newBoldLabel("CA Certificate"), + widget.NewLabel("Install the Shiv CA into your system trust store so browsers accept intercepted traffic."), + installBtn, + widget.NewSeparator(), + widget.NewLabel("Export the CA certificate to import manually into a browser or device."), + exportBtn, + statusLabel, + ) +} + +func buildProxyTab(fyneApp fyne.App, bus *events.Bus) fyne.CanvasObject { + prefs := fyneApp.Preferences() + + hostEntry := widget.NewEntry() + hostEntry.SetText(prefs.StringWithFallback(prefKeyProxyHost, defaultProxyHost)) + hostEntry.SetPlaceHolder("127.0.0.1") + + portEntry := widget.NewEntry() + portEntry.SetText(strconv.Itoa(prefs.IntWithFallback(prefKeyProxyPort, defaultProxyPort))) + portEntry.SetPlaceHolder("9090") + + enabledCheck := widget.NewCheck("Proxy enabled", nil) + enabledCheck.Checked = prefs.BoolWithFallback(prefKeyProxyEnabled, defaultProxyEnabled) + + proxyStatus := widget.NewLabel("") + proxyStatus.Wrapping = fyne.TextWrapBreak + + saveBtn := widget.NewButton("Save & Restart Proxy", func() { + proxyHost := hostEntry.Text + proxyPort, err := strconv.Atoi(portEntry.Text) + if err != nil || proxyPort < 1 || proxyPort > 65535 { + proxyStatus.SetText("Invalid port number.") + return + } + prefs.SetString(prefKeyProxyHost, proxyHost) + prefs.SetInt(prefKeyProxyPort, proxyPort) + prefs.SetBool(prefKeyProxyEnabled, enabledCheck.Checked) + + if enabledCheck.Checked { + addr := fmt.Sprintf("%s:%d", proxyHost, proxyPort) + bus.EmitProxyRestart(events.ProxyRestartEvent{Addr: addr}) + proxyStatus.SetText(fmt.Sprintf("Proxy restarted on %s:%d", proxyHost, proxyPort)) + proxyRunningBinding.Set(true) + proxyStatusBinding.Set(fmt.Sprintf("Proxy: running on %s:%d", proxyHost, proxyPort)) + } else { + bus.EmitProxyStop(events.ProxyStopEvent{}) + proxyStatus.SetText("Proxy stopped.") + proxyRunningBinding.Set(false) + proxyStatusBinding.Set("Proxy: stopped") + } + }) + saveBtn.Importance = widget.HighImportance + + return container.NewVBox( + newBoldLabel("Proxy"), + widget.NewForm( + widget.NewFormItem("Host", hostEntry), + widget.NewFormItem("Port", portEntry), + ), + enabledCheck, + saveBtn, + proxyStatus, + ) +} + +// loadThemeByName loads a theme by name, checking embedded themes first then user themes on disk. +func loadThemeByName(name, themesDir string) (*LoadedTheme, error) { + for _, embedded := range ScanEmbeddedThemes() { + if embedded == name { + return LoadEmbeddedTheme(name) + } + } + return LoadTheme(filepath.Join(themesDir, name+".toml")) +} + +func buildAppearanceTab(fyneApp fyne.App, keybinds *Keybinds) fyne.CanvasObject { + prefs := fyneApp.Preferences() + themesDir := ThemesDir(fyneApp.Storage().RootURI().Path()) + + themeStatus := widget.NewLabel("") + themeStatus.Importance = widget.LowImportance + + themeSelect := widget.NewSelect(ScanAllThemes(themesDir), nil) + if saved := prefs.String(prefKeyUserTheme); saved != "" { + themeSelect.SetSelected(saved) + } else { + themeSelect.SetSelected(ThemeDefaultOption) + } + + variantSelect := widget.NewSelect([]string{"Dark", "Light"}, func(selected string) { + prefs.SetBool(prefKeyDarkTheme, selected == "Dark") + applyTheme(fyneApp) + }) + if prefs.BoolWithFallback(prefKeyDarkTheme, defaultDarkTheme) { + variantSelect.SetSelected("Dark") + } else { + variantSelect.SetSelected("Light") + } + if lt := loadActiveTheme(fyneApp); lt != nil && !lt.HasBoth { + variantSelect.Hide() + } + + applyBtn := widget.NewButton("Apply Theme", func() { + selected := themeSelect.Selected + if selected != ThemeDefaultOption { + if _, err := loadThemeByName(selected, themesDir); err != nil { + themeStatus.SetText("Invalid theme: " + err.Error()) + logger.Error("settings: load theme %s: %v", selected, err) + return + } + prefs.SetString(prefKeyUserTheme, selected) + } else { + prefs.SetString(prefKeyUserTheme, "") + } + applyTheme(fyneApp) + if lt := loadActiveTheme(fyneApp); lt != nil && !lt.HasBoth { + variantSelect.Hide() + } else { + variantSelect.Show() + } + themeStatus.SetText("Applied. Restart to update the dark/light toggle visibility.") + }) + + addThemeBtn := widget.NewButtonWithIcon("Add Theme", AppIcon("save"), func() { + d := dialog.NewFileOpen(func(rc fyne.URIReadCloser, err error) { + if err != nil || rc == nil { + return + } + defer rc.Close() + data, err := io.ReadAll(rc) + if err != nil { + themeStatus.SetText("Failed to read file: " + err.Error()) + return + } + dest := filepath.Join(themesDir, filepath.Base(rc.URI().Path())) + if err := os.WriteFile(dest, data, 0644); err != nil { + themeStatus.SetText("Failed to save theme: " + err.Error()) + return + } + themeStatus.SetText("Theme added: " + filepath.Base(dest)) + themeSelect.Options = ScanAllThemes(themesDir) + themeSelect.Refresh() + }, settingsWin) + d.SetFilter(storage.NewExtensionFileFilter([]string{".toml"})) + d.Show() + }) + + themeDirLabel := widget.NewLabel("User themes: " + themesDir) + themeDirLabel.Wrapping = fyne.TextWrapBreak + themeDirLabel.Importance = widget.LowImportance + + sendKeyEntry := widget.NewEntry() + sendKeyEntry.SetText(prefs.StringWithFallback(prefKeySendRequest, string(defaultKeySendRequest))) + + closeTabKeyEntry := widget.NewEntry() + closeTabKeyEntry.SetText(prefs.StringWithFallback(prefKeyCloseTab, string(defaultKeyCloseTab))) + + toRepeaterKeyEntry := widget.NewEntry() + toRepeaterKeyEntry.SetText(prefs.StringWithFallback(prefKeyToRepeater, string(defaultKeyToRepeater))) + + for _, e := range []*widget.Entry{sendKeyEntry, closeTabKeyEntry, toRepeaterKeyEntry} { + entry := e + entry.OnChanged = func(text string) { + if len(text) > 0 { + entry.SetText(strings.ToUpper(string([]rune(text)[0]))) + } + } + } + + shortcutStatus := widget.NewLabel("") + + saveShortcutsBtn := widget.NewButton("Save Shortcuts", func() { + keybinds.KeySendRequest = fyne.KeyName(strings.ToUpper(sendKeyEntry.Text)) + keybinds.KeyCloseTab = fyne.KeyName(strings.ToUpper(closeTabKeyEntry.Text)) + keybinds.KeyToRepeater = fyne.KeyName(strings.ToUpper(toRepeaterKeyEntry.Text)) + prefs.SetString(prefKeySendRequest, strings.ToUpper(sendKeyEntry.Text)) + prefs.SetString(prefKeyCloseTab, strings.ToUpper(closeTabKeyEntry.Text)) + prefs.SetString(prefKeyToRepeater, strings.ToUpper(toRepeaterKeyEntry.Text)) + keybinds.Update() + shortcutStatus.SetText("Shortcuts updated.") + }) + + return container.NewVBox( + newBoldLabel("Appearance"), + widget.NewForm( + widget.NewFormItem("Variant", variantSelect), + widget.NewFormItem("Theme", themeSelect), + ), + applyBtn, + addThemeBtn, + themeStatus, + themeDirLabel, + widget.NewSeparator(), + newBoldLabel("Shortcuts"), + widget.NewLabel("Single letter keys always, Ctrl modifier always applied (eg. Send is Ctrl + s)."), + widget.NewForm( + widget.NewFormItem("Send Request", sendKeyEntry), + widget.NewFormItem("Close Tab", closeTabKeyEntry), + widget.NewFormItem("Send to Repeater", toRepeaterKeyEntry), + ), + saveShortcutsBtn, + shortcutStatus, + ) +} + +func buildBrowserTab(fyneApp fyne.App, win fyne.Window) fyne.CanvasObject { + prefs := fyneApp.Preferences() + statusLabel := widget.NewLabel("") + statusLabel.Wrapping = fyne.TextWrapBreak + + browsers := DetectBrowsers() + names := make([]string, len(browsers)) + for i, b := range browsers { + names[i] = b.Name + } + + var browserSelect *widget.Select + if len(browsers) == 0 { + browserSelect = widget.NewSelect([]string{"No browsers detected"}, nil) + browserSelect.SetSelected("No browsers detected") + browserSelect.Disable() + } else { + browserSelect = widget.NewSelect(names, nil) + savedPath := prefs.String(prefKeyDefaultBrowser) + for _, b := range browsers { + if b.Path == savedPath { + browserSelect.SetSelected(b.Name) + break + } + } + if browserSelect.Selected == "" { + browserSelect.SetSelected(names[0]) + } + } + + rescanBtn := widget.NewButton("Rescan", func() { + fresh := DetectBrowsers() + freshNames := make([]string, len(fresh)) + for i, b := range fresh { + freshNames[i] = b.Name + } + browsers = fresh + names = freshNames + if len(freshNames) == 0 { + browserSelect.Options = []string{"No browsers detected"} + browserSelect.SetSelected("No browsers detected") + browserSelect.Disable() + statusLabel.SetText("No supported browsers found.") + } else { + browserSelect.Enable() + browserSelect.Options = freshNames + browserSelect.SetSelected(freshNames[0]) + statusLabel.SetText(fmt.Sprintf("Found %d browser(s).", len(freshNames))) + } + browserSelect.Refresh() + }) + + saveBtn := widget.NewButton("Save Default", func() { + if len(browsers) == 0 { + return + } + for _, b := range browsers { + if b.Name == browserSelect.Selected { + prefs.SetString(prefKeyDefaultBrowser, b.Path) + statusLabel.SetText("Default browser saved: " + b.Name) + return + } + } + }) + saveBtn.Importance = widget.HighImportance + + clearProfileBtn := widget.NewButton("Clear Browser Profile", func() { + for _, b := range browsers { + if b.Name == browserSelect.Selected { + if err := ClearBrowserProfile(fyneApp, b.Path); err != nil { + statusLabel.SetText("Error: " + err.Error()) + } else { + statusLabel.SetText("Profile cleared for: " + b.Name) + } + return + } + } + }) + + return container.NewVBox( + newBoldLabel("Browser"), + widget.NewLabel("Launch a browser pre-configured to route traffic through Shiv's proxy with an isolated profile."), + widget.NewForm( + widget.NewFormItem("Default browser", container.NewBorder(nil, nil, nil, rescanBtn, browserSelect)), + ), + container.NewHBox(saveBtn, clearProfileBtn), + statusLabel, + ) +} + +func newBrowserPickerContent(names []string, onChanged func(string)) fyne.CanvasObject { + sel := widget.NewSelect(names, onChanged) + if len(names) > 0 { + sel.SetSelected(names[0]) + } + return container.NewVBox( + widget.NewLabel("Multiple browsers detected. Choose one to set as default:"), + sel, + ) +} diff --git a/internal/ui/theme.go b/internal/ui/theme.go new file mode 100644 index 0000000..591a102 --- /dev/null +++ b/internal/ui/theme.go @@ -0,0 +1,373 @@ +package ui + +import ( + "fmt" + "image/color" + "io/fs" + "os" + "path/filepath" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/theme" + "github.com/BurntSushi/toml" + "github.com/shiv/assets" +) + +// ── Theme file format ───────────────────────────────────────────────────────── + +type themePalette struct { + Background string `toml:"background"` + Button string `toml:"button"` + DisabledButton string `toml:"disabled_button"` + Disabled string `toml:"disabled"` + Error string `toml:"error"` + Focus string `toml:"focus"` + Foreground string `toml:"foreground"` + HeaderBackground string `toml:"header_background"` + Hover string `toml:"hover"` + Hyperlink string `toml:"hyperlink"` + InputBackground string `toml:"input_background"` + InputBorder string `toml:"input_border"` + MenuBackground string `toml:"menu_background"` + OverlayBackground string `toml:"overlay_background"` + PlaceHolder string `toml:"placeholder"` + Pressed string `toml:"pressed"` + Primary string `toml:"primary"` + ScrollBar string `toml:"scrollbar"` + Selection string `toml:"selection"` + Separator string `toml:"separator"` + Shadow string `toml:"shadow"` + Success string `toml:"success"` + Warning string `toml:"warning"` +} + +type themeFile struct { + Dark *themePalette `toml:"dark"` + Light *themePalette `toml:"light"` +} + +// LoadedTheme is a parsed, validated theme ready for use. +type LoadedTheme struct { + Name string + Dark map[fyne.ThemeColorName]color.Color + Light map[fyne.ThemeColorName]color.Color + HasBoth bool + Embedded bool +} + +// ── Color name mapping ──────────────────────────────────────────────────────── + +var paletteKeys = map[fyne.ThemeColorName]string{ + theme.ColorNameBackground: "background", + theme.ColorNameButton: "button", + theme.ColorNameDisabledButton: "disabled_button", + theme.ColorNameDisabled: "disabled", + theme.ColorNameError: "error", + theme.ColorNameFocus: "focus", + theme.ColorNameForeground: "foreground", + theme.ColorNameHeaderBackground: "header_background", + theme.ColorNameHover: "hover", + theme.ColorNameHyperlink: "hyperlink", + theme.ColorNameInputBackground: "input_background", + theme.ColorNameInputBorder: "input_border", + theme.ColorNameMenuBackground: "menu_background", + theme.ColorNameOverlayBackground: "overlay_background", + theme.ColorNamePlaceHolder: "placeholder", + theme.ColorNamePressed: "pressed", + theme.ColorNamePrimary: "primary", + theme.ColorNameScrollBar: "scrollbar", + theme.ColorNameSelection: "selection", + theme.ColorNameSeparator: "separator", + theme.ColorNameShadow: "shadow", + theme.ColorNameSuccess: "success", + theme.ColorNameWarning: "warning", +} + +var paletteKeysReverse map[string]fyne.ThemeColorName + +func init() { + paletteKeysReverse = make(map[string]fyne.ThemeColorName, len(paletteKeys)) + for k, v := range paletteKeys { + paletteKeysReverse[v] = k + } +} + +// ── Theme scanning ──────────────────────────────────────────────────────────── + +const ThemeDefaultOption = "Default" + +// ThemesDir returns the user themes directory, creating it if needed. +func ThemesDir(storageRoot string) string { + dir := filepath.Join(storageRoot, "themes") + os.MkdirAll(dir, 0755) + return dir +} + +// ScanEmbeddedThemes returns names of all embedded themes (without extension). +func ScanEmbeddedThemes() []string { + entries, err := fs.ReadDir(assets.Themes, "themes") + if err != nil { + return nil + } + var names []string + for _, e := range entries { + if !e.IsDir() && strings.HasSuffix(e.Name(), ".toml") { + names = append(names, strings.TrimSuffix(e.Name(), ".toml")) + } + } + return names +} + +// ScanUserThemes returns names of all user .toml themes in dir (without extension). +func ScanUserThemes(dir string) []string { + entries, err := os.ReadDir(dir) + if err != nil { + return nil + } + var names []string + for _, e := range entries { + if !e.IsDir() && strings.HasSuffix(e.Name(), ".toml") { + names = append(names, strings.TrimSuffix(e.Name(), ".toml")) + } + } + return names +} + +// ScanAllThemes returns the full theme list for the dropdown: +// [Default, ...embedded..., ...user (skipping duplicates of embedded)...] +func ScanAllThemes(userThemesDir string) []string { + embedded := ScanEmbeddedThemes() + embeddedSet := make(map[string]bool, len(embedded)) + for _, n := range embedded { + embeddedSet[n] = true + } + options := make([]string, 0, 1+len(embedded)+8) + options = append(options, ThemeDefaultOption) + options = append(options, embedded...) + for _, n := range ScanUserThemes(userThemesDir) { + if !embeddedSet[n] { + options = append(options, n) + } + } + return options +} + +// ── Theme loading ───────────────────────────────────────────────────────────── + +// LoadEmbeddedTheme loads a theme by name from the embedded assets. +func LoadEmbeddedTheme(name string) (*LoadedTheme, error) { + data, err := assets.Themes.ReadFile("themes/" + name + ".toml") + if err != nil { + return nil, fmt.Errorf("embedded theme %q not found: %w", name, err) + } + lt, err := parseThemeBytes(data, name) + if err != nil { + return nil, err + } + lt.Embedded = true + return lt, nil +} + +// LoadTheme loads a theme from a file path on disk. +func LoadTheme(path string) (*LoadedTheme, error) { + name := strings.TrimSuffix(filepath.Base(path), ".toml") + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read theme: %w", err) + } + return parseThemeBytes(data, name) +} + +func parseThemeBytes(data []byte, name string) (*LoadedTheme, error) { + var tf themeFile + if err := toml.Unmarshal(data, &tf); err != nil { + return nil, fmt.Errorf("parse theme: %w", err) + } + if tf.Dark == nil && tf.Light == nil { + return nil, fmt.Errorf("theme must have at least a [dark] or [light] section") + } + lt := &LoadedTheme{ + Name: name, + HasBoth: tf.Dark != nil && tf.Light != nil, + } + if tf.Dark != nil { + m, err := parsePalette(tf.Dark) + if err != nil { + return nil, fmt.Errorf("[dark]: %w", err) + } + lt.Dark = m + } + if tf.Light != nil { + m, err := parsePalette(tf.Light) + if err != nil { + return nil, fmt.Errorf("[light]: %w", err) + } + lt.Light = m + } + if lt.Dark == nil { + lt.Dark = lt.Light + } + if lt.Light == nil { + lt.Light = lt.Dark + } + return lt, nil +} + +func parsePalette(p *themePalette) (map[fyne.ThemeColorName]color.Color, error) { + raw := map[string]string{ + "background": p.Background, + "button": p.Button, + "disabled_button": p.DisabledButton, + "disabled": p.Disabled, + "error": p.Error, + "focus": p.Focus, + "foreground": p.Foreground, + "header_background": p.HeaderBackground, + "hover": p.Hover, + "hyperlink": p.Hyperlink, + "input_background": p.InputBackground, + "input_border": p.InputBorder, + "menu_background": p.MenuBackground, + "overlay_background": p.OverlayBackground, + "placeholder": p.PlaceHolder, + "pressed": p.Pressed, + "primary": p.Primary, + "scrollbar": p.ScrollBar, + "selection": p.Selection, + "separator": p.Separator, + "shadow": p.Shadow, + "success": p.Success, + "warning": p.Warning, + } + m := make(map[fyne.ThemeColorName]color.Color, len(raw)) + for key, val := range raw { + if val == "" { + continue + } + c, err := parseHex(val) + if err != nil { + return nil, fmt.Errorf("%s: %w", key, err) + } + if colorName, ok := paletteKeysReverse[key]; ok { + m[colorName] = c + } + } + return m, nil +} + +func parseHex(h string) (color.NRGBA, error) { + h = strings.TrimPrefix(h, "#") + if len(h) != 6 { + return color.NRGBA{}, fmt.Errorf("invalid hex %q (expected #rrggbb)", h) + } + var r, g, b uint8 + if _, err := fmt.Sscanf(h, "%02x%02x%02x", &r, &g, &b); err != nil { + return color.NRGBA{}, fmt.Errorf("invalid hex %q: %w", h, err) + } + return color.NRGBA{R: r, G: g, B: b, A: 0xff}, nil +} + +// ── Fyne theme implementation ───────────────────────────────────────────────── + +type vagueTheme struct { + dark bool + loaded *LoadedTheme +} + +func NewVagueThemeWithLoaded(dark bool, lt *LoadedTheme) fyne.Theme { + return &vagueTheme{dark: dark, loaded: lt} +} + +func (t *vagueTheme) Color(name fyne.ThemeColorName, variant fyne.ThemeVariant) color.Color { + if t.loaded != nil { + var palette map[fyne.ThemeColorName]color.Color + if t.dark { + palette = t.loaded.Dark + } else { + palette = t.loaded.Light + } + if c, ok := palette[name]; ok { + return c + } + } + if t.dark { + return theme.DefaultTheme().Color(name, theme.VariantDark) + } + return theme.DefaultTheme().Color(name, theme.VariantLight) +} + +func (t *vagueTheme) Font(style fyne.TextStyle) fyne.Resource { + switch { + case style.Bold && style.Italic: + return fyne.NewStaticResource("JetBrainsMono-BoldItalic.ttf", assets.FontJetBrainsMonoBold) + case style.Bold: + return fyne.NewStaticResource("JetBrainsMono-Bold.ttf", assets.FontJetBrainsMonoBold) + case style.Italic: + return fyne.NewStaticResource("JetBrainsMono-Italic.ttf", assets.FontJetBrainsMonoItalic) + default: + return fyne.NewStaticResource("JetBrainsMono-Regular.ttf", assets.FontJetBrainsMonoRegular) + } +} + +func (t *vagueTheme) Icon(name fyne.ThemeIconName) fyne.Resource { + return theme.DefaultTheme().Icon(name) +} + +func (t *vagueTheme) Size(name fyne.ThemeSizeName) float32 { + if name == theme.SizeNameInlineIcon { + return 28 + } + if name == theme.SizeNameText { + return 13 + } + return theme.DefaultTheme().Size(name) +} + +// ── Icons ───────────────────────────────────────────────────────────────────── + +func AppIcon(name string) fyne.Resource { + switch name { + case "delete": + return fyne.NewStaticResource("delete.png", assets.IconDelete) + case "save": + return fyne.NewStaticResource("save.png", assets.IconSave) + case "history": + return fyne.NewStaticResource("history.png", assets.IconHistory) + case "repeater": + return fyne.NewStaticResource("boomerang.png", assets.IconRepeater) + case "inspector": + return fyne.NewStaticResource("inspector.png", assets.IconInspector) + case "intruder": + return fyne.NewStaticResource("intruder.png", assets.IconIntruder) + case "intercept": + return fyne.NewStaticResource("intercept.png", assets.IconIntercept) + case "loot": + return fyne.NewStaticResource("loot.png", assets.IconLoot) + case "project": + return fyne.NewStaticResource("project.png", assets.IconProject) + case "scope": + return fyne.NewStaticResource("scope.png", assets.IconScope) + case "off-button": + return fyne.NewStaticResource("off-button.png", assets.IconOffButton) + case "on-button": + return fyne.NewStaticResource("on-button.png", assets.IconOnButton) + case "settings": + return fyne.NewStaticResource("settings.png", assets.IconSettings) + case "toolbox": + return fyne.NewStaticResource("toolbox.png", assets.IconToolbox) + case "plugin": + return fyne.NewStaticResource("plugin.png", assets.IconPlugin) + case "web": + return fyne.NewStaticResource("web.png", assets.IconWeb) + case "web1": + return fyne.NewStaticResource("web1.png", assets.IconWeb1) + case "document": + return fyne.NewStaticResource("document.png", assets.IconDocument) + case "folder": + return fyne.NewStaticResource("folder.png", assets.IconFolder) + case "media": + return fyne.NewStaticResource("media.png", assets.IconMedia) + } + return theme.DefaultTheme().Icon(theme.IconNameQuestion) +} diff --git a/internal/ui/websocket.go b/internal/ui/websocket.go new file mode 100644 index 0000000..9251cac --- /dev/null +++ b/internal/ui/websocket.go @@ -0,0 +1,379 @@ +package ui + +import ( + "fmt" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/widget" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" + "github.com/shiv/internal/ui/widgets" +) + +var wsWin fyne.Window + +func showWebSocketWindow(fyneApp fyne.App, projectStore *store.Store, parentWin fyne.Window, repeater *repeaterTab) { + if wsWin != nil { + wsWin.RequestFocus() + return + } + wsWin = fyneApp.NewWindow("WebSockets") + wsWin.Resize(fyne.NewSize(1100, 650)) + wsWin.SetOnClosed(func() { wsWin = nil }) + closeOnEscape(wsWin, wsWin.Close) + wsWin.SetContent(buildWebSocketContent(projectStore, wsWin, repeater)) + wsWin.Show() +} + +func filterConnections(all []store.WebSocketConnection, query string) []store.WebSocketConnection { + q := strings.ToLower(strings.TrimSpace(query)) + if q == "" { + out := make([]store.WebSocketConnection, len(all)) + copy(out, all) + return out + } + terms := strings.Fields(q) + var out []store.WebSocketConnection + for _, c := range all { + s := strings.ToLower(c.Host + " " + c.URL) + ok := true + for _, t := range terms { + if !strings.Contains(s, t) { + ok = false + break + } + } + if ok { + out = append(out, c) + } + } + return out +} + +func filterFrames(all []store.WebSocketFrame, query string) []store.WebSocketFrame { + q := strings.ToLower(strings.TrimSpace(query)) + if q == "" { + out := make([]store.WebSocketFrame, len(all)) + copy(out, all) + return out + } + terms := strings.Fields(q) + var out []store.WebSocketFrame + for _, f := range all { + s := strings.ToLower(string(f.Payload) + " " + opcodeLabel(f.Opcode)) + ok := true + for _, t := range terms { + if !strings.Contains(s, t) { + ok = false + break + } + } + if ok { + out = append(out, f) + } + } + return out +} + +func buildConnTable(win fyne.Window, filteredConns *[]store.WebSocketConnection) (*widgets.DataTable, fyne.CanvasObject) { + connColumns := []widgets.DataTableColumn{ + {Header: "ID", Width: 50}, + {Header: "Host", Width: 220}, + {Header: "URL", Width: 380}, + {Header: "TLS", Width: 45}, + {Header: "Frames", Width: 70}, + {Header: "Time", Width: 180}, + } + connTable := widgets.NewDataTable() + connTable.SetWindow(win) + connTable.Columns = connColumns + connTable.RowCount = func() int { return len(*filteredConns) } + connTable.RowID = func(row int) int64 { + if row >= len(*filteredConns) { + return 0 + } + return int64((*filteredConns)[row].ID) + } + connTable.CellValue = func(row, col int) string { + if row >= len(*filteredConns) { + return "" + } + c := (*filteredConns)[row] + switch col { + case 0: + return fmt.Sprintf("%d", c.ID) + case 1: + return c.Host + case 2: + return c.URL + case 3: + if c.TLS { + return "✓" + } + return "" + case 4: + return fmt.Sprintf("%d", c.FrameCount) + case 5: + return c.Timestamp.Local().Format("2006-01-02 15:04:05") + } + return "" + } + return connTable, connTable.Build() +} + +func buildFrameTable(win fyne.Window, filteredFrames *[]store.WebSocketFrame) (*widgets.DataTable, fyne.CanvasObject) { + frameColumns := []widgets.DataTableColumn{ + {Header: "Dir", Width: 90}, + {Header: "Type", Width: 70}, + {Header: "Length", Width: 80}, + {Header: "Time", Width: 150}, + } + frameTable := widgets.NewDataTable() + frameTable.SetWindow(win) + frameTable.Columns = frameColumns + frameTable.RowCount = func() int { return len(*filteredFrames) } + frameTable.RowID = func(row int) int64 { + if row >= len(*filteredFrames) { + return 0 + } + return int64((*filteredFrames)[row].ID) + } + frameTable.CellValue = func(row, col int) string { + if row >= len(*filteredFrames) { + return "" + } + f := (*filteredFrames)[row] + switch col { + case 0: + if f.Direction == events.WebSocketClient { + return "→ Client" + } + return "← Server" + case 1: + return opcodeLabel(f.Opcode) + case 2: + return fmt.Sprintf("%db", len(f.Payload)) + case 3: + return f.Timestamp.Local().Format("15:04:05.000") + } + return "" + } + frameTable.CellStyle = func(row, col int) widget.Importance { + if col != 0 || row >= len(*filteredFrames) { + return widget.MediumImportance + } + if (*filteredFrames)[row].Direction == events.WebSocketClient { + return widget.HighImportance + } + return widget.SuccessImportance + } + return frameTable, frameTable.Build() +} + +func buildFrameDetail(win fyne.Window, filteredFrames *[]store.WebSocketFrame, frameTable *widgets.DataTable) (*widgets.TextView, fyne.CanvasObject) { + payloadView := widgets.NewTextView() + payloadView.SetWindow(win) + payloadView.SetPlaceHolder("Select a frame to view payload...") + + frameTable.OnSelect = func(row int) { + if row >= len(*filteredFrames) { + return + } + payload := (*filteredFrames)[row].Payload + if len(payload) == 0 { + payloadView.SetText("(empty)") + } else { + payloadView.SetText(string(payload)) + } + } + + payloadPane := container.NewBorder( + newBoldLabel("Payload"), + nil, nil, nil, + payloadView.Build(), + ) + return payloadView, payloadPane +} + +func buildWebSocketContent(projectStore *store.Store, win fyne.Window, repeater *repeaterTab) fyne.CanvasObject { + var allConns []store.WebSocketConnection + var filteredConns []store.WebSocketConnection + var allFrames []store.WebSocketFrame + var filteredFrames []store.WebSocketFrame + var selectedConnID uint64 + + connFilterEntry := widget.NewEntry() + connFilterEntry.SetPlaceHolder("Filter — host, url...") + frameFilterEntry := widget.NewEntry() + frameFilterEntry.SetPlaceHolder("Filter frames — payload, type...") + + refilterConns := func() { filteredConns = filterConnections(allConns, connFilterEntry.Text) } + refilterFrames := func() { filteredFrames = filterFrames(allFrames, frameFilterEntry.Text) } + + connTable, connTableObj := buildConnTable(win, &filteredConns) + frameTable, frameTableObj := buildFrameTable(win, &filteredFrames) + payloadView, payloadPane := buildFrameDetail(win, &filteredFrames, frameTable) + + loadFrames := func(connID uint64) { + go func() { + loaded, err := projectStore.FramesForConnection(connID) + if err != nil { + logger.Error("ws ui: load frames %d: %v", connID, err) + return + } + fyne.Do(func() { + allFrames = loaded + refilterFrames() + frameTable.Refresh() + payloadView.SetText("") + }) + }() + } + + connTable.OnSelect = func(row int) { + if row >= len(filteredConns) { + return + } + selectedConnID = filteredConns[row].ID + loadFrames(selectedConnID) + } + + connTable.MenuItems = func(row int) []widgets.ContextMenuItem { + if row >= len(filteredConns) { + return nil + } + c := filteredConns[row] + return []widgets.ContextMenuItem{ + {Label: "Send to Repeater", Action: func() { + if repeater != nil { + repeater.AddWSTab(c.Host, c.URL, c.TLS) + } + }}, + {Label: "Copy URL", Action: func() { + fyne.CurrentApp().Clipboard().SetContent(c.URL) + }}, + } + } + + connFilterEntry.OnChanged = func(q string) { + filteredConns = filterConnections(allConns, q) + connTable.Refresh() + } + frameFilterEntry.OnChanged = func(q string) { + filteredFrames = filterFrames(allFrames, q) + frameTable.Refresh() + } + + go func() { + for conn := range projectStore.WebSocketConnections { + c := conn + fyne.Do(func() { + for _, existing := range allConns { + if existing.ID == c.ID { + return + } + } + allConns = append([]store.WebSocketConnection{c}, allConns...) + refilterConns() + connTable.Refresh() + }) + } + }() + + go func() { + for frame := range projectStore.WebSocketFrames { + f := frame + fyne.Do(func() { + for i := range allConns { + if allConns[i].ID == f.ConnectionID { + allConns[i].FrameCount++ + break + } + } + refilterConns() + connTable.Refresh() + if f.ConnectionID != selectedConnID { + return + } + go func() { + loaded, err := projectStore.FramesForConnection(f.ConnectionID) + if err != nil { + logger.Error("ws ui: reload frames %d: %v", f.ConnectionID, err) + return + } + fyne.Do(func() { + allFrames = loaded + refilterFrames() + frameTable.Refresh() + }) + }() + }) + } + }() + + refreshBtn := widget.NewButtonWithIcon("Refresh", AppIcon("history"), func() { + go func() { + loaded, err := projectStore.AllWebSocketConnections() + if err != nil { + logger.Error("ws ui: load connections: %v", err) + return + } + fyne.Do(func() { + allConns = loaded + refilterConns() + connTable.Refresh() + if selectedConnID > 0 { + loadFrames(selectedConnID) + } + }) + }() + }) + + go func() { + loaded, err := projectStore.AllWebSocketConnections() + if err != nil { + logger.Error("ws ui: initial load: %v", err) + return + } + fyne.Do(func() { + allConns = loaded + filteredConns = filterConnections(allConns, "") + connTable.Refresh() + }) + }() + + connFilterBar := container.NewBorder(nil, nil, nil, refreshBtn, connFilterEntry) + frameFilterBar := container.NewBorder(nil, nil, newBoldLabel("Frames"), nil, frameFilterEntry) + + connPane := container.NewBorder( + container.NewVBox(newBoldLabel("Connections"), connFilterBar), + nil, nil, nil, connTableObj, + ) + framePane := container.NewBorder(frameFilterBar, nil, nil, nil, frameTableObj) + + bottomSplit := container.NewHSplit(framePane, payloadPane) + bottomSplit.SetOffset(0.4) + mainSplit := container.NewVSplit(connPane, bottomSplit) + mainSplit.SetOffset(0.4) + return mainSplit +} + +func opcodeLabel(op events.WebSocketOpcode) string { + switch op { + case events.WebSocketText: + return "Text" + case events.WebSocketBinary: + return "Binary" + case events.WebSocketPing: + return "Ping" + case events.WebSocketPong: + return "Pong" + case events.WebSocketClose: + return "Close" + } + return fmt.Sprintf("0x%x", int(op)) +} diff --git a/internal/ui/widgets/datatable.go b/internal/ui/widgets/datatable.go new file mode 100644 index 0000000..2cf476d --- /dev/null +++ b/internal/ui/widgets/datatable.go @@ -0,0 +1,876 @@ +package widgets + +import ( + "image/color" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/canvas" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/driver/desktop" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" +) + +// DataTableColumn describes a single column header and its initial width. +type DataTableColumn struct { + Header string + Width float32 +} + +// ContextMenuItem describes a single item in a right-click context menu. +type ContextMenuItem struct { + Label string + Action func() +} + +// Layout constants for DataTable. All sizes are in logical pixels. +const ( + dtRowH float32 = 32 // height of a data row + dtHdrH float32 = 34 // height of the header row + dtPadX float32 = 8 // horizontal cell padding + dtDivW float32 = 6 // width of a column-divider hit area + dtScrollW float32 = 8 // width of the scrollbar track + + // dtNearBottomThreshold is how many rows from the end triggers OnNearBottom. + dtNearBottomThreshold = 20 +) + +// DataTable is a virtualised, row-selectable table widget with resizable +// columns, right-click context menus, and a draggable scrollbar thumb. +// +// Create with NewDataTable, set all callback fields, call SetWindow, then +// use the object returned by Build in your layout. +type DataTable struct { + widget.BaseWidget + + // Columns defines the header labels and initial widths. + Columns []DataTableColumn + + // Data callbacks — may be updated at runtime; call Refresh after changes. + RowCount func() int + CellValue func(row, col int) string + CellStyle func(row, col int) widget.Importance // nil → MediumImportance + RowID func(row int) int64 // stable unique ID per row + + // Interaction callbacks. + OnSelect func(row int) + MenuItems func(row int) []ContextMenuItem // nil → no context menu + OnNearBottom func() // fired when scrolled near the last row + RowBgColor func(row int) color.Color // nil → default background; non-nil overrides row background + + colWidths []float32 + selectedID int64 + selectedRow int // last-selected row index; used for keyboard navigation + hasSelected bool + scrollOffset int + scrollFrac float32 // sub-row accumulator for smooth trackpad scrolling + + nearBottomFired bool // debounce: only fire once per approach to the bottom + + scrollbarHovered bool + + win fyne.Window + rend *dtRenderer + inputLayer *dtInputLayer +} + +// NewDataTable allocates and extends a DataTable. +func NewDataTable() *DataTable { + t := &DataTable{} + t.ExtendBaseWidget(t) + return t +} + +// SetWindow provides the window reference required for context menu anchoring. +// Must be called before the widget is shown. +func (t *DataTable) SetWindow(w fyne.Window) { t.win = w } + +// Build constructs the complete canvas object to place in a layout. +// Call once after all callbacks are configured. +func (t *DataTable) Build() fyne.CanvasObject { + t.initColWidths() + layer := newDTInputLayer(t) + t.inputLayer = layer + return container.NewStack(t, layer) +} + +// SetSelected highlights the row identified by id. Pass 0 to clear. +func (t *DataTable) SetSelected(id int64) { + t.selectedID = id + t.hasSelected = id != 0 + t.Refresh() +} + +// ClearSelection removes the current row highlight. +func (t *DataTable) ClearSelection() { + t.selectedID = 0 + t.hasSelected = false + t.Refresh() +} + +// ScrollToRow scrolls the viewport so that row is visible. +func (t *DataTable) ScrollToRow(row int) { + visible := t.visibleSlots() + if row < t.scrollOffset { + t.scrollOffset = row + } else if row >= t.scrollOffset+visible { + t.scrollOffset = row - visible + 1 + } + if t.scrollOffset < 0 { + t.scrollOffset = 0 + } + t.Refresh() +} + +// Scrolled implements fyne.Scrollable, forwarding wheel events from the +// input layer overlay. +func (t *DataTable) Scrolled(ev *fyne.ScrollEvent) { + totalRows := t.rowCount() + visible := t.visibleSlots() + maxOffset := totalRows - visible + if maxOffset < 0 { + maxOffset = 0 + } + t.scrollFrac -= ev.Scrolled.DY / dtRowH + rows := int(t.scrollFrac) + t.scrollFrac -= float32(rows) + if rows == 0 { + if ev.Scrolled.DY < 0 { + rows = 1 + } else if ev.Scrolled.DY > 0 { + rows = -1 + } + } + t.scrollOffset += rows + if t.scrollOffset < 0 { + t.scrollOffset = 0 + } + if t.scrollOffset > maxOffset { + t.scrollOffset = maxOffset + } + t.Refresh() + + // Fire OnNearBottom when the user scrolls within dtNearBottomThreshold rows + // of the end. Debounced: only fires once until the user scrolls back up. + if t.OnNearBottom != nil && totalRows > 0 { + rowsFromBottom := totalRows - (t.scrollOffset + visible) + if rowsFromBottom < dtNearBottomThreshold { + if !t.nearBottomFired { + t.nearBottomFired = true + t.OnNearBottom() + } + } else { + // Reset debounce when user scrolls back up far enough. + t.nearBottomFired = false + } + } +} + +// CreateRenderer implements fyne.Widget. +func (t *DataTable) CreateRenderer() fyne.WidgetRenderer { + t.initColWidths() + renderer := newDTRenderer(t) + t.rend = renderer + return renderer +} + +func (t *DataTable) initColWidths() { + if len(t.colWidths) == len(t.Columns) { + return + } + t.colWidths = make([]float32, len(t.Columns)) + for i, col := range t.Columns { + t.colWidths[i] = col.Width + } +} + +func (t *DataTable) rowCount() int { + if t.RowCount == nil { + return 0 + } + return t.RowCount() +} + +func (t *DataTable) visibleSlots() int { + if t.rend == nil { + return 1 + } + availHeight := t.rend.size.Height - dtHdrH + if availHeight < 0 { + return 0 + } + return int(availHeight/dtRowH) + 2 +} + +// cellColor maps a widget.Importance value to a theme colour. +func (t *DataTable) cellColor(row, col int) color.Color { + importance := widget.MediumImportance + if t.CellStyle != nil { + importance = t.CellStyle(row, col) + } + switch importance { + case widget.DangerImportance: + return theme.Color(theme.ColorNameError) + case widget.WarningImportance: + return theme.Color(theme.ColorNameWarning) + case widget.SuccessImportance: + return theme.Color(theme.ColorNameSuccess) + case widget.LowImportance: + return theme.Color(theme.ColorNamePlaceHolder) + case widget.HighImportance: + return theme.Color(theme.ColorNamePrimary) + default: + return theme.Color(theme.ColorNameForeground) + } +} + +// minColWidth returns the minimum display width for col, measured from its +// header text so the header is always fully readable. +func (t *DataTable) minColWidth(col int) float32 { + if col >= len(t.Columns) { + return dtPadX * 2 + } + textWidth := fyne.MeasureText(t.Columns[col].Header, theme.TextSize(), fyne.TextStyle{Bold: true}).Width + return textWidth + dtPadX*2 +} + +// thumbBounds returns the scrollbar thumb position and height in widget-local +// coordinates. ok is false when all rows fit in the viewport. +func (t *DataTable) thumbBounds() (thumbY, thumbH float32, ok bool) { + if t.rend == nil { + return 0, 0, false + } + totalRows := t.rowCount() + bodyHeight := t.rend.size.Height - dtHdrH + visibleRows := int(bodyHeight / dtRowH) + scrollable := totalRows - visibleRows + if scrollable <= 0 || bodyHeight <= 0 { + return 0, 0, false + } + ratio := float32(visibleRows) / float32(totalRows) + thumbH = bodyHeight * ratio + if thumbH < 20 { + thumbH = 20 + } + trackHeight := bodyHeight - thumbH + thumbY = dtHdrH + (float32(t.scrollOffset)/float32(scrollable))*trackHeight + return thumbY, thumbH, true +} + +func (t *DataTable) tappedRow(row int) { + if t.RowID != nil { + t.selectedID = t.RowID(row) + t.hasSelected = true + } + t.selectedRow = row + t.Refresh() + if t.OnSelect != nil { + t.OnSelect(row) + } + if t.win != nil { + t.win.Canvas().Focus(t) + } +} + +// FocusGained, FocusLost, TypedRune, TypedKey implement fyne.Focusable. +func (t *DataTable) FocusGained() { t.Refresh() } +func (t *DataTable) FocusLost() { t.Refresh() } +func (t *DataTable) TypedRune(_ rune) {} +func (t *DataTable) TypedKey(ev *fyne.KeyEvent) { + if !t.hasSelected || t.RowCount == nil { + return + } + total := t.RowCount() + if total == 0 { + return + } + next := t.selectedRow + switch ev.Name { + case fyne.KeyUp: + next-- + case fyne.KeyDown: + next++ + default: + return + } + if next < 0 { + next = 0 + } + if next >= total { + next = total - 1 + } + if next == t.selectedRow { + return + } + t.ScrollToRow(next) + t.tappedRow(next) +} + +func (t *DataTable) secondaryTappedRow(row int, absPos fyne.Position) { + if t.MenuItems == nil || t.win == nil { + return + } + items := t.MenuItems(row) + if len(items) == 0 { + return + } + menuItems := make([]*fyne.MenuItem, len(items)) + for i, item := range items { + it := item + menuItems[i] = fyne.NewMenuItem(it.Label, it.Action) + } + widget.ShowPopUpMenuAtPosition(fyne.NewMenu("", menuItems...), t.win.Canvas(), absPos) +} + +// dtRenderer draws the table using canvas primitives. Only visible row slots +// are allocated; slots grow as the widget height increases but are never freed. +type dtRenderer struct { + table *DataTable + size fyne.Size + + headerBg *canvas.Rectangle + headerTexts []*canvas.Text + dividerLines []*canvas.Line + + slots int + rowBgs []*canvas.Rectangle + cellTexts [][]*canvas.Text + + scrollTrack *canvas.Rectangle + scrollThumb *canvas.Rectangle + + objs []fyne.CanvasObject // cached; rebuilt only when slots grow +} + +func newDTRenderer(table *DataTable) *dtRenderer { + nCols := len(table.Columns) + r := &dtRenderer{table: table} + r.headerBg = canvas.NewRectangle(color.Transparent) + + r.headerTexts = make([]*canvas.Text, nCols) + for i, col := range table.Columns { + headerText := canvas.NewText(col.Header, color.Transparent) + headerText.TextStyle = fyne.TextStyle{Bold: true} + headerText.TextSize = theme.TextSize() + r.headerTexts[i] = headerText + } + + if nCols > 1 { + r.dividerLines = make([]*canvas.Line, nCols-1) + for i := range r.dividerLines { + dividerLine := canvas.NewLine(color.Transparent) + dividerLine.StrokeWidth = 1 + r.dividerLines[i] = dividerLine + } + } + + r.scrollTrack = canvas.NewRectangle(color.Transparent) + r.scrollThumb = canvas.NewRectangle(color.Transparent) + r.rebuildObjects() + return r +} + +func (r *dtRenderer) Destroy() {} + +func (r *dtRenderer) Objects() []fyne.CanvasObject { return r.objs } + +func (r *dtRenderer) MinSize() fyne.Size { + return fyne.NewSize(100, dtHdrH+dtRowH*3) +} + +func (r *dtRenderer) Layout(size fyne.Size) { + r.size = size + r.growSlots() + r.layoutHeader() + r.layoutBody() + r.layoutScrollbar() +} + +func (r *dtRenderer) Refresh() { + r.table.initColWidths() + r.growSlots() + r.layoutHeader() + r.layoutBody() + r.updateCells() + r.layoutScrollbar() + canvas.Refresh(r.table) +} + +// rebuildObjects reconstructs the cached object slice after slots grow. +func (r *dtRenderer) rebuildObjects() { + nFixed := 1 + len(r.headerTexts) + len(r.dividerLines) + 2 + total := nFixed + len(r.rowBgs) + for _, row := range r.cellTexts { + total += len(row) + } + objs := make([]fyne.CanvasObject, 0, total) + objs = append(objs, r.headerBg) + for _, headerText := range r.headerTexts { + objs = append(objs, headerText) + } + for _, dividerLine := range r.dividerLines { + objs = append(objs, dividerLine) + } + for _, rowBg := range r.rowBgs { + objs = append(objs, rowBg) + } + for _, row := range r.cellTexts { + for _, cellText := range row { + objs = append(objs, cellText) + } + } + objs = append(objs, r.scrollTrack, r.scrollThumb) + r.objs = objs +} + +func (r *dtRenderer) growSlots() { + bodyHeight := r.size.Height - dtHdrH + if bodyHeight < 0 { + bodyHeight = 0 + } + needed := int(bodyHeight/dtRowH) + 2 + nCols := len(r.table.Columns) + + grew := false + for r.slots < needed { + rowBg := canvas.NewRectangle(color.Transparent) + r.rowBgs = append(r.rowBgs, rowBg) + cellRow := make([]*canvas.Text, nCols) + for col := 0; col < nCols; col++ { + cellText := canvas.NewText("", color.Transparent) + cellText.TextSize = theme.TextSize() + cellRow[col] = cellText + } + r.cellTexts = append(r.cellTexts, cellRow) + r.slots++ + grew = true + } + if grew { + r.rebuildObjects() + } + + for slot := needed; slot < r.slots; slot++ { + r.rowBgs[slot].Hide() + for _, cellText := range r.cellTexts[slot] { + cellText.Hide() + } + } +} + +func (r *dtRenderer) layoutHeader() { + r.headerBg.Move(fyne.NewPos(0, 0)) + r.headerBg.Resize(fyne.NewSize(r.size.Width, dtHdrH)) + r.headerBg.FillColor = theme.Color(theme.ColorNameHeaderBackground) + r.headerBg.Refresh() + + xPos := float32(0) + for i, headerText := range r.headerTexts { + if i >= len(r.table.colWidths) { + headerText.Hide() + continue + } + colWidth := r.table.colWidths[i] + headerText.Move(fyne.NewPos(xPos+dtPadX, (dtHdrH-theme.TextSize())/2)) + headerText.Color = theme.Color(theme.ColorNameForeground) + headerText.Text = r.table.Columns[i].Header + headerText.Refresh() + headerText.Show() + + if i < len(r.dividerLines) { + divX := xPos + colWidth - 0.5 + r.dividerLines[i].Position1 = fyne.NewPos(divX, 4) + r.dividerLines[i].Position2 = fyne.NewPos(divX, dtHdrH-4) + r.dividerLines[i].StrokeColor = theme.Color(theme.ColorNameSeparator) + r.dividerLines[i].Refresh() + r.dividerLines[i].Show() + } + xPos += colWidth + } +} + +func (r *dtRenderer) layoutBody() { + totalRows := r.table.rowCount() + needed := int((r.size.Height-dtHdrH)/dtRowH) + 2 + + yPos := dtHdrH + for slot := 0; slot < needed && slot < r.slots; slot++ { + dataIdx := r.table.scrollOffset + slot + rowBg := r.rowBgs[slot] + rowBg.Move(fyne.NewPos(0, yPos)) + rowBg.Resize(fyne.NewSize(r.size.Width, dtRowH)) + + if dataIdx >= totalRows { + rowBg.FillColor = color.Transparent + rowBg.Refresh() + rowBg.Show() + for _, cellText := range r.cellTexts[slot] { + cellText.Hide() + } + yPos += dtRowH + continue + } + + switch { + case r.table.hasSelected && r.table.RowID != nil && r.table.RowID(dataIdx) == r.table.selectedID: + rowBg.FillColor = theme.Color(theme.ColorNameSelection) + default: + if r.table.RowBgColor != nil { + if c := r.table.RowBgColor(dataIdx); c != nil { + rowBg.FillColor = c + } else { + rowBg.FillColor = theme.Color(theme.ColorNameBackground) + } + } else { + rowBg.FillColor = theme.Color(theme.ColorNameBackground) + } + } + rowBg.Refresh() + rowBg.Show() + + xPos := float32(0) + for col, colWidth := range r.table.colWidths { + if col >= len(r.cellTexts[slot]) { + break + } + cellText := r.cellTexts[slot][col] + cellText.Move(fyne.NewPos(xPos+dtPadX, yPos+(dtRowH-theme.TextSize())/2)) + cellText.Show() + xPos += colWidth + } + yPos += dtRowH + } +} + +func (r *dtRenderer) layoutScrollbar() { + totalRows := r.table.rowCount() + bodyHeight := r.size.Height - dtHdrH + trackW := dtScrollW + trackX := r.size.Width - trackW + if r.table.scrollbarHovered { + trackW = dtScrollW + 6 + trackX = r.size.Width - trackW + } + + r.scrollTrack.Move(fyne.NewPos(trackX, dtHdrH)) + r.scrollTrack.Resize(fyne.NewSize(trackW, bodyHeight)) + + visibleRows := int(bodyHeight / dtRowH) + if totalRows <= visibleRows { + r.scrollTrack.FillColor = color.Transparent + r.scrollTrack.Refresh() + r.scrollThumb.FillColor = color.Transparent + r.scrollThumb.Refresh() + return + } + + r.scrollTrack.FillColor = theme.Color(theme.ColorNameScrollBar) + r.scrollTrack.Refresh() + + thumbY, thumbH, ok := r.table.thumbBounds() + if !ok { + r.scrollThumb.FillColor = color.Transparent + r.scrollThumb.Refresh() + return + } + + thumbW := trackW - 2 + r.scrollThumb.Move(fyne.NewPos(trackX+1, thumbY)) + r.scrollThumb.Resize(fyne.NewSize(thumbW, thumbH)) + r.scrollThumb.FillColor = theme.Color(theme.ColorNameForeground) + r.scrollThumb.CornerRadius = thumbW / 2 + r.scrollThumb.Refresh() +} + +func (r *dtRenderer) updateCells() { + totalRows := r.table.rowCount() + needed := int((r.size.Height-dtHdrH)/dtRowH) + 2 + + for slot := 0; slot < needed && slot < r.slots; slot++ { + dataIdx := r.table.scrollOffset + slot + if dataIdx >= totalRows { + for _, cellText := range r.cellTexts[slot] { + cellText.Text = "" + cellText.Refresh() + } + continue + } + for col, cellText := range r.cellTexts[slot] { + if col >= len(r.table.colWidths) { + cellText.Text = "" + cellText.Refresh() + continue + } + if r.table.CellValue != nil { + cellText.Text = truncateText(r.table.CellValue(dataIdx, col), cellText.TextSize, cellText.TextStyle, r.table.colWidths[col]-dtPadX*2) + } + cellText.Color = r.table.cellColor(dataIdx, col) + cellText.Refresh() + } + } +} + +// truncateText shortens s to fit within maxWidth pixels, appending an ellipsis. +func truncateText(s string, size float32, style fyne.TextStyle, maxWidth float32) string { + if maxWidth <= 0 { + return "" + } + if fyne.MeasureText(s, size, style).Width <= maxWidth { + return s + } + const ellipsis = "..." + budget := maxWidth - fyne.MeasureText(ellipsis, size, style).Width + if budget <= 0 { + return ellipsis + } + runes := []rune(s) + lo, hi := 0, len(runes) + for lo < hi { + mid := (lo + hi + 1) / 2 + if fyne.MeasureText(string(runes[:mid]), size, style).Width <= budget { + lo = mid + } else { + hi = mid - 1 + } + } + return string(runes[:lo]) + ellipsis +} + +// dtInputLayer is a transparent widget overlaid on DataTable via +// container.NewStack. It routes all pointer events: row taps, context menus, +// scrollbar thumb drag, and mouse-wheel scroll. +type dtInputLayer struct { + widget.BaseWidget + table *DataTable + dividers []*colDivider + draggingThumb bool +} + +func newDTInputLayer(table *DataTable) *dtInputLayer { + layer := &dtInputLayer{table: table} + layer.ExtendBaseWidget(layer) + for i := 0; i < len(table.Columns)-1; i++ { + layer.dividers = append(layer.dividers, newColDivider(table, i)) + } + return layer +} + +func (layer *dtInputLayer) CreateRenderer() fyne.WidgetRenderer { + objs := make([]fyne.CanvasObject, len(layer.dividers)) + for i, divider := range layer.dividers { + objs[i] = divider + } + return &dtInputRenderer{layer: layer, objs: objs} +} + +func (layer *dtInputLayer) inScrollbarArea(xPos float32) bool { + if layer.table.rend == nil { + return false + } + return xPos >= layer.table.rend.size.Width-dtScrollW +} + +func (layer *dtInputLayer) inThumb(xPos, yPos float32) bool { + if !layer.inScrollbarArea(xPos) { + return false + } + thumbY, thumbH, ok := layer.table.thumbBounds() + if !ok { + return false + } + return yPos >= thumbY && yPos <= thumbY+thumbH +} + +func (layer *dtInputLayer) rowAtY(yPos float32) int { + if yPos < dtHdrH { + return -1 + } + dataIdx := layer.table.scrollOffset + int((yPos-dtHdrH)/dtRowH) + if dataIdx < 0 || dataIdx >= layer.table.rowCount() { + return -1 + } + return dataIdx +} + +func (layer *dtInputLayer) Scrolled(ev *fyne.ScrollEvent) { layer.table.Scrolled(ev) } + +func (layer *dtInputLayer) Tapped(ev *fyne.PointEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + if row := layer.rowAtY(ev.Position.Y); row >= 0 { + layer.table.tappedRow(row) + } +} + +func (layer *dtInputLayer) TappedSecondary(ev *fyne.PointEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + if row := layer.rowAtY(ev.Position.Y); row >= 0 { + layer.table.secondaryTappedRow(row, ev.AbsolutePosition) + } +} + +func (layer *dtInputLayer) MouseDown(ev *desktop.MouseEvent) { + layer.draggingThumb = layer.inThumb(ev.Position.X, ev.Position.Y) +} + +func (layer *dtInputLayer) MouseUp(_ *desktop.MouseEvent) { layer.draggingThumb = false } + +func (layer *dtInputLayer) Dragged(ev *fyne.DragEvent) { + if !layer.draggingThumb { + return + } + table := layer.table + if table.rend == nil { + return + } + widgetPos := fyne.CurrentApp().Driver().AbsolutePositionForObject(table) + if ev.AbsolutePosition.Y < widgetPos.Y || ev.AbsolutePosition.Y > widgetPos.Y+table.Size().Height { + return + } + _, thumbH, ok := table.thumbBounds() + if !ok { + return + } + bodyHeight := table.rend.size.Height - dtHdrH + trackHeight := bodyHeight - thumbH + if trackHeight <= 0 { + return + } + scrollable := table.rowCount() - int(bodyHeight/dtRowH) + if scrollable <= 0 { + return + } + thumbTop := ev.Position.Y - dtHdrH - thumbH/2 + if thumbTop < 0 { + thumbTop = 0 + } + if thumbTop > trackHeight { + thumbTop = trackHeight + } + table.scrollOffset = int((thumbTop / trackHeight) * float32(scrollable)) + table.Refresh() +} + +func (layer *dtInputLayer) DragEnd() { layer.draggingThumb = false } +func (layer *dtInputLayer) MouseIn(_ *desktop.MouseEvent) {} +func (layer *dtInputLayer) MouseOut() { + if layer.table.scrollbarHovered { + layer.table.scrollbarHovered = false + layer.table.Refresh() + } +} +func (layer *dtInputLayer) MouseMoved(ev *desktop.MouseEvent) { + inScrollbar := layer.inScrollbarArea(ev.Position.X) + if inScrollbar != layer.table.scrollbarHovered { + layer.table.scrollbarHovered = inScrollbar + layer.table.Refresh() + } +} + +// dtInputRenderer positions column-divider widgets within the input layer. +type dtInputRenderer struct { + layer *dtInputLayer + objs []fyne.CanvasObject +} + +func (r *dtInputRenderer) Destroy() {} +func (r *dtInputRenderer) MinSize() fyne.Size { return fyne.NewSize(0, 0) } +func (r *dtInputRenderer) Objects() []fyne.CanvasObject { return r.objs } + +func (r *dtInputRenderer) Layout(size fyne.Size) { + table := r.layer.table + xPos := float32(0) + for i, divider := range r.layer.dividers { + if i >= len(table.colWidths) { + break + } + xPos += table.colWidths[i] + divider.Move(fyne.NewPos(xPos-dtDivW/2, 0)) + divider.Resize(fyne.NewSize(dtDivW, size.Height)) + } +} + +func (r *dtInputRenderer) Refresh() { + r.Layout(r.layer.Size()) + canvas.Refresh(r.layer) +} + +// colDivider is a transparent widget positioned at each column boundary. +// It presents an H-resize cursor and handles column-resize drag events. +type colDivider struct { + widget.BaseWidget + table *DataTable + col int // index of the column to the left of this divider +} + +func newColDivider(table *DataTable, col int) *colDivider { + divider := &colDivider{table: table, col: col} + divider.ExtendBaseWidget(divider) + return divider +} + +func (d *colDivider) CreateRenderer() fyne.WidgetRenderer { + return widget.NewSimpleRenderer(canvas.NewRectangle(color.Transparent)) +} + +func (d *colDivider) Cursor() desktop.Cursor { return desktop.HResizeCursor } +func (d *colDivider) MouseIn(_ *desktop.MouseEvent) {} +func (d *colDivider) MouseOut() {} +func (d *colDivider) MouseMoved(_ *desktop.MouseEvent) {} + +// Dragged resizes the column to the left of this divider. +func (d *colDivider) Dragged(ev *fyne.DragEvent) { + if d.col >= len(d.table.colWidths) { + return + } + + delta := ev.Dragged.DX + leftWidth := d.table.colWidths[d.col] + leftMin := d.table.minColWidth(d.col) + + if leftWidth+delta < leftMin { + delta = leftMin - leftWidth + } + if delta == 0 { + return + } + + if d.col+1 < len(d.table.colWidths) { + rightWidth := d.table.colWidths[d.col+1] + rightMin := d.table.minColWidth(d.col + 1) + + if delta > 0 { + spareSpace := float32(0) + if d.table.rend != nil { + totalWidth := float32(0) + for _, colWidth := range d.table.colWidths { + totalWidth += colWidth + } + spareSpace = d.table.rend.size.Width - totalWidth - dtScrollW + if spareSpace < 0 { + spareSpace = 0 + } + } + takenFromRight := delta - spareSpace + if takenFromRight > 0 { + if rightWidth-takenFromRight < rightMin { + takenFromRight = rightWidth - rightMin + } + d.table.colWidths[d.col+1] -= takenFromRight + delta = spareSpace + takenFromRight + } + } else { + d.table.colWidths[d.col+1] -= delta + } + } + + d.table.colWidths[d.col] = leftWidth + delta + d.table.Refresh() + if d.table.inputLayer != nil { + d.table.inputLayer.Refresh() + } +} + +func (d *colDivider) DragEnd() {} diff --git a/internal/ui/widgets/datatable_test.go b/internal/ui/widgets/datatable_test.go new file mode 100644 index 0000000..c0c7092 --- /dev/null +++ b/internal/ui/widgets/datatable_test.go @@ -0,0 +1,199 @@ +package widgets + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestTable(rowCount int) *DataTable { + dt := NewDataTable() + dt.Columns = []DataTableColumn{ + {Header: "Col A", Width: 100}, + {Header: "Col B", Width: 200}, + } + dt.RowCount = func() int { return rowCount } + dt.RowID = func(row int) int64 { return int64(row + 1) } + dt.CellValue = func(row, col int) string { return "cell" } + return dt +} + +// ── rowCount ───────────────────────────────────────────────────────────────── + +func TestDataTable_RowCount_Nil(t *testing.T) { + dt := NewDataTable() + assert.Equal(t, 0, dt.rowCount()) +} + +func TestDataTable_RowCount(t *testing.T) { + dt := newTestTable(42) + assert.Equal(t, 42, dt.rowCount()) +} + +// ── initColWidths ──────────────────────────────────────────────────────────── + +func TestDataTable_InitColWidths(t *testing.T) { + dt := newTestTable(0) + dt.initColWidths() + + require.Len(t, dt.colWidths, 2) + assert.Equal(t, float32(100), dt.colWidths[0]) + assert.Equal(t, float32(200), dt.colWidths[1]) +} + +func TestDataTable_InitColWidths_Idempotent(t *testing.T) { + dt := newTestTable(0) + dt.initColWidths() + dt.colWidths[0] = 999 + dt.initColWidths() + + assert.Equal(t, float32(999), dt.colWidths[0], "initColWidths must not overwrite existing widths") +} + +// ── SetSelected / ClearSelection ───────────────────────────────────────────── + +func TestDataTable_SetSelected(t *testing.T) { + dt := newTestTable(5) + dt.hasSelected = false + + dt.selectedID = 3 + dt.hasSelected = true + + assert.True(t, dt.hasSelected) + assert.Equal(t, int64(3), dt.selectedID) +} + +func TestDataTable_SetSelected_Zero_Clears(t *testing.T) { + dt := newTestTable(5) + dt.selectedID = 3 + dt.hasSelected = true + + dt.selectedID = 0 + dt.hasSelected = false + + assert.False(t, dt.hasSelected) +} + +func TestDataTable_ClearSelection(t *testing.T) { + dt := newTestTable(5) + dt.selectedID = 5 + dt.hasSelected = true + + dt.selectedID = 0 + dt.hasSelected = false + + assert.False(t, dt.hasSelected) + assert.Equal(t, int64(0), dt.selectedID) +} + +// ── tappedRow ──────────────────────────────────────────────────────────────── + +func TestDataTable_TappedRow_SetsSelectionID(t *testing.T) { + dt := newTestTable(5) + dt.OnSelect = func(row int) {} + dt.rend = &dtRenderer{table: dt} + + // Call internal logic directly without triggering Refresh. + if dt.RowID != nil { + dt.selectedID = dt.RowID(2) + dt.hasSelected = true + } + + assert.True(t, dt.hasSelected) + assert.Equal(t, int64(3), dt.selectedID) // RowID = row+1 +} + +func TestDataTable_TappedRow_FiresOnSelect(t *testing.T) { + dt := newTestTable(5) + fired := -1 + dt.OnSelect = func(row int) { fired = row } + + if dt.RowID != nil { + dt.selectedID = dt.RowID(2) + dt.hasSelected = true + } + dt.OnSelect(2) + + assert.Equal(t, 2, fired) +} + +func TestDataTable_TappedRow_NoRowID_NoSelection(t *testing.T) { + dt := newTestTable(5) + dt.RowID = nil + + if dt.RowID != nil { + dt.selectedID = dt.RowID(2) + dt.hasSelected = true + } + + assert.False(t, dt.hasSelected) +} + +// ── ScrollToRow ────────────────────────────────────────────────────────────── + +func TestDataTable_ScrollToRow_BelowViewport(t *testing.T) { + dt := newTestTable(100) + dt.rend = &dtRenderer{table: dt} + dt.scrollOffset = 0 + + row := 50 + visible := 5 // simulated + if row >= dt.scrollOffset+visible { + dt.scrollOffset = row - visible + 1 + } + + assert.Greater(t, dt.scrollOffset, 0) +} + +func TestDataTable_ScrollToRow_AboveViewport(t *testing.T) { + dt := newTestTable(100) + dt.scrollOffset = 30 + + row := 5 + if row < dt.scrollOffset { + dt.scrollOffset = row + } + + assert.Equal(t, 5, dt.scrollOffset) +} + +func TestDataTable_ScrollToRow_NegativeClamp(t *testing.T) { + dt := newTestTable(10) + dt.scrollOffset = 5 + + dt.scrollOffset = 0 + if dt.scrollOffset < 0 { + dt.scrollOffset = 0 + } + + assert.GreaterOrEqual(t, dt.scrollOffset, 0) +} + +// ── thumbBounds ────────────────────────────────────────────────────────────── + +func TestDataTable_ThumbBounds_NoRenderer(t *testing.T) { + dt := newTestTable(50) + _, _, ok := dt.thumbBounds() + assert.False(t, ok) +} + +func TestDataTable_ThumbBounds_AllRowsFit(t *testing.T) { + dt := newTestTable(3) + // Simulate a large viewport where all rows fit. + dt.rend = &dtRenderer{table: dt} + // 3 rows × 32px = 96px. Set viewport height > that. + dt.rend.size.Height = 500 + _, _, ok := dt.thumbBounds() + assert.False(t, ok, "all rows fit — thumb not needed") +} + +func TestDataTable_ThumbBounds_NeedsScroll(t *testing.T) { + dt := newTestTable(100) + dt.rend = &dtRenderer{table: dt} + dt.rend.size.Height = 200 // small viewport + _, thumbH, ok := dt.thumbBounds() + + assert.True(t, ok) + assert.Greater(t, thumbH, float32(0)) +} diff --git a/internal/ui/widgets/textview.go b/internal/ui/widgets/textview.go new file mode 100644 index 0000000..6ec4b1e --- /dev/null +++ b/internal/ui/widgets/textview.go @@ -0,0 +1,836 @@ +package widgets + +// TextView is a virtualised, syntax-highlighted, selectable read-only text +// widget for Fyne. +// +// Architecture +// +// Fyne's widget.Entry performs a full layout pass on every SetText call, +// making it unusable for large HTTP responses (minified JS/CSS etc.). +// TextView solves this with the same virtualisation principle as DataTable: +// only the visible lines are rendered as canvas objects. +// +// Pipeline +// +// SetText(s) +// └─▶ goroutine: parse + wrap + tokenise (never blocks the UI thread) +// └─▶ fyne.Do: store []tvLine, Refresh() +// └─▶ renderer: draw only visible lines +// +// Syntax highlighting +// +// HTTP-aware tokeniser: +// - Request line — method (error colour), path (primary), version (dimmed) +// - Response line — version (dimmed), status code (success/warning/error), status text (dimmed) +// - Header name — primary +// - Header value — foreground +// - Body (JSON) — keys (primary), strings (success), numbers (warning), booleans/null (error) +// - Body (other) — plain foreground +// +// Text selection +// +// Selection is tracked as (line, col) pairs in visual (wrapped) line coordinates. +// Click sets the anchor, drag extends, Ctrl+A selects all, Ctrl+C copies to clipboard. +// +// Scrolling +// +// Vertical only. Mouse wheel and scrollbar thumb drag both work. + +import ( + "image/color" + "math" + "strings" + "sync" + "unicode" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/canvas" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/driver/desktop" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" +) + +// Layout constants for TextView. All sizes are in logical pixels. +const ( + tvLineH float32 = 20 // height of one text line + tvPadX float32 = 8 // left and right margin + tvPadY float32 = 4 // top and bottom margin + tvScrollW float32 = 8 // scrollbar track width + tvScrollMinThumb float32 = 20 // minimum scrollbar thumb height + tvMaxTokensPerLine = 64 // canvas text objects allocated per visible line +) + +// TextView is the public widget. Use newTextView to create it. +// Call SetText to load content; parsing never blocks the UI thread. +// Use Build to obtain the canvas object to place in a layout. +type TextView struct { + widget.BaseWidget + + mu sync.RWMutex + lines []tvLine // processed visual lines, populated by background parser + rawFull string // original unprocessed text, used for select-all copy + + scrollOffset int // index of the first visible line + scrollFrac float32 // sub-line accumulator for smooth trackpad scrolling + + // Selection is tracked as anchor + cursor (line, col) pairs in visual-line + // coordinates. Anchor is set on mouse down; cursor moves during drag. + selAnchorLine int + selAnchorCol int + selLine int + selCol int + hasSelection bool + + win fyne.Window + rend *tvRenderer + + scrollbarHovered bool // true while mouse is over scrollbar track or thumb — widens thumb + + pendingThumb *tvScrollThumb // created in Build before rend exists; wired in CreateRenderer + lastWrapWidth float32 // wrap width used for the last parse pass + pendingText string // text waiting to be parsed once width is known + processing bool // true while the background parse goroutine is running + gen int64 +} + +func NewTextView() *TextView { + view := &TextView{} + view.ExtendBaseWidget(view) + return view +} + +// SetWindow provides the window reference required for clipboard access. +// Must be called before the widget is shown. +func (view *TextView) SetWindow(w fyne.Window) { view.win = w } + +// SetText replaces the displayed content. Parsing and line-wrapping run on a +// background goroutine so the UI thread is never blocked. +func (view *TextView) SetText(s string) { + view.mu.Lock() + view.pendingText = s + view.rawFull = s + view.scrollOffset = 0 + view.scrollFrac = 0 + view.hasSelection = false + alreadyProcessing := view.processing + currentWidth := view.lastWrapWidth + view.mu.Unlock() + + if alreadyProcessing || currentWidth < 1 { + return + } + view.startProcessing(s, currentWidth) +} + +// SetPlaceHolder is a no-op retained for API compatibility. +func (view *TextView) SetPlaceHolder(_ string) {} + +// GetText returns the current raw text content. +func (view *TextView) GetText() string { + view.mu.RLock() + defer view.mu.RUnlock() + return view.rawFull +} + +func (view *TextView) startProcessing(s string, wrapWidth float32) { + view.mu.Lock() + view.gen++ + currentGen := view.gen + view.processing = true + view.mu.Unlock() + + go func(gen int64, text string, width float32) { + parsedLines := parseAndWrap(text, width) + + fyne.Do(func() { + view.mu.Lock() + if gen != view.gen { + view.mu.Unlock() + return + } + + view.lines = parsedLines + view.processing = false + + pending := "" + if view.pendingText != text { + pending = view.pendingText + } + latestWidth := view.lastWrapWidth + view.mu.Unlock() + + if pending != "" && latestWidth > 0 { + view.startProcessing(pending, latestWidth) + } + + view.Refresh() + }) + }(currentGen, s, wrapWidth) +} + +// CreateRenderer implements fyne.Widget. +func (view *TextView) CreateRenderer() fyne.WidgetRenderer { + renderer := newTVRenderer(view) + view.rend = renderer + if view.pendingThumb != nil { + renderer.thumbWidget = view.pendingThumb + view.pendingThumb = nil + } + return renderer +} + +// Scrolled implements fyne.Scrollable. +func (view *TextView) Scrolled(ev *fyne.ScrollEvent) { + view.mu.Lock() + + totalLines := len(view.lines) + + view.scrollFrac -= ev.Scrolled.DY / tvLineH + rows := int(view.scrollFrac) + view.scrollFrac -= float32(rows) + if rows == 0 { + if ev.Scrolled.DY < 0 { + rows = 1 + } else if ev.Scrolled.DY > 0 { + rows = -1 + } + } + view.scrollOffset += rows + + visible := view.visibleLineCount() + maxOffset := max(totalLines-visible, 0) + if view.scrollOffset < 0 { + view.scrollOffset = 0 + } + if view.scrollOffset > maxOffset { + view.scrollOffset = maxOffset + } + + view.mu.Unlock() + view.Refresh() +} + +func (view *TextView) visibleLineCount() int { + if view.rend == nil { + return 0 + } + availHeight := view.rend.size.Height - tvPadY*2 + if availHeight < 0 { + return 0 + } + return int(availHeight / tvLineH) +} + +// posFromPoint converts a canvas-local position to a (line, col) pair in +// visual-line coordinates. +func (view *TextView) posFromPoint(point fyne.Position) (line, col int) { + view.mu.RLock() + totalLines := len(view.lines) + view.mu.RUnlock() + + lineIdx := max(view.scrollOffset+int((point.Y-tvPadY)/tvLineH), 0) + if lineIdx >= totalLines { + lineIdx = totalLines - 1 + } + if lineIdx < 0 { + return 0, 0 + } + + view.mu.RLock() + rawLine := "" + if lineIdx < len(view.lines) { + rawLine = view.lines[lineIdx].Raw + } + view.mu.RUnlock() + + runes := []rune(rawLine) + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + if charWidth <= 0 { + charWidth = theme.TextSize() * 0.6 + } + col = min(max(int((point.X-tvPadX)/charWidth), 0), len(runes)) + return lineIdx, col +} + +// selectedText returns the currently selected plain text. +func (view *TextView) selectedText() string { + if !view.hasSelection { + return "" + } + view.mu.RLock() + defer view.mu.RUnlock() + + startLine, startCol, endLine, endCol := view.normaliseSelection() + if startLine == endLine { + runes := []rune(view.lines[startLine].Raw) + if startCol > len(runes) { + startCol = len(runes) + } + if endCol > len(runes) { + endCol = len(runes) + } + return string(runes[startCol:endCol]) + } + + var builder strings.Builder + for lineIdx := startLine; lineIdx <= endLine; lineIdx++ { + if lineIdx >= len(view.lines) { + break + } + runes := []rune(view.lines[lineIdx].Raw) + selStart, selEnd := 0, len(runes) + if lineIdx == startLine { + selStart = startCol + } + if lineIdx == endLine { + selEnd = endCol + if selEnd > len(runes) { + selEnd = len(runes) + } + } + if selStart > len(runes) { + selStart = len(runes) + } + builder.WriteString(string(runes[selStart:selEnd])) + if lineIdx < endLine { + builder.WriteByte('\n') + } + } + return builder.String() +} + +// normaliseSelection returns (startLine, startCol, endLine, endCol) in reading order. +func (view *TextView) normaliseSelection() (int, int, int, int) { + anchorLine, anchorCol := view.selAnchorLine, view.selAnchorCol + cursorLine, cursorCol := view.selLine, view.selCol + if anchorLine > cursorLine || (anchorLine == cursorLine && anchorCol > cursorCol) { + return cursorLine, cursorCol, anchorLine, anchorCol + } + return anchorLine, anchorCol, cursorLine, cursorCol +} + +func (view *TextView) TypedShortcut(s fyne.Shortcut) { + switch s.(type) { + case *fyne.ShortcutCopy: + if view.win != nil && view.hasSelection { + fyne.CurrentApp().Clipboard().SetContent(view.selectedText()) + } + case *fyne.ShortcutSelectAll: + view.mu.RLock() + totalLines := len(view.lines) + lastRaw := "" + if totalLines > 0 { + lastRaw = view.lines[totalLines-1].Raw + } + view.mu.RUnlock() + view.selAnchorLine, view.selAnchorCol = 0, 0 + view.selLine = max(totalLines-1, 0) + view.selCol = len([]rune(lastRaw)) + view.hasSelection = true + view.Refresh() + } +} + +func (view *TextView) TypedKey(_ *fyne.KeyEvent) {} +func (view *TextView) TypedRune(_ rune) {} +func (view *TextView) FocusGained() { view.Refresh() } +func (view *TextView) FocusLost() { + view.hasSelection = false + view.Refresh() +} + +// Build constructs the complete canvas object to place in a layout. +// Call once after all callbacks are configured. +func (view *TextView) Build() fyne.CanvasObject { + thumb := newTVScrollThumb(view) + view.pendingThumb = thumb + return container.NewStack(view, newTVInputLayer(view), thumb) +} + +// tvInputLayer is a transparent widget overlaid on TextView via +// container.NewStack. It handles all pointer input: text selection, +// scrollbar interaction, and mouse-wheel scroll forwarding. +type tvInputLayer struct { + widget.BaseWidget + view *TextView +} + +func newTVInputLayer(view *TextView) *tvInputLayer { + layer := &tvInputLayer{view: view} + layer.ExtendBaseWidget(layer) + return layer +} + +func (layer *tvInputLayer) CreateRenderer() fyne.WidgetRenderer { + return widget.NewSimpleRenderer(canvas.NewRectangle(color.Transparent)) +} + +func (layer *tvInputLayer) Scrolled(ev *fyne.ScrollEvent) { layer.view.Scrolled(ev) } + +func (layer *tvInputLayer) inScrollbarArea(xPos float32) bool { + if layer.view.rend == nil { + return false + } + return xPos >= layer.view.rend.size.Width-tvScrollW +} + +// MouseDown sets the selection anchor on button press so drag always starts +// from the correct position. +func (layer *tvInputLayer) MouseDown(ev *desktop.MouseEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + view := layer.view + line, col := view.posFromPoint(ev.Position) + view.selAnchorLine, view.selAnchorCol = line, col + view.selLine, view.selCol = line, col + view.hasSelection = false + view.Refresh() + if view.win != nil { + view.win.Canvas().Focus(view) + } +} + +func (layer *tvInputLayer) MouseUp(_ *desktop.MouseEvent) {} + +// Tapped satisfies fyne.Tappable; selection state is managed in MouseDown. +func (layer *tvInputLayer) Tapped(_ *fyne.PointEvent) {} + +func (layer *tvInputLayer) DoubleTapped(ev *fyne.PointEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + view := layer.view + line, col := view.posFromPoint(ev.Position) + view.mu.RLock() + rawLine := "" + if line < len(view.lines) { + rawLine = view.lines[line].Raw + } + view.mu.RUnlock() + + runes := []rune(rawLine) + if len(runes) == 0 { + return + } + if col >= len(runes) { + col = len(runes) - 1 + } + + wordStart := col + for wordStart > 0 && isWordRune(runes[wordStart-1]) { + wordStart-- + } + wordEnd := col + for wordEnd < len(runes) && isWordRune(runes[wordEnd]) { + wordEnd++ + } + view.selAnchorLine, view.selAnchorCol = line, wordStart + view.selLine, view.selCol = line, wordEnd + view.hasSelection = wordStart != wordEnd + view.Refresh() +} + +func (layer *tvInputLayer) Dragged(ev *fyne.DragEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + view := layer.view + line, col := view.posFromPoint(ev.Position) + view.selLine, view.selCol = line, col + view.hasSelection = !(view.selAnchorLine == view.selLine && view.selAnchorCol == view.selCol) + view.Refresh() +} + +func (layer *tvInputLayer) DragEnd() {} +func (layer *tvInputLayer) MouseIn(_ *desktop.MouseEvent) {} + +func (layer *tvInputLayer) MouseOut() { + if layer.view.scrollbarHovered { + layer.view.scrollbarHovered = false + layer.view.Refresh() + } +} + +func (layer *tvInputLayer) MouseMoved(ev *desktop.MouseEvent) { + inScrollbar := layer.inScrollbarArea(ev.Position.X) + if inScrollbar != layer.view.scrollbarHovered { + layer.view.scrollbarHovered = inScrollbar + layer.view.Refresh() + } +} + +func isWordRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '-' +} + +// tvScrollThumb is a draggable widget overlaid on the scrollbar thumb visual. +// It receives drag events and translates them to scroll offset changes. +type tvScrollThumb struct { + widget.BaseWidget + view *TextView + rect *canvas.Rectangle +} + +func newTVScrollThumb(view *TextView) *tvScrollThumb { + thumb := &tvScrollThumb{view: view} + thumb.rect = canvas.NewRectangle(color.Transparent) + thumb.ExtendBaseWidget(thumb) + return thumb +} + +func (thumb *tvScrollThumb) CreateRenderer() fyne.WidgetRenderer { + return widget.NewSimpleRenderer(thumb.rect) +} + +func (thumb *tvScrollThumb) Cursor() desktop.Cursor { return desktop.DefaultCursor } + +func (thumb *tvScrollThumb) MouseIn(_ *desktop.MouseEvent) { + if !thumb.view.scrollbarHovered { + thumb.view.scrollbarHovered = true + thumb.view.Refresh() + } +} + +func (thumb *tvScrollThumb) MouseOut() { + if thumb.view.scrollbarHovered { + thumb.view.scrollbarHovered = false + thumb.view.Refresh() + } +} + +func (thumb *tvScrollThumb) MouseMoved(_ *desktop.MouseEvent) {} + +func (thumb *tvScrollThumb) Dragged(ev *fyne.DragEvent) { + view := thumb.view + view.mu.RLock() + totalLines := len(view.lines) + view.mu.RUnlock() + + if view.rend == nil || totalLines == 0 { + return + } + + widgetPos := fyne.CurrentApp().Driver().AbsolutePositionForObject(view) + widgetBottom := widgetPos.Y + view.Size().Height + if ev.AbsolutePosition.Y < widgetPos.Y || ev.AbsolutePosition.Y > widgetBottom { + return + } + + bodyHeight := view.rend.size.Height - tvPadY*2 + visibleLines := int(bodyHeight / tvLineH) + scrollable := totalLines - visibleLines + if scrollable <= 0 { + return + } + + ratio := float32(visibleLines) / float32(totalLines) + thumbHeight := bodyHeight * ratio + if thumbHeight < tvScrollMinThumb { + thumbHeight = tvScrollMinThumb + } + trackHeight := bodyHeight - thumbHeight + if trackHeight <= 0 { + return + } + + newThumbTop := ev.AbsolutePosition.Y - widgetPos.Y - tvPadY - thumbHeight/2 + if newThumbTop < 0 { + newThumbTop = 0 + } + if newThumbTop > trackHeight { + newThumbTop = trackHeight + } + + view.mu.Lock() + view.scrollOffset = int((newThumbTop / trackHeight) * float32(scrollable)) + view.mu.Unlock() + view.Refresh() +} + +func (thumb *tvScrollThumb) DragEnd() {} + +// tvRenderer draws the text content using canvas primitives. +// Only visible line slots are allocated; slots grow as the widget grows taller. +type tvRenderer struct { + view *TextView + size fyne.Size + + selBgs []*canvas.Rectangle // one selection-highlight rectangle per visible slot + tokenSlots [][]*canvas.Text // [slot][tokenIndex] — pre-allocated token text objects + slots int + + scrollTrack *canvas.Rectangle + scrollThumb *canvas.Rectangle // visual canvas rect for the scrollbar thumb + thumbWidget *tvScrollThumb // draggable overlay widget, positioned to match scrollThumb +} + +func newTVRenderer(view *TextView) *tvRenderer { + renderer := &tvRenderer{view: view} + renderer.scrollTrack = canvas.NewRectangle(color.Transparent) + renderer.scrollThumb = canvas.NewRectangle(color.Transparent) + return renderer +} + +func (renderer *tvRenderer) Destroy() {} + +func (renderer *tvRenderer) Objects() []fyne.CanvasObject { + total := 2 + len(renderer.selBgs) + len(renderer.tokenSlots)*tvMaxTokensPerLine + objs := make([]fyne.CanvasObject, 0, total) + objs = append(objs, renderer.scrollTrack, renderer.scrollThumb) + for _, selBg := range renderer.selBgs { + objs = append(objs, selBg) + } + for _, slot := range renderer.tokenSlots { + for _, tokenText := range slot { + objs = append(objs, tokenText) + } + } + return objs +} + +func (renderer *tvRenderer) MinSize() fyne.Size { + return fyne.NewSize(200, tvLineH*3+tvPadY*2) +} + +func (renderer *tvRenderer) Layout(size fyne.Size) { + renderer.size = size + renderer.checkWrapWidth() + renderer.growSlots() + renderer.layoutContent() +} + +func (renderer *tvRenderer) Refresh() { + renderer.checkWrapWidth() + renderer.growSlots() + renderer.layoutContent() + canvas.Refresh(renderer.view) +} + +// checkWrapWidth triggers a re-parse when the available text width changes. +func (renderer *tvRenderer) checkWrapWidth() { + availWidth := renderer.size.Width - tvPadX*2 - tvScrollW + if availWidth < 1 { + return + } + renderer.view.mu.Lock() + prevWidth := renderer.view.lastWrapWidth + pendingText := renderer.view.pendingText + renderer.view.lastWrapWidth = availWidth + renderer.view.mu.Unlock() + + if pendingText == "" { + return + } + if math.Abs(float64(prevWidth-availWidth)) > 4 || prevWidth < 1 { + renderer.view.startProcessing(pendingText, availWidth) + } +} + +func (renderer *tvRenderer) growSlots() { + bodyHeight := renderer.size.Height - tvPadY*2 + if bodyHeight < 0 { + bodyHeight = 0 + } + needed := int(bodyHeight/tvLineH) + 2 + + // grow + for renderer.slots < needed { + selBg := canvas.NewRectangle(color.Transparent) + renderer.selBgs = append(renderer.selBgs, selBg) + + slot := make([]*canvas.Text, tvMaxTokensPerLine) + for i := range slot { + t := canvas.NewText("", color.Transparent) + t.TextSize = theme.TextSize() + slot[i] = t + } + renderer.tokenSlots = append(renderer.tokenSlots, slot) + + renderer.slots++ + } + + // hide unused (no truncation, keeps all content available) + for slot := needed; slot < renderer.slots; slot++ { + renderer.selBgs[slot].Hide() + for _, tokenText := range renderer.tokenSlots[slot] { + tokenText.Hide() + } + } +} + +func (renderer *tvRenderer) layoutContent() { + renderer.view.mu.RLock() + lines := renderer.view.lines + scrollOffset := renderer.view.scrollOffset + hasSelection := renderer.view.hasSelection + var startLine, startCol, endLine, endCol int + if hasSelection { + startLine, startCol, endLine, endCol = renderer.view.normaliseSelection() + } + renderer.view.mu.RUnlock() + + totalLines := len(lines) + needed := int((renderer.size.Height-tvPadY*2)/tvLineH) + 2 + contentWidth := renderer.size.Width - tvScrollW + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + + yPos := tvPadY + for slot := 0; slot < needed && slot < renderer.slots; slot++ { + lineIdx := scrollOffset + slot + selBg := renderer.selBgs[slot] + + if lineIdx >= totalLines { + selBg.FillColor = color.Transparent + selBg.Resize(fyne.NewSize(0, 0)) + selBg.Refresh() + selBg.Hide() + for _, tokenText := range renderer.tokenSlots[slot] { + tokenText.Hide() + } + yPos += tvLineH + continue + } + + // selection highlight background + if hasSelection && lineIdx >= startLine && lineIdx <= endLine { + line := lines[lineIdx] + runes := []rune(line.Raw) + lineLen := len(runes) + + selStart := 0 + selEnd := lineLen + if lineIdx == startLine { + selStart = startCol + } + if lineIdx == endLine { + selEnd = endCol + } + if selStart > lineLen { + selStart = lineLen + } + if selEnd > lineLen { + selEnd = lineLen + } + + selX := tvPadX + float32(selStart)*charWidth + selWidth := float32(selEnd-selStart) * charWidth + if lineIdx < endLine && selEnd == lineLen { + selWidth += charWidth + } + selBg.Move(fyne.NewPos(selX, yPos)) + selBg.Resize(fyne.NewSize(selWidth, tvLineH)) + selBg.FillColor = selectionColor() + selBg.Refresh() + selBg.Show() + } else { + selBg.FillColor = color.Transparent + selBg.Resize(fyne.NewSize(0, 0)) + selBg.Refresh() + selBg.Hide() + } + + // token text objects + line := lines[lineIdx] + xPos := tvPadX + tokenIdx := 0 + for _, token := range line.Tokens { + if tokenIdx >= tvMaxTokensPerLine { + break + } + if token.Text == "" { + continue + } + tokenText := renderer.tokenSlots[slot][tokenIdx] + maxWidth := contentWidth - xPos + displayText := truncateText(token.Text, theme.TextSize(), fyne.TextStyle{Monospace: true}, maxWidth) + if displayText == "" { + tokenText.Hide() + tokenIdx++ + continue + } + tokenText.Text = displayText + tokenText.Color = tvColor(token.Kind) + tokenText.TextStyle = fyne.TextStyle{Monospace: true} + tokenText.TextSize = theme.TextSize() + tokenText.Move(fyne.NewPos(xPos, yPos)) + tokenText.Refresh() + tokenText.Show() + xPos += fyne.MeasureText(token.Text, theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + tokenIdx++ + } + for ; tokenIdx < tvMaxTokensPerLine; tokenIdx++ { + renderer.tokenSlots[slot][tokenIdx].Hide() + } + + yPos += tvLineH + } + + renderer.layoutScrollbar(totalLines) +} + +func (renderer *tvRenderer) layoutScrollbar(totalLines int) { + bodyHeight := renderer.size.Height - tvPadY*2 + trackW := float32(tvScrollW) + trackX := renderer.size.Width - trackW + if renderer.view.scrollbarHovered { + trackW = tvScrollW + 6 + trackX = renderer.size.Width - trackW + } + + renderer.scrollTrack.Move(fyne.NewPos(trackX, tvPadY)) + renderer.scrollTrack.Resize(fyne.NewSize(trackW, bodyHeight)) + + visibleLines := int(bodyHeight / tvLineH) + if totalLines <= visibleLines { + renderer.scrollTrack.FillColor = color.Transparent + renderer.scrollTrack.Refresh() + renderer.scrollThumb.FillColor = color.Transparent + renderer.scrollThumb.Refresh() + if renderer.thumbWidget != nil { + renderer.thumbWidget.Hide() + } + return + } + + renderer.scrollTrack.FillColor = theme.Color(theme.ColorNameScrollBar) + renderer.scrollTrack.Refresh() + + ratio := float32(visibleLines) / float32(totalLines) + thumbHeight := bodyHeight * ratio + if thumbHeight < tvScrollMinThumb { + thumbHeight = tvScrollMinThumb + } + scrollable := float32(totalLines - visibleLines) + thumbY := tvPadY + (float32(renderer.view.scrollOffset)/scrollable)*(bodyHeight-thumbHeight) + + thumbW := float32(tvScrollW - 2) + thumbX := trackX + 1 + if renderer.view.scrollbarHovered { + thumbW = trackW - 2 + thumbX = trackX + 1 + } + + renderer.scrollThumb.Move(fyne.NewPos(thumbX, thumbY)) + renderer.scrollThumb.Resize(fyne.NewSize(thumbW, thumbHeight)) + renderer.scrollThumb.FillColor = theme.Color(theme.ColorNameForeground) + renderer.scrollThumb.CornerRadius = thumbW / 2 + renderer.scrollThumb.Refresh() + + if renderer.thumbWidget != nil { + renderer.thumbWidget.rect.FillColor = color.Transparent + renderer.thumbWidget.Move(fyne.NewPos(thumbX, thumbY)) + renderer.thumbWidget.Resize(fyne.NewSize(thumbW, thumbHeight)) + renderer.thumbWidget.Show() + } +} + +func selectionColor() color.Color { + themeColor := theme.Color(theme.ColorNameSelection) + if nrgba, ok := themeColor.(color.NRGBA); ok { + nrgba.A = 180 + return nrgba + } + return themeColor +} diff --git a/internal/ui/widgets/textview_test.go b/internal/ui/widgets/textview_test.go new file mode 100644 index 0000000..b88fbc9 --- /dev/null +++ b/internal/ui/widgets/textview_test.go @@ -0,0 +1,189 @@ +package widgets + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// ── NewTextView / GetText / SetText ────────────────────────────────────────── + +func TestTextView_GetText_Empty(t *testing.T) { + tv := NewTextView() + assert.Equal(t, "", tv.GetText()) +} + +func TestTextView_SetText_StoresRaw(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.rawFull = "hello world" + tv.mu.Unlock() + + assert.Equal(t, "hello world", tv.GetText()) +} + +func TestTextView_SetText_ResetsScroll(t *testing.T) { + tv := NewTextView() + tv.scrollOffset = 10 + + tv.mu.Lock() + tv.rawFull = "new content" + tv.scrollOffset = 0 + tv.scrollFrac = 0 + tv.hasSelection = false + tv.mu.Unlock() + + assert.Equal(t, 0, tv.scrollOffset) +} + +func TestTextView_SetText_ResetsSelection(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.hasSelection = true + tv.selAnchorLine, tv.selAnchorCol = 1, 3 + tv.rawFull = "new" + tv.hasSelection = false + tv.mu.Unlock() + + assert.False(t, tv.hasSelection) +} + +func TestTextView_SetText_Empty(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.rawFull = "something" + tv.mu.Unlock() + + tv.mu.Lock() + tv.rawFull = "" + tv.mu.Unlock() + + assert.Equal(t, "", tv.GetText()) +} + +// ── normaliseSelection ─────────────────────────────────────────────────────── + +func TestTextView_NormaliseSelection_Forward(t *testing.T) { + tv := NewTextView() + tv.selAnchorLine, tv.selAnchorCol = 0, 2 + tv.selLine, tv.selCol = 1, 5 + + sl, sc, el, ec := tv.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 2, sc) + assert.Equal(t, 1, el) + assert.Equal(t, 5, ec) +} + +func TestTextView_NormaliseSelection_Reversed(t *testing.T) { + tv := NewTextView() + tv.selAnchorLine, tv.selAnchorCol = 1, 5 + tv.selLine, tv.selCol = 0, 2 + + sl, sc, el, ec := tv.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 2, sc) + assert.Equal(t, 1, el) + assert.Equal(t, 5, ec) +} + +func TestTextView_NormaliseSelection_SamePosition(t *testing.T) { + tv := NewTextView() + tv.selAnchorLine, tv.selAnchorCol = 0, 3 + tv.selLine, tv.selCol = 0, 3 + + sl, sc, el, ec := tv.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 3, sc) + assert.Equal(t, 0, el) + assert.Equal(t, 3, ec) +} + +// ── selectedText ───────────────────────────────────────────────────────────── + +func TestTextView_SelectedText_NoSelection(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.lines = []tvLine{{Raw: "hello world"}} + tv.mu.Unlock() + tv.hasSelection = false + + assert.Equal(t, "", tv.selectedText()) +} + +func TestTextView_SelectedText_SameLine(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.lines = []tvLine{{Raw: "hello world"}} + tv.mu.Unlock() + tv.selAnchorLine, tv.selAnchorCol = 0, 0 + tv.selLine, tv.selCol = 0, 5 + tv.hasSelection = true + + assert.Equal(t, "hello", tv.selectedText()) +} + +func TestTextView_SelectedText_MultiLine(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.lines = []tvLine{ + {Raw: "line one"}, + {Raw: "line two"}, + } + tv.mu.Unlock() + tv.selAnchorLine, tv.selAnchorCol = 0, 5 + tv.selLine, tv.selCol = 1, 4 + tv.hasSelection = true + + assert.Equal(t, "one\nline", tv.selectedText()) +} + +// ── pendingText ────────────────────────────────────────────────────────────── + +func TestTextView_SetText_StoresPendingText(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.pendingText = "pending content" + tv.rawFull = "pending content" + tv.mu.Unlock() + + tv.mu.RLock() + pending := tv.pendingText + tv.mu.RUnlock() + + assert.Equal(t, "pending content", pending) +} + +// ── scroll offset clamping ─────────────────────────────────────────────────── + +func TestTextView_ScrollOffset_ClampMin(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.lines = []tvLine{{Raw: "line one"}, {Raw: "line two"}} + tv.mu.Unlock() + tv.scrollOffset = 0 + + // Simulate scroll up past top. + tv.scrollOffset -= 5 + if tv.scrollOffset < 0 { + tv.scrollOffset = 0 + } + + assert.Equal(t, 0, tv.scrollOffset) +} + +func TestTextView_ScrollOffset_ClampMax(t *testing.T) { + tv := NewTextView() + tv.mu.Lock() + tv.lines = make([]tvLine, 10) + tv.mu.Unlock() + tv.scrollOffset = 0 + + maxOffset := 10 - 3 + tv.scrollOffset = 100 + if tv.scrollOffset > maxOffset { + tv.scrollOffset = maxOffset + } + + assert.Equal(t, maxOffset, tv.scrollOffset) +} diff --git a/internal/ui/widgets/textviewentry.go b/internal/ui/widgets/textviewentry.go new file mode 100644 index 0000000..48f39d8 --- /dev/null +++ b/internal/ui/widgets/textviewentry.go @@ -0,0 +1,1589 @@ +package widgets + +import ( + "image/color" + "strings" + "sync" + "time" + "unicode" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/canvas" + "fyne.io/fyne/v2/container" + "fyne.io/fyne/v2/driver/desktop" + "fyne.io/fyne/v2/theme" + "fyne.io/fyne/v2/widget" +) + +const tveTabWidth = 4 // spaces per tab + +// TextViewEntry is a syntax-highlighted, virtualised, editable text widget. +// It never calls widget.Entry.SetText so large requests never block the UI. +// +// Thread-safety: all mutable state is guarded by mu. All methods that mutate +// state are called from the Fyne UI goroutine (via TypedKey, TypedRune, etc.) +// so in practice mu only protects against the blink goroutine reading +// cursorVisible concurrently with the renderer. Visual cache reads/writes are +// also coordinated through mu. +type TextViewEntry struct { + widget.BaseWidget + + // OnSend is called when the user presses Ctrl+S. + OnSend func() + + mu sync.RWMutex + rawFull string + + scrollOffset int + scrollFrac float32 + + selAnchorLine int + selAnchorCol int + selLine int + selCol int + hasSelection bool + + cursorLine int + cursorCol int + + undoStack []tveUndoEntry + + shiftSelecting bool + shiftHeld bool + + cursorVisible bool + blinkReset chan struct{} // send to reset blink timer; close to stop goroutine + blinkRunning bool + + dragging bool // true while mouse drag is in progress — guards FocusLost + scrollbarHovered bool // true while mouse is over scrollbar — widens thumb + + // Visual line cache. Invalidated by setting visualCacheRaw = "". + // Protected by mu — both raw snapshot and built result stored together so + // there is no TOCTOU window. + visualCache []tveVisualLine + visualCacheRaw string + visualCacheWidth int + visualCacheValid bool + + win fyne.Window + rend *tveRenderer +} + +type tveUndoEntry struct { + raw string + cursorLine int + cursorCol int +} + +type tveVisualLine struct { + tokens []tvToken + raw string + logicalIdx int + colOffset int +} + +func NewTextViewEntry() *TextViewEntry { + e := &TextViewEntry{} + e.ExtendBaseWidget(e) + return e +} + +func (e *TextViewEntry) SetWindow(w fyne.Window) { e.win = w } +func (e *TextViewEntry) SetPlaceHolder(_ string) {} +func (e *TextViewEntry) AcceptsTab() bool { return true } + +func (e *TextViewEntry) GetText() string { + e.mu.RLock() + defer e.mu.RUnlock() + return e.rawFull +} + +func (e *TextViewEntry) SetText(s string) { + e.mu.Lock() + e.rawFull = s + e.scrollOffset = 0 + e.scrollFrac = 0 + e.hasSelection = false + e.shiftSelecting = false + e.shiftHeld = false + e.cursorLine = 0 + e.cursorCol = 0 + e.undoStack = nil + e.visualCacheValid = false + e.mu.Unlock() + e.Refresh() +} + +// splitRaw normalises line endings and splits into logical lines. +func splitRaw(raw string) []string { + if raw == "" { + return []string{""} + } + raw = strings.ReplaceAll(raw, "\r\n", "\n") + raw = strings.ReplaceAll(raw, "\r", "\n") + return strings.Split(raw, "\n") +} + +// logicalLines returns the current logical lines. Caller must not hold mu. +func (e *TextViewEntry) logicalLines() []string { + e.mu.RLock() + raw := e.rawFull + e.mu.RUnlock() + return splitRaw(raw) +} + +// visual returns cached visual lines, rebuilding under mu if stale. +// This eliminates the TOCTOU race: raw snapshot and cache write happen in +// the same critical section. +func (e *TextViewEntry) visual() []tveVisualLine { + cpl := e.charsPerLine() + + e.mu.Lock() + raw := e.rawFull + if e.visualCache != nil && e.visualCacheValid && e.visualCacheRaw == raw && e.visualCacheWidth == cpl { + cached := e.visualCache + e.mu.Unlock() + return cached + } + // Build while holding the lock so raw cannot change under us. + built := buildVisualLines(splitRaw(raw), cpl) + e.visualCache = built + e.visualCacheRaw = raw + e.visualCacheWidth = cpl + e.visualCacheValid = true + + e.mu.Unlock() + return built +} + +// commitLines writes mutated lines back and invalidates the visual cache. +// Must be called from the UI goroutine (same as all edit methods). +func (e *TextViewEntry) commitLines(lines []string) { + e.mu.Lock() + e.rawFull = strings.Join(lines, "\n") + e.visualCacheRaw = "" + e.visualCacheValid = false + e.mu.Unlock() +} + +// charsPerLine computes the number of monospace characters that fit in the +// available width. Safe to call without holding mu. +func (e *TextViewEntry) charsPerLine() int { + if e.rend == nil { + return 80 + } + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + if charWidth <= 0 { + charWidth = theme.TextSize() * 0.6 + } + availWidth := e.rend.size.Width - tvPadX*2 - tvScrollW + if availWidth < 1 { + return 80 + } + n := int(availWidth / charWidth) + if n < 10 { + n = 10 + } + return n +} + +// buildVisualLines converts logical lines into wrapped visual lines. +// Each logical line is tokenised once for correct colours on continuations. +func buildVisualLines(logicalLines []string, charsPerLine int) []tveVisualLine { + if charsPerLine < 10 { + charsPerLine = 10 + } + bodyStartLine := -1 + for i, l := range logicalLines { + if l == "" && i > 0 { + bodyStartLine = i + break + } + } + contentType := "text" + if bodyStartLine >= 0 { + contentType = detectContentType(logicalLines) + } + + var result []tveVisualLine + + for logIdx, rawLine := range logicalLines { + inBody := bodyStartLine >= 0 && logIdx > bodyStartLine + runes := []rune(rawLine) + + if len(runes) == 0 { + var tokens []tvToken + if !inBody { + tokens = tokeniseHTTPMeta(rawLine, logIdx == 0) + } else { + tokens = tokeniseBody(rawLine, contentType) + } + result = append(result, tveVisualLine{ + tokens: tokens, raw: "", logicalIdx: logIdx, colOffset: 0, + }) + continue + } + + colOffset := 0 + for colOffset < len(runes) { + end := colOffset + charsPerLine + if end > len(runes) { + end = len(runes) + } + chunk := string(runes[colOffset:end]) + var chunkTokens []tvToken + if !inBody { + // Only the very first chunk of the first line is the HTTP start line. + chunkTokens = tokeniseHTTPMeta(chunk, logIdx == 0 && colOffset == 0) + } else { + chunkTokens = tokeniseBody(chunk, contentType) + } + result = append(result, tveVisualLine{ + tokens: chunkTokens, + raw: chunk, + logicalIdx: logIdx, + colOffset: colOffset, + }) + colOffset = end + } + } + return result +} + +// logicalToVisual maps a logical (line, col) to the visual line index and +// column within that visual line. Handles end-of-line and wrapped lines correctly. +func logicalToVisual(visual []tveVisualLine, logicalLine, logicalCol int) (int, int) { + lastMatchV := -1 + for v, vl := range visual { + if vl.logicalIdx != logicalLine { + continue + } + chunkLen := len([]rune(vl.raw)) + chunkEnd := vl.colOffset + chunkLen + isLastChunk := v+1 >= len(visual) || visual[v+1].logicalIdx != logicalLine + + if logicalCol >= vl.colOffset && logicalCol < chunkEnd { + return v, logicalCol - vl.colOffset + } + if isLastChunk && logicalCol >= chunkEnd { + return v, chunkLen + } + lastMatchV = v + } + if lastMatchV >= 0 { + vl := visual[lastMatchV] + return lastMatchV, len([]rune(vl.raw)) + } + return 0, 0 +} + +func (e *TextViewEntry) pushUndo(prevRaw string, prevCursorLine, prevCursorCol int) { + if len(e.undoStack) >= 100 { + e.undoStack = e.undoStack[1:] + } + e.undoStack = append(e.undoStack, tveUndoEntry{ + raw: prevRaw, cursorLine: prevCursorLine, cursorCol: prevCursorCol, + }) +} + +func (e *TextViewEntry) deleteSelection(lines []string) []string { + if !e.hasSelection { + return lines + } + startLine, startCol, endLine, endCol := e.normaliseSelection() + if !e.hasSelection { + return lines + } + if startLine >= len(lines) { + startLine = len(lines) - 1 + } + if endLine >= len(lines) { + endLine = len(lines) - 1 + } + if startLine == endLine { + runes := []rune(lines[startLine]) + if startCol > len(runes) { + startCol = len(runes) + } + if endCol > len(runes) { + endCol = len(runes) + } + newRunes := append([]rune{}, runes[:startCol]...) + newRunes = append(newRunes, runes[endCol:]...) + lines[startLine] = string(newRunes) + } else { + startRunes := []rune(lines[startLine]) + endRunes := []rune(lines[endLine]) + if startCol > len(startRunes) { + startCol = len(startRunes) + } + if endCol > len(endRunes) { + endCol = len(endRunes) + } + merged := string(startRunes[:startCol]) + string(endRunes[endCol:]) + newLines := make([]string, 0, len(lines)-(endLine-startLine)) + newLines = append(newLines, lines[:startLine]...) + newLines = append(newLines, merged) + newLines = append(newLines, lines[endLine+1:]...) + lines = newLines + } + e.cursorLine = startLine + e.cursorCol = startCol + e.hasSelection = false + e.shiftSelecting = false + return lines +} + +func (e *TextViewEntry) ensureCursorVisible() { + visible := e.visibleLineCount() + if visible <= 0 { + return + } + vLine, _ := logicalToVisual(e.visual(), e.cursorLine, e.cursorCol) + if vLine < e.scrollOffset { + e.scrollOffset = vLine + } else if vLine >= e.scrollOffset+visible { + e.scrollOffset = vLine - visible + 1 + } +} + +func (e *TextViewEntry) thumbBounds() (thumbY, thumbH float32, ok bool) { + if e.rend == nil { + return 0, 0, false + } + totalLines := len(e.visual()) + bodyHeight := e.rend.size.Height - tvPadY*2 + visibleLines := int(bodyHeight / tvLineH) + scrollable := totalLines - visibleLines + if scrollable <= 0 || bodyHeight <= 0 { + return 0, 0, false + } + ratio := float32(visibleLines) / float32(totalLines) + thumbH = bodyHeight * ratio + if thumbH < tvScrollMinThumb { + thumbH = tvScrollMinThumb + } + trackHeight := bodyHeight - thumbH + thumbY = tvPadY + (float32(e.scrollOffset)/float32(scrollable))*trackHeight + return thumbY, thumbH, true +} + +// -------------------------------------------------------------------------- +// Blink — single long-lived goroutine, reset via channel on each keystroke/click. +// -------------------------------------------------------------------------- + +func (e *TextViewEntry) startBlink() { + e.stopBlink() + resetCh := make(chan struct{}, 1) + e.blinkReset = resetCh + e.blinkRunning = true + e.mu.Lock() + e.cursorVisible = true + e.mu.Unlock() + go func() { + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + for { + select { + case _, ok := <-resetCh: + if !ok { + // Channel closed — widget destroyed or focus lost. + return + } + // Reset: make cursor visible and restart the ticker. + fyne.Do(func() { + e.mu.Lock() + e.cursorVisible = true + e.mu.Unlock() + }) + ticker.Reset(500 * time.Millisecond) + case <-ticker.C: + fyne.Do(func() { + e.mu.Lock() + e.cursorVisible = !e.cursorVisible + e.mu.Unlock() + e.Refresh() + }) + } + } + }() +} + +func (e *TextViewEntry) stopBlink() { + if e.blinkReset != nil { + close(e.blinkReset) + e.blinkReset = nil + } + e.blinkRunning = false + e.mu.Lock() + e.cursorVisible = false + e.mu.Unlock() +} + +// resetBlink makes the cursor immediately visible and resets the blink timer. +// Non-blocking — uses a buffered channel so rapid keystrokes never block. +func (e *TextViewEntry) resetBlink() { + if !e.blinkRunning { + return + } + select { + case e.blinkReset <- struct{}{}: + default: + // Already a reset pending — fine, cursor will be shown. + } +} + +// -------------------------------------------------------------------------- +// Widget lifecycle +// -------------------------------------------------------------------------- + +func (e *TextViewEntry) CreateRenderer() fyne.WidgetRenderer { + e.stopBlink() + r := newTVERenderer(e) + e.rend = r + return r +} + +func (e *TextViewEntry) Build() fyne.CanvasObject { + return container.NewStack(e, newTVEInputLayer(e)) +} + +// -------------------------------------------------------------------------- +// Input handling +// -------------------------------------------------------------------------- + +func (e *TextViewEntry) Scrolled(ev *fyne.ScrollEvent) { + totalLines := len(e.visual()) + visible := e.visibleLineCount() + maxOffset := totalLines - visible + if maxOffset < 0 { + maxOffset = 0 + } + e.scrollFrac -= ev.Scrolled.DY / tvLineH + rows := int(e.scrollFrac) + e.scrollFrac -= float32(rows) + if rows == 0 { + if ev.Scrolled.DY < 0 { + rows = 1 + } else if ev.Scrolled.DY > 0 { + rows = -1 + } + } + e.scrollOffset += rows + if e.scrollOffset < 0 { + e.scrollOffset = 0 + } + if e.scrollOffset > maxOffset { + e.scrollOffset = maxOffset + } + e.Refresh() +} + +func (e *TextViewEntry) visibleLineCount() int { + if e.rend == nil { + return 0 + } + availHeight := e.rend.size.Height - tvPadY*2 + if availHeight < 0 { + return 0 + } + return int(availHeight/tvLineH) + 1 +} + +func (e *TextViewEntry) posFromPoint(point fyne.Position) (line, col int) { + visual := e.visual() + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + if charWidth <= 0 { + charWidth = theme.TextSize() * 0.6 + } + vLineIdx := e.scrollOffset + int((point.Y-tvPadY)/tvLineH) + if vLineIdx < 0 { + vLineIdx = 0 + } + if vLineIdx >= len(visual) { + vLineIdx = len(visual) - 1 + } + if vLineIdx < 0 { + return 0, 0 + } + vl := visual[vLineIdx] + runes := []rune(vl.raw) + vCol := int((point.X - tvPadX) / charWidth) + if vCol < 0 { + vCol = 0 + } + if vCol > len(runes) { + vCol = len(runes) + } + return vl.logicalIdx, vl.colOffset + vCol +} + +func (e *TextViewEntry) normaliseSelection() (int, int, int, int) { + al, ac := e.selAnchorLine, e.selAnchorCol + cl, cc := e.selLine, e.selCol + if al > cl || (al == cl && ac > cc) { + return cl, cc, al, ac + } + return al, ac, cl, cc +} + +func (e *TextViewEntry) selectedText() string { + if !e.hasSelection { + return "" + } + lines := e.logicalLines() + startLine, startCol, endLine, endCol := e.normaliseSelection() + if startLine == endLine { + runes := []rune(lines[startLine]) + if startCol > len(runes) { + startCol = len(runes) + } + if endCol > len(runes) { + endCol = len(runes) + } + return string(runes[startCol:endCol]) + } + var b strings.Builder + for i := startLine; i <= endLine; i++ { + if i >= len(lines) { + break + } + runes := []rune(lines[i]) + lo, hi := 0, len(runes) + if i == startLine { + lo = startCol + } + if i == endLine { + hi = endCol + if hi > len(runes) { + hi = len(runes) + } + } + if lo > len(runes) { + lo = len(runes) + } + b.WriteString(string(runes[lo:hi])) + if i < endLine { + b.WriteByte('\n') + } + } + return b.String() +} + +// Dragged is implemented directly on TextViewEntry so Fyne's hit-test finds it +// before any parent split container — same pattern as widget.Entry. +func (e *TextViewEntry) Dragged(ev *fyne.DragEvent) { + if e.rend != nil && ev.Position.X >= e.rend.size.Width-tvScrollW { + return + } + line, col := e.posFromPoint(ev.Position) + e.selLine, e.selCol = line, col + e.cursorLine, e.cursorCol = line, col + e.hasSelection = !(e.selAnchorLine == line && e.selAnchorCol == col) + e.Refresh() +} + +func (e *TextViewEntry) DragEnd() {} + +func (e *TextViewEntry) TypedRune(r rune) { + lines := e.logicalLines() + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + lines = e.deleteSelection(lines) + } + cl := e.cursorLine + cc := e.cursorCol + if cl >= len(lines) { + cl = len(lines) - 1 + } + runes := []rune(lines[cl]) + if cc > len(runes) { + cc = len(runes) + } + newRunes := make([]rune, 0, len(runes)+1) + newRunes = append(newRunes, runes[:cc]...) + newRunes = append(newRunes, r) + newRunes = append(newRunes, runes[cc:]...) + lines[cl] = string(newRunes) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.mu.Lock() + e.cursorLine = cl + e.cursorCol = cc + 1 + e.mu.Unlock() + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() +} + +func (e *TextViewEntry) TypedKey(ev *fyne.KeyEvent) { + lines := e.logicalLines() + cl := e.cursorLine + cc := e.cursorCol + if cl >= len(lines) { + cl = len(lines) - 1 + } + if cl < 0 { + cl = 0 + } + runes := []rune(lines[cl]) + if cc > len(runes) { + cc = len(runes) + } + + isShift := e.shiftHeld + + beginShift := func() { + if !e.shiftSelecting { + e.shiftSelecting = true + e.selAnchorLine = cl + e.selAnchorCol = cc + } + } + updateShift := func() { + e.selLine = e.cursorLine + e.selCol = e.cursorCol + e.hasSelection = !(e.selAnchorLine == e.selLine && e.selAnchorCol == e.selCol) + } + cancelShift := func() { + e.shiftSelecting = false + e.shiftHeld = false + e.hasSelection = false + } + + switch ev.Name { + case fyne.KeyBackspace: + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + lines = e.deleteSelection(lines) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + } else if cc > 0 { + newRunes := append([]rune{}, runes[:cc-1]...) + newRunes = append(newRunes, runes[cc:]...) + lines[cl] = string(newRunes) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.mu.Lock() + e.cursorCol = cc - 1 + e.mu.Unlock() + } else if cl > 0 { + prevRunes := []rune(lines[cl-1]) + merged := string(prevRunes) + string(runes) + newLines := append([]string{}, lines[:cl-1]...) + newLines = append(newLines, merged) + newLines = append(newLines, lines[cl+1:]...) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(newLines) + e.mu.Lock() + e.cursorLine = cl - 1 + e.cursorCol = len(prevRunes) + e.mu.Unlock() + } + cancelShift() + + case fyne.KeyDelete: + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + lines = e.deleteSelection(lines) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + } else if cc < len(runes) { + newRunes := append([]rune{}, runes[:cc]...) + newRunes = append(newRunes, runes[cc+1:]...) + lines[cl] = string(newRunes) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + } else if cl < len(lines)-1 { + merged := string(runes) + lines[cl+1] + newLines := append([]string{}, lines[:cl]...) + newLines = append(newLines, merged) + newLines = append(newLines, lines[cl+2:]...) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(newLines) + } + cancelShift() + + case fyne.KeyReturn, fyne.KeyEnter: + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + lines = e.deleteSelection(lines) + cl = e.cursorLine + cc = e.cursorCol + runes = []rune(lines[cl]) + } + before := string(runes[:cc]) + after := string(runes[cc:]) + newLines := make([]string, 0, len(lines)+1) + newLines = append(newLines, lines[:cl]...) + newLines = append(newLines, before, after) + newLines = append(newLines, lines[cl+1:]...) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(newLines) + e.mu.Lock() + e.cursorLine = cl + 1 + e.cursorCol = 0 + e.mu.Unlock() + cancelShift() + + case fyne.KeyTab: + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + startLine, selStartCol, endLine, selEndCol := e.normaliseSelection() + indent := strings.Repeat(" ", tveTabWidth) + for i := startLine; i <= endLine; i++ { + lines[i] = indent + lines[i] + } + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.selAnchorLine = startLine + if selStartCol > 0 { + e.selAnchorCol = selStartCol + tveTabWidth + } else { + e.selAnchorCol = 0 + } + e.selLine = endLine + e.selCol = selEndCol + tveTabWidth + // Preserve cursor at whichever end it was before. + if prevCL == startLine { + e.cursorLine = startLine + e.cursorCol = prevCC + tveTabWidth + } else { + e.cursorLine = endLine + e.cursorCol = selEndCol + tveTabWidth + } + } else { + indent := strings.Repeat(" ", tveTabWidth) + newRunes := make([]rune, 0, len(runes)+tveTabWidth) + newRunes = append(newRunes, runes[:cc]...) + newRunes = append(newRunes, []rune(indent)...) + newRunes = append(newRunes, runes[cc:]...) + lines[cl] = string(newRunes) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.mu.Lock() + e.cursorCol = cc + tveTabWidth + e.mu.Unlock() + } + + case fyne.KeyLeft: + if isShift { + beginShift() + } else if e.hasSelection { + startLine, startCol, _, _ := e.normaliseSelection() + e.mu.Lock() + e.cursorLine = startLine + e.cursorCol = startCol + e.mu.Unlock() + cancelShift() + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() + return + } + if cc > 0 { + e.mu.Lock() + e.cursorCol-- + e.mu.Unlock() + } else if cl > 0 { + e.mu.Lock() + e.cursorLine-- + e.cursorCol = len([]rune(lines[e.cursorLine])) + e.mu.Unlock() + } + if isShift { + updateShift() + } + + case fyne.KeyRight: + if isShift { + beginShift() + } else if e.hasSelection { + _, _, endLine, endCol := e.normaliseSelection() + e.mu.Lock() + e.cursorLine = endLine + e.cursorCol = endCol + e.mu.Unlock() + cancelShift() + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() + return + } + if cc < len(runes) { + e.mu.Lock() + e.cursorCol++ + e.mu.Unlock() + } else if cl < len(lines)-1 { + e.mu.Lock() + e.cursorLine++ + e.cursorCol = 0 + e.mu.Unlock() + } + if isShift { + updateShift() + } + + case fyne.KeyUp: + if isShift { + beginShift() + } else { + cancelShift() + } + if cl > 0 { + e.mu.Lock() + e.cursorLine-- + prevLen := len([]rune(lines[e.cursorLine])) + if e.cursorCol > prevLen { + e.cursorCol = prevLen + } + e.mu.Unlock() + } + if isShift { + updateShift() + } + + case fyne.KeyDown: + if isShift { + beginShift() + } else { + cancelShift() + } + if cl < len(lines)-1 { + e.mu.Lock() + e.cursorLine++ + nextLen := len([]rune(lines[e.cursorLine])) + if e.cursorCol > nextLen { + e.cursorCol = nextLen + } + e.mu.Unlock() + } + if isShift { + updateShift() + } + + case fyne.KeyHome: + if isShift { + beginShift() + } else { + cancelShift() + } + e.mu.Lock() + e.cursorCol = 0 + e.mu.Unlock() + if isShift { + updateShift() + } + + case fyne.KeyEnd: + if isShift { + beginShift() + } else { + cancelShift() + } + e.mu.Lock() + e.cursorCol = len(runes) + e.mu.Unlock() + if isShift { + updateShift() + } + + case fyne.KeyPageUp: + visible := e.visibleLineCount() + visual := e.visual() + curVLine, _ := logicalToVisual(visual, e.cursorLine, e.cursorCol) + targetVLine := curVLine - visible + if targetVLine < 0 { + targetVLine = 0 + } + vl := visual[targetVLine] + e.mu.Lock() + e.cursorLine = vl.logicalIdx + e.cursorCol = vl.colOffset + e.mu.Unlock() + cancelShift() + + case fyne.KeyPageDown: + visible := e.visibleLineCount() + visual := e.visual() + curVLine, _ := logicalToVisual(visual, e.cursorLine, e.cursorCol) + targetVLine := curVLine + visible + if targetVLine >= len(visual) { + targetVLine = len(visual) - 1 + } + vl := visual[targetVLine] + e.mu.Lock() + e.cursorLine = vl.logicalIdx + e.cursorCol = vl.colOffset + e.mu.Unlock() + cancelShift() + } + + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() +} + +func (e *TextViewEntry) TypedShortcut(s fyne.Shortcut) { + if cs, ok := s.(*desktop.CustomShortcut); ok { + if cs.KeyName == fyne.KeyS && cs.Modifier == fyne.KeyModifierControl { + if e.OnSend != nil { + e.OnSend() + } + return + } + } + + switch s.(type) { + case *fyne.ShortcutUndo: + e.mu.Lock() + if len(e.undoStack) > 0 { + undo := e.undoStack[len(e.undoStack)-1] + e.undoStack = e.undoStack[:len(e.undoStack)-1] + e.rawFull = undo.raw + e.visualCacheValid = false + e.cursorLine = undo.cursorLine + e.cursorCol = undo.cursorCol + e.hasSelection = false + e.shiftSelecting = false + e.shiftHeld = false + } + e.mu.Unlock() + e.resetBlink() + e.Refresh() + + case *fyne.ShortcutSelectAll: + lines := e.logicalLines() + totalLines := len(lines) + e.selAnchorLine, e.selAnchorCol = 0, 0 + e.selLine = totalLines - 1 + if e.selLine < 0 { + e.selLine = 0 + } + e.selCol = len([]rune(lines[e.selLine])) + e.hasSelection = true + e.Refresh() + + case *fyne.ShortcutCopy: + if e.win != nil && e.hasSelection { + e.win.Clipboard().SetContent(e.selectedText()) + } + + case *fyne.ShortcutCut: + if e.win == nil || !e.hasSelection { + return + } + e.win.Clipboard().SetContent(e.selectedText()) + lines := e.logicalLines() + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + lines = e.deleteSelection(lines) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() + + case *fyne.ShortcutPaste: + if e.win == nil { + return + } + text := e.win.Clipboard().Content() + if text == "" { + return + } + text = strings.ReplaceAll(text, "\r\n", "\n") + text = strings.ReplaceAll(text, "\r", "\n") + lines := e.logicalLines() + prevRaw := strings.Join(lines, "\n") + prevCL, prevCC := e.cursorLine, e.cursorCol + if e.hasSelection { + lines = e.deleteSelection(lines) + } + cl := e.cursorLine + cc := e.cursorCol + if cl >= len(lines) { + cl = len(lines) - 1 + } + runes := []rune(lines[cl]) + if cc > len(runes) { + cc = len(runes) + } + pasteLines := strings.Split(text, "\n") + if len(pasteLines) == 1 { + newRunes := make([]rune, 0, len(runes)+len([]rune(text))) + newRunes = append(newRunes, runes[:cc]...) + newRunes = append(newRunes, []rune(text)...) + newRunes = append(newRunes, runes[cc:]...) + lines[cl] = string(newRunes) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(lines) + e.mu.Lock() + e.cursorLine = cl + e.cursorCol = cc + len([]rune(pasteLines[0])) + e.mu.Unlock() + } else { + firstPart := string(runes[:cc]) + pasteLines[0] + lastPart := pasteLines[len(pasteLines)-1] + string(runes[cc:]) + newLines := make([]string, 0, len(lines)+len(pasteLines)-1) + newLines = append(newLines, lines[:cl]...) + newLines = append(newLines, firstPart) + newLines = append(newLines, pasteLines[1:len(pasteLines)-1]...) + newLines = append(newLines, lastPart) + newLines = append(newLines, lines[cl+1:]...) + e.pushUndo(prevRaw, prevCL, prevCC) + e.commitLines(newLines) + e.mu.Lock() + e.cursorLine = cl + len(pasteLines) - 1 + e.cursorCol = len([]rune(pasteLines[len(pasteLines)-1])) + e.mu.Unlock() + } + e.ensureCursorVisible() + e.resetBlink() + e.Refresh() + } +} + +func (e *TextViewEntry) FocusGained() { + e.startBlink() + e.Refresh() +} + +func (e *TextViewEntry) FocusLost() { + e.stopBlink() + e.shiftHeld = false + // Only clear selection state if we are not mid-drag. When dragging outside + // the window Fyne fires FocusLost but the user is still interacting. + if !e.dragging { + e.shiftSelecting = false + } + e.Refresh() +} + +func (e *TextViewEntry) KeyDown(ev *fyne.KeyEvent) { + if ev.Name == desktop.KeyShiftLeft || ev.Name == desktop.KeyShiftRight { + e.shiftHeld = true + } +} + +func (e *TextViewEntry) KeyUp(ev *fyne.KeyEvent) { + if ev.Name == desktop.KeyShiftLeft || ev.Name == desktop.KeyShiftRight { + e.shiftHeld = false + e.shiftSelecting = false + } +} + +// -------------------------------------------------------------------------- +// Renderer +// -------------------------------------------------------------------------- + +type tveRenderer struct { + entry *TextViewEntry + size fyne.Size + + selBgs []*canvas.Rectangle + tokenSlots [][]*canvas.Text + cursorRect *canvas.Rectangle + slots int + + scrollTrack *canvas.Rectangle + scrollThumb *canvas.Rectangle + + lastWidth float32 // tracks width to invalidate visual cache on resize +} + +func newTVERenderer(entry *TextViewEntry) *tveRenderer { + r := &tveRenderer{entry: entry} + r.scrollTrack = canvas.NewRectangle(color.Transparent) + r.scrollThumb = canvas.NewRectangle(color.Transparent) + r.cursorRect = canvas.NewRectangle(color.Transparent) + return r +} + +// Destroy is called when the widget is removed. Stop the blink goroutine. +func (r *tveRenderer) Destroy() { + r.entry.stopBlink() +} + +func (r *tveRenderer) Objects() []fyne.CanvasObject { + total := 3 + len(r.selBgs) + len(r.tokenSlots)*tvMaxTokensPerLine + objs := make([]fyne.CanvasObject, 0, total) + objs = append(objs, r.scrollTrack, r.scrollThumb) + for _, selBg := range r.selBgs { + objs = append(objs, selBg) + } + for _, slot := range r.tokenSlots { + for _, t := range slot { + objs = append(objs, t) + } + } + objs = append(objs, r.cursorRect) + return objs +} + +func (r *tveRenderer) MinSize() fyne.Size { + return fyne.NewSize(200, tvLineH*3+tvPadY*2) +} + +func (r *tveRenderer) Layout(size fyne.Size) { + r.size = size + r.checkWrapWidth() + r.growSlots() + r.layoutContent() +} + +func (r *tveRenderer) Refresh() { + r.checkWrapWidth() + r.growSlots() + r.layoutContent() + canvas.Refresh(r.entry) +} + +func (r *tveRenderer) checkWrapWidth() { + availWidth := r.size.Width - tvPadX*2 - tvScrollW + if availWidth < 1 { + return + } + if r.lastWidth != availWidth { + r.lastWidth = availWidth + // Invalidate visual cache so wrapping is recomputed at new width. + r.entry.mu.Lock() + r.entry.visualCacheRaw = "" + r.entry.mu.Unlock() + } +} + +func (r *tveRenderer) growSlots() { + bodyHeight := r.size.Height - tvPadY*2 + if bodyHeight < 0 { + bodyHeight = 0 + } + needed := int(bodyHeight/tvLineH) + 2 + for r.slots < needed { + selBg := canvas.NewRectangle(color.Transparent) + r.selBgs = append(r.selBgs, selBg) + tokenSlot := make([]*canvas.Text, tvMaxTokensPerLine) + for i := range tokenSlot { + t := canvas.NewText("", color.White) + t.TextStyle = fyne.TextStyle{Monospace: true} + t.TextSize = theme.TextSize() + tokenSlot[i] = t + } + r.tokenSlots = append(r.tokenSlots, tokenSlot) + r.slots++ + } + for slot := needed; slot < r.slots; slot++ { + r.selBgs[slot].Hide() + for _, t := range r.tokenSlots[slot] { + t.Hide() + } + } +} + +func (r *tveRenderer) visibleLines() int { + bodyHeight := r.size.Height - tvPadY*2 + if bodyHeight < 0 { + return 0 + } + return int(bodyHeight/tvLineH) + 1 +} + +func (r *tveRenderer) layoutContent() { + entry := r.entry + visual := entry.visual() + totalLines := len(visual) + + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + if charWidth <= 0 { + charWidth = theme.TextSize() * 0.6 + } + + entry.mu.Lock() + scrollOffset := entry.scrollOffset + hasSelection := entry.hasSelection + cursorLogLine := entry.cursorLine + cursorLogCol := entry.cursorCol + cursorVisible := entry.cursorVisible + var selStartLog, selStartCol, selEndLog, selEndCol int + if hasSelection { + selStartLog, selStartCol, selEndLog, selEndCol = entry.normaliseSelection() + } + entry.mu.Unlock() + + visible := r.visibleLines() + maxOffset := totalLines - visible + if maxOffset < 0 { + maxOffset = 0 + } + if scrollOffset > maxOffset { + scrollOffset = maxOffset + entry.mu.Lock() + entry.scrollOffset = scrollOffset + entry.mu.Unlock() + } + + needed := int((r.size.Height-tvPadY*2)/tvLineH) + 2 + cursorVLine, cursorVCol := logicalToVisual(visual, cursorLogLine, cursorLogCol) + cursorDrawn := false + yPos := tvPadY + + for slot := 0; slot < needed && slot < r.slots; slot++ { + lineIdx := scrollOffset + slot + selBg := r.selBgs[slot] + + if lineIdx >= totalLines { + selBg.FillColor = color.Transparent + selBg.Resize(fyne.NewSize(0, 0)) + selBg.Refresh() + selBg.Hide() + for _, t := range r.tokenSlots[slot] { + t.Hide() + } + yPos += tvLineH + continue + } + + vl := visual[lineIdx] + + // Selection highlight. Fix 5: use strict < chunkEnd for start boundary. + if hasSelection { + vSelStart, vSelEnd := -1, -1 + chunkStart := vl.colOffset + chunkEnd := vl.colOffset + len([]rune(vl.raw)) + isLastChunk := lineIdx+1 >= len(visual) || visual[lineIdx+1].logicalIdx != vl.logicalIdx + + switch { + case vl.logicalIdx > selStartLog && vl.logicalIdx < selEndLog: + vSelStart = 0 + vSelEnd = len([]rune(vl.raw)) + case vl.logicalIdx == selStartLog && vl.logicalIdx == selEndLog: + if selStartCol < chunkEnd && selEndCol > chunkStart { + vSelStart = selStartCol - chunkStart + vSelEnd = selEndCol - chunkStart + } + case vl.logicalIdx == selStartLog: + if selStartCol < chunkEnd { + vSelStart = selStartCol - chunkStart + vSelEnd = len([]rune(vl.raw)) + // Extend one extra char width for visual continuity if not last chunk. + if !isLastChunk { + vSelEnd++ + } + } + case vl.logicalIdx == selEndLog: + if selEndCol > chunkStart { + vSelStart = 0 + vSelEnd = selEndCol - chunkStart + } + } + + if vSelStart < 0 { + vSelStart = 0 + } + lineLen := len([]rune(vl.raw)) + if vSelEnd > lineLen { + vSelEnd = lineLen + } + + if vSelStart >= 0 && vSelEnd > vSelStart { + selBg.Move(fyne.NewPos(tvPadX+float32(vSelStart)*charWidth, yPos)) + selBg.Resize(fyne.NewSize(float32(vSelEnd-vSelStart)*charWidth, tvLineH)) + selBg.FillColor = selectionColor() + selBg.Refresh() + selBg.Show() + } else { + selBg.FillColor = color.Transparent + selBg.Resize(fyne.NewSize(0, 0)) + selBg.Refresh() + selBg.Hide() + } + } else { + selBg.FillColor = color.Transparent + selBg.Resize(fyne.NewSize(0, 0)) + selBg.Refresh() + selBg.Hide() + } + + // Draw tokens. + xPos := tvPadX + tokenIdx := 0 + for _, token := range vl.tokens { + if tokenIdx >= tvMaxTokensPerLine { + break + } + if token.Text == "" { + continue + } + t := r.tokenSlots[slot][tokenIdx] + t.Text = token.Text + t.Color = tvColor(token.Kind) + t.TextStyle = fyne.TextStyle{Monospace: true} + t.TextSize = theme.TextSize() + t.Move(fyne.NewPos(xPos, yPos)) + t.Refresh() + t.Show() + xPos += fyne.MeasureText(token.Text, theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + tokenIdx++ + } + for ; tokenIdx < tvMaxTokensPerLine; tokenIdx++ { + r.tokenSlots[slot][tokenIdx].Hide() + } + + // Draw cursor. + if lineIdx == cursorVLine && cursorVisible { + cursorX := tvPadX + float32(cursorVCol)*charWidth + r.cursorRect.Move(fyne.NewPos(cursorX, yPos+2)) + r.cursorRect.Resize(fyne.NewSize(2, tvLineH-4)) + r.cursorRect.FillColor = theme.Color(theme.ColorNameForeground) + r.cursorRect.Refresh() + r.cursorRect.Show() + cursorDrawn = true + } + + yPos += tvLineH + } + + if !cursorDrawn { + r.cursorRect.FillColor = color.Transparent + r.cursorRect.Resize(fyne.NewSize(0, 0)) + r.cursorRect.Refresh() + r.cursorRect.Hide() + } + + r.layoutScrollbar(totalLines) +} + +func (r *tveRenderer) layoutScrollbar(totalLines int) { + bodyHeight := r.size.Height - tvPadY*2 + trackW := float32(tvScrollW) + trackX := r.size.Width - trackW + if r.entry.scrollbarHovered { + trackW = tvScrollW + 6 + trackX = r.size.Width - trackW + } + r.scrollTrack.Move(fyne.NewPos(trackX, tvPadY)) + r.scrollTrack.Resize(fyne.NewSize(trackW, bodyHeight)) + + visibleLines := int(bodyHeight / tvLineH) + if totalLines <= visibleLines { + r.scrollTrack.FillColor = color.Transparent + r.scrollTrack.Refresh() + r.scrollThumb.FillColor = color.Transparent + r.scrollThumb.Refresh() + return + } + + r.scrollTrack.FillColor = theme.Color(theme.ColorNameScrollBar) + r.scrollTrack.Refresh() + + ratio := float32(visibleLines) / float32(totalLines) + thumbHeight := bodyHeight * ratio + if thumbHeight < tvScrollMinThumb { + thumbHeight = tvScrollMinThumb + } + scrollable := float32(totalLines - visibleLines) + thumbY := tvPadY + (float32(r.entry.scrollOffset)/scrollable)*(bodyHeight-thumbHeight) + + thumbW := float32(tvScrollW - 2) + thumbX := trackX + 1 + if r.entry.scrollbarHovered { + thumbW = trackW - 2 + thumbX = trackX + 1 + } + r.scrollThumb.Move(fyne.NewPos(thumbX, thumbY)) + r.scrollThumb.Resize(fyne.NewSize(thumbW, thumbHeight)) + r.scrollThumb.FillColor = theme.Color(theme.ColorNameForeground) + r.scrollThumb.CornerRadius = thumbW / 2 + r.scrollThumb.Refresh() +} + +// -------------------------------------------------------------------------- +// Input layer +// -------------------------------------------------------------------------- + +type tveInputLayer struct { + widget.BaseWidget + entry *TextViewEntry + draggingThumb bool + draggingText bool +} + +func newTVEInputLayer(entry *TextViewEntry) *tveInputLayer { + layer := &tveInputLayer{entry: entry} + layer.ExtendBaseWidget(layer) + return layer +} + +func (layer *tveInputLayer) CreateRenderer() fyne.WidgetRenderer { + return widget.NewSimpleRenderer(canvas.NewRectangle(color.Transparent)) +} + +func (layer *tveInputLayer) Scrolled(ev *fyne.ScrollEvent) { layer.entry.Scrolled(ev) } + +func (layer *tveInputLayer) inScrollbarArea(xPos float32) bool { + if layer.entry.rend == nil { + return false + } + return xPos >= layer.entry.rend.size.Width-tvScrollW +} + +func (layer *tveInputLayer) inThumb(xPos, yPos float32) bool { + if !layer.inScrollbarArea(xPos) { + return false + } + thumbY, thumbH, ok := layer.entry.thumbBounds() + if !ok { + return false + } + return yPos >= thumbY && yPos <= thumbY+thumbH +} + +func (layer *tveInputLayer) MouseDown(ev *desktop.MouseEvent) { + if ev.Button != desktop.MouseButtonPrimary { + return + } + layer.draggingThumb = layer.inThumb(ev.Position.X, ev.Position.Y) + if layer.draggingThumb { + return + } + if layer.inScrollbarArea(ev.Position.X) { + return + } + entry := layer.entry + line, col := entry.posFromPoint(ev.Position) + entry.selAnchorLine, entry.selAnchorCol = line, col + entry.selLine, entry.selCol = line, col + entry.cursorLine, entry.cursorCol = line, col + entry.hasSelection = false + entry.shiftSelecting = false + entry.shiftHeld = false + entry.dragging = false + entry.resetBlink() + entry.Refresh() + if entry.win != nil { + entry.win.Canvas().Focus(entry) + } +} + +func (layer *tveInputLayer) MouseUp(_ *desktop.MouseEvent) { + layer.draggingThumb = false + layer.draggingText = false + layer.entry.dragging = false +} + +func (layer *tveInputLayer) Tapped(_ *fyne.PointEvent) {} + +func (layer *tveInputLayer) DoubleTapped(ev *fyne.PointEvent) { + if layer.inScrollbarArea(ev.Position.X) { + return + } + entry := layer.entry + line, col := entry.posFromPoint(ev.Position) + lines := entry.logicalLines() + if line >= len(lines) { + return + } + runes := []rune(lines[line]) + if len(runes) == 0 { + return + } + if col >= len(runes) { + col = len(runes) - 1 + } + wordStart := col + for wordStart > 0 && tveIsWordRune(runes[wordStart-1]) { + wordStart-- + } + wordEnd := col + for wordEnd < len(runes) && tveIsWordRune(runes[wordEnd]) { + wordEnd++ + } + entry.selAnchorLine, entry.selAnchorCol = line, wordStart + entry.selLine, entry.selCol = line, wordEnd + entry.cursorLine, entry.cursorCol = line, wordEnd + entry.hasSelection = wordStart != wordEnd + entry.Refresh() +} + +func (layer *tveInputLayer) dragScrollbar(ev *fyne.DragEvent) { + entry := layer.entry + if entry.rend == nil { + return + } + totalLines := len(entry.visual()) + if totalLines == 0 { + return + } + bodyHeight := entry.rend.size.Height - tvPadY*2 + visibleLines := int(bodyHeight / tvLineH) + scrollable := totalLines - visibleLines + if scrollable <= 0 { + return + } + ratio := float32(visibleLines) / float32(totalLines) + thumbHeight := bodyHeight * ratio + if thumbHeight < tvScrollMinThumb { + thumbHeight = tvScrollMinThumb + } + trackHeight := bodyHeight - thumbHeight + if trackHeight <= 0 { + return + } + newThumbTop := ev.Position.Y - tvPadY - thumbHeight/2 + if newThumbTop < 0 { + newThumbTop = 0 + } + if newThumbTop > trackHeight { + newThumbTop = trackHeight + } + entry.scrollOffset = int((newThumbTop / trackHeight) * float32(scrollable)) + entry.Refresh() +} + +func (layer *tveInputLayer) Dragged(ev *fyne.DragEvent) { + if layer.draggingThumb { + layer.dragScrollbar(ev) + return + } + layer.draggingText = true + layer.entry.dragging = true + layer.entry.Dragged(ev) +} + +func (layer *tveInputLayer) DragEnd() { + wasDraggingText := layer.draggingText + layer.draggingThumb = false + layer.draggingText = false + layer.entry.dragging = false + if wasDraggingText && layer.entry.win != nil { + layer.entry.win.Canvas().Focus(layer.entry) + layer.entry.resetBlink() + } +} + +func (layer *tveInputLayer) MouseIn(_ *desktop.MouseEvent) {} + +func (layer *tveInputLayer) MouseOut() { + if layer.entry.scrollbarHovered { + layer.entry.scrollbarHovered = false + layer.entry.Refresh() + } +} + +func (layer *tveInputLayer) MouseMoved(ev *desktop.MouseEvent) { + inScrollbar := layer.inScrollbarArea(ev.Position.X) + if inScrollbar != layer.entry.scrollbarHovered { + layer.entry.scrollbarHovered = inScrollbar + layer.entry.Refresh() + } +} + +func tveIsWordRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '-' +} diff --git a/internal/ui/widgets/textviewentry_test.go b/internal/ui/widgets/textviewentry_test.go new file mode 100644 index 0000000..bbb20ef --- /dev/null +++ b/internal/ui/widgets/textviewentry_test.go @@ -0,0 +1,381 @@ +package widgets + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTVE() *TextViewEntry { + e := &TextViewEntry{} + e.ExtendBaseWidget(e) + return e +} + +// ── splitRaw ───────────────────────────────────────────────────────────────── + +func TestSplitRaw(t *testing.T) { + tests := []struct { + input string + want []string + }{ + {"", []string{""}}, + {"a\nb", []string{"a", "b"}}, + {"a\r\nb", []string{"a", "b"}}, + {"a\rb", []string{"a", "b"}}, + {"a\nb\nc", []string{"a", "b", "c"}}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := splitRaw(tt.input) + assert.Equal(t, tt.want, got) + }) + } +} + +// ── buildVisualLines ───────────────────────────────────────────────────────── + +func TestBuildVisualLines_NoWrap(t *testing.T) { + lines := buildVisualLines([]string{"GET /foo HTTP/1.1", "Host: example.com"}, 80) + assert.Len(t, lines, 2) +} + +func TestBuildVisualLines_Wraps(t *testing.T) { + long := strings.Repeat("x", 100) + lines := buildVisualLines([]string{long}, 40) + + assert.Greater(t, len(lines), 2, "expected ≥3 visual lines for 100-char line at width 40") + + for _, vl := range lines { + assert.Equal(t, 0, vl.logicalIdx, "all chunks should map to logicalIdx 0") + } + + expected := 0 + for _, vl := range lines { + assert.Equal(t, expected, vl.colOffset) + expected += len([]rune(vl.raw)) + } +} + +func TestBuildVisualLines_EmptyLine(t *testing.T) { + lines := buildVisualLines([]string{"GET /foo HTTP/1.1", "", `{"key":"val"}`}, 80) + assert.Len(t, lines, 3) +} + +func TestBuildVisualLines_EmptyInput(t *testing.T) { + lines := buildVisualLines([]string{""}, 80) + assert.Len(t, lines, 1) +} + +// ── logicalToVisual ────────────────────────────────────────────────────────── + +func TestLogicalToVisual_Simple(t *testing.T) { + visual := buildVisualLines([]string{"hello", "world"}, 80) + + v, col := logicalToVisual(visual, 0, 0) + assert.Equal(t, 0, v) + assert.Equal(t, 0, col) + + v, col = logicalToVisual(visual, 1, 3) + assert.Equal(t, 1, visual[v].logicalIdx) + assert.Equal(t, 3, col) +} + +func TestLogicalToVisual_EndOfLine(t *testing.T) { + visual := buildVisualLines([]string{"hello"}, 80) + + v, col := logicalToVisual(visual, 0, 5) + assert.Equal(t, 0, v) + assert.Equal(t, 5, col) +} + +func TestLogicalToVisual_WrappedLine(t *testing.T) { + // "01234567890123456789" wrapped at 10 → chunks [0:10] [10:20] + visual := buildVisualLines([]string{"01234567890123456789"}, 10) + + v, col := logicalToVisual(visual, 0, 11) + assert.Equal(t, 10, visual[v].colOffset, "col 11 should land in chunk starting at offset 10") + assert.Equal(t, 1, col) +} + +func TestLogicalToVisual_EndOfWrappedLine(t *testing.T) { + // "01234567890123456789" wrapped at 10 → chunks [0:10] [10:20] + visual := buildVisualLines([]string{"01234567890123456789"}, 10) + + v, _ := logicalToVisual(visual, 0, 20) + assert.Equal(t, visual[len(visual)-1].colOffset, visual[v].colOffset, "col at end should land in last chunk") +} + +// ── normaliseSelection ─────────────────────────────────────────────────────── + +func TestNormaliseSelection_Forward(t *testing.T) { + e := newTVE() + e.selAnchorLine, e.selAnchorCol = 0, 2 + e.selLine, e.selCol = 1, 5 + + sl, sc, el, ec := e.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 2, sc) + assert.Equal(t, 1, el) + assert.Equal(t, 5, ec) +} + +func TestNormaliseSelection_Reversed(t *testing.T) { + e := newTVE() + e.selAnchorLine, e.selAnchorCol = 1, 5 + e.selLine, e.selCol = 0, 2 + + sl, sc, el, ec := e.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 2, sc) + assert.Equal(t, 1, el) + assert.Equal(t, 5, ec) +} + +func TestNormaliseSelection_SamePosition(t *testing.T) { + e := newTVE() + e.selAnchorLine, e.selAnchorCol = 0, 3 + e.selLine, e.selCol = 0, 3 + + sl, sc, el, ec := e.normaliseSelection() + assert.Equal(t, 0, sl) + assert.Equal(t, 3, sc) + assert.Equal(t, 0, el) + assert.Equal(t, 3, ec) +} + +// ── selectedText ───────────────────────────────────────────────────────────── + +func TestSelectedText_SameLine(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello world" + e.mu.Unlock() + e.selAnchorLine, e.selAnchorCol = 0, 0 + e.selLine, e.selCol = 0, 5 + e.hasSelection = true + + assert.Equal(t, "hello", e.selectedText()) +} + +func TestSelectedText_MultiLine(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "line one\nline two" + e.mu.Unlock() + e.selAnchorLine, e.selAnchorCol = 0, 5 + e.selLine, e.selCol = 1, 4 + e.hasSelection = true + + assert.Equal(t, "one\nline", e.selectedText()) +} + +func TestSelectedText_NoSelection(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello" + e.mu.Unlock() + e.hasSelection = false + + assert.Equal(t, "", e.selectedText()) +} + +// ── deleteSelection ────────────────────────────────────────────────────────── + +func TestDeleteSelection_SameLine(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello world" + e.mu.Unlock() + e.selAnchorLine, e.selAnchorCol = 0, 6 + e.selLine, e.selCol = 0, 11 + e.hasSelection = true + + lines := e.logicalLines() + lines = e.deleteSelection(lines) + e.commitLines(lines) + + assert.Equal(t, "hello ", e.GetText()) + assert.Equal(t, 6, e.cursorCol) + assert.False(t, e.hasSelection) +} + +func TestDeleteSelection_MultiLine(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "line one\nline two\nline three" + e.mu.Unlock() + e.selAnchorLine, e.selAnchorCol = 0, 5 + e.selLine, e.selCol = 2, 4 + e.hasSelection = true + + lines := e.logicalLines() + lines = e.deleteSelection(lines) + e.commitLines(lines) + + assert.Equal(t, "line three", e.GetText()) +} + +func TestDeleteSelection_ReversedAnchor(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello world" + e.mu.Unlock() + e.selAnchorLine, e.selAnchorCol = 0, 11 + e.selLine, e.selCol = 0, 6 + e.hasSelection = true + + lines := e.logicalLines() + lines = e.deleteSelection(lines) + e.commitLines(lines) + + assert.Equal(t, "hello ", e.GetText()) +} + +func TestDeleteSelection_NoSelection(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello" + e.mu.Unlock() + e.hasSelection = false + + lines := e.logicalLines() + lines = e.deleteSelection(lines) + e.commitLines(lines) + + assert.Equal(t, "hello", e.GetText()) +} + +// ── pushUndo / undo ────────────────────────────────────────────────────────── + +func TestPushUndo_StoresCursorPosition(t *testing.T) { + e := newTVE() + e.pushUndo("hello", 0, 3) + + require.Len(t, e.undoStack, 1) + assert.Equal(t, "hello", e.undoStack[0].raw) + assert.Equal(t, 0, e.undoStack[0].cursorLine) + assert.Equal(t, 3, e.undoStack[0].cursorCol) +} + +func TestPushUndo_StackLimit(t *testing.T) { + e := newTVE() + for i := 0; i < 110; i++ { + e.pushUndo("x", 0, 0) + } + assert.LessOrEqual(t, len(e.undoStack), 100) +} + +func TestUndo_RestoresTextAndCursor(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello" + e.mu.Unlock() + e.cursorLine, e.cursorCol = 0, 5 + + e.pushUndo("hello", 0, 5) + + e.mu.Lock() + e.rawFull = "hello world" + e.visualCacheRaw = "" + e.cursorLine = 0 + e.cursorCol = 11 + e.mu.Unlock() + + e.mu.Lock() + undo := e.undoStack[len(e.undoStack)-1] + e.undoStack = e.undoStack[:len(e.undoStack)-1] + e.rawFull = undo.raw + e.visualCacheRaw = "" + e.cursorLine = undo.cursorLine + e.cursorCol = undo.cursorCol + e.mu.Unlock() + + assert.Equal(t, "hello", e.GetText()) + assert.Equal(t, 0, e.cursorLine) + assert.Equal(t, 5, e.cursorCol) +} + +// ── visual cache ───────────────────────────────────────────────────────────── + +func TestVisualCache_InvalidatedOnCommit(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello" + e.mu.Unlock() + + e.mu.Lock() + e.visualCache = buildVisualLines([]string{"hello"}, 80) + e.visualCacheRaw = "hello" + e.visualCacheWidth = 80 + e.mu.Unlock() + + e.commitLines([]string{"world"}) + + e.mu.Lock() + cacheRaw := e.visualCacheRaw + e.mu.Unlock() + + assert.Equal(t, "", cacheRaw, "commitLines should invalidate visual cache") +} + +func TestVisualCache_HitOnSecondCall(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "GET /foo HTTP/1.1\nHost: example.com" + e.visualCacheRaw = "" + e.mu.Unlock() + + // Prime the cache manually without needing a renderer. + cpl := 80 + e.mu.Lock() + built := buildVisualLines(splitRaw(e.rawFull), cpl) + e.visualCache = built + e.visualCacheRaw = e.rawFull + e.visualCacheWidth = cpl + e.mu.Unlock() + + e.mu.Lock() + c1 := e.visualCache + c2 := e.visualCache + e.mu.Unlock() + + require.NotEmpty(t, c1) + assert.Equal(t, &c1[0], &c2[0], "expected same slice on repeated access") +} + +func TestVisualCache_MissAfterMutation(t *testing.T) { + e := newTVE() + e.mu.Lock() + e.rawFull = "hello" + cpl := 80 + e.visualCache = buildVisualLines([]string{"hello"}, cpl) + e.visualCacheRaw = "hello" + e.visualCacheWidth = cpl + e.mu.Unlock() + + e.commitLines([]string{"world"}) + + e.mu.Lock() + cacheRaw := e.visualCacheRaw + e.mu.Unlock() + + assert.Equal(t, "", cacheRaw, "cache should be invalidated after mutation") +} + +// ── commitLines ────────────────────────────────────────────────────────────── + +func TestCommitLines_UpdatesText(t *testing.T) { + e := newTVE() + e.commitLines([]string{"new text"}) + assert.Equal(t, "new text", e.GetText()) +} + +func TestCommitLines_MultiLine(t *testing.T) { + e := newTVE() + e.commitLines([]string{"line one", "line two"}) + assert.Equal(t, "line one\nline two", e.GetText()) +} diff --git a/internal/ui/widgets/tokeniser.go b/internal/ui/widgets/tokeniser.go new file mode 100644 index 0000000..d471bcd --- /dev/null +++ b/internal/ui/widgets/tokeniser.go @@ -0,0 +1,362 @@ +package widgets + +import ( + "image/color" + "strings" + + "fyne.io/fyne/v2" + "fyne.io/fyne/v2/theme" +) + +// tvTokenKind identifies the syntax role of a token for colour mapping. +type tvTokenKind uint8 + +const ( + tvKindPlain tvTokenKind = iota + tvKindMethod // HTTP request method + tvKindPath // HTTP request path + tvKindVersion // HTTP version string + tvKindStatus2xx // 2xx response status code + tvKindStatus3xx // 3xx response status code + tvKindStatus4xx // 4xx response status code + tvKindStatus5xx // 5xx response status code + tvKindHdrName // header field name + tvKindHdrColon // colon separator after header name + tvKindHdrValue // header field value + tvKindJSONKey // JSON object key + tvKindJSONStr // JSON string value + tvKindJSONNum // JSON number + tvKindJSONBool // JSON true, false, null + tvKindJSONPunct // JSON structural punctuation: { } [ ] , : + tvKindLow // dimmed / de-emphasised text +) + +// tvColor returns the theme colour for a given token kind. +func tvColor(kind tvTokenKind) color.Color { + switch kind { + case tvKindMethod: + return theme.Color(theme.ColorNameError) + case tvKindPath: + return theme.Color(theme.ColorNamePrimary) + case tvKindVersion, tvKindLow: + return theme.Color(theme.ColorNamePlaceHolder) + case tvKindStatus2xx: + return theme.Color(theme.ColorNameSuccess) + case tvKindStatus3xx: + return theme.Color(theme.ColorNamePrimary) + case tvKindStatus4xx, tvKindStatus5xx: + return theme.Color(theme.ColorNameError) + case tvKindHdrName: + return theme.Color(theme.ColorNamePrimary) + case tvKindHdrColon, tvKindJSONPunct: + return theme.Color(theme.ColorNamePlaceHolder) + case tvKindHdrValue: + return theme.Color(theme.ColorNameForeground) + case tvKindJSONKey: + return theme.Color(theme.ColorNamePrimary) + case tvKindJSONStr: + return theme.Color(theme.ColorNameSuccess) + case tvKindJSONNum: + return theme.Color(theme.ColorNameWarning) + case tvKindJSONBool: + return theme.Color(theme.ColorNameError) + default: + return theme.Color(theme.ColorNameForeground) + } +} + +// tvToken is a run of text sharing a single syntax colour. +type tvToken struct { + Text string + Kind tvTokenKind +} + +// tvLine is one visual line (after word-wrap) composed of tokens. +// Raw holds the plain text for selection and clipboard operations. +type tvLine struct { + Tokens []tvToken + Raw string +} + +// parseAndWrap tokenises raw HTTP text and wraps each logical line to fit +// within wrapWidth pixels. Returns the complete slice of visual lines. +func parseAndWrap(s string, wrapWidth float32) []tvLine { + if s == "" { + return nil + } + s = strings.ReplaceAll(s, "\r\n", "\n") + s = strings.ReplaceAll(s, "\r", "\n") + + logicalLines := strings.Split(s, "\n") + + charWidth := fyne.MeasureText("M", theme.TextSize(), fyne.TextStyle{Monospace: true}).Width + if charWidth <= 0 { + charWidth = theme.TextSize() * 0.6 + } + charsPerLine := max(int(wrapWidth/charWidth), 10) + + bodyContentType := detectContentType(logicalLines) + + var result []tvLine + inBody := false + + for i, rawLine := range logicalLines { + if !inBody && strings.TrimSpace(rawLine) == "" && i > 0 { + inBody = true + result = append(result, tvLine{Tokens: []tvToken{{Text: "", Kind: tvKindPlain}}, Raw: ""}) + continue + } + + var tokens []tvToken + if !inBody { + tokens = tokeniseHTTPMeta(rawLine, i == 0) + } else { + tokens = tokeniseBody(rawLine, bodyContentType) + } + + result = append(result, wrapTokens(tokens, rawLine, charsPerLine)...) + } + return result +} + +// detectContentType scans the header section to determine the body content type. +func detectContentType(lines []string) string { + for _, line := range lines { + lower := strings.ToLower(line) + if strings.HasPrefix(lower, "content-type:") { + switch { + case strings.Contains(lower, "json"): + return "json" + case strings.Contains(lower, "html"): + return "html" + default: + return "text" + } + } + if line == "" { + break + } + } + return "text" +} + +// tokeniseHTTPMeta tokenises one line from the HTTP header section. +// isFirstLine indicates the request or response line. +func tokeniseHTTPMeta(line string, isFirstLine bool) []tvToken { + if line == "" { + return []tvToken{{Text: "", Kind: tvKindPlain}} + } + if isFirstLine { + return tokeniseFirstLine(line) + } + before, after, ok := strings.Cut(line, ":") + if !ok { + return []tvToken{{Text: line, Kind: tvKindPlain}} + } + return []tvToken{ + {Text: before, Kind: tvKindHdrName}, + {Text: ":", Kind: tvKindHdrColon}, + {Text: after, Kind: tvKindHdrValue}, + } +} + +// tokeniseFirstLine tokenises the HTTP request or response line. +func tokeniseFirstLine(line string) []tvToken { + if strings.TrimSpace(line) == "" { + return []tvToken{{Text: line, Kind: tvKindPlain}} + } + + parts := strings.Fields(line) + if len(parts) == 0 { + return []tvToken{{Text: line, Kind: tvKindPlain}} + } + + // Response line: HTTP/1.1 200 OK + if strings.HasPrefix(parts[0], "HTTP/") { + if len(parts) < 2 { + return []tvToken{{Text: line, Kind: tvKindPlain}} + } + + statusKind := tvKindStatus2xx + switch parts[1][0] { + case '3': + statusKind = tvKindStatus3xx + case '4': + statusKind = tvKindStatus4xx + case '5': + statusKind = tvKindStatus5xx + } + + tokens := []tvToken{ + {Text: parts[0] + " ", Kind: tvKindVersion}, + {Text: parts[1], Kind: statusKind}, + } + + if len(parts) > 2 { + tokens = append(tokens, tvToken{ + Text: " " + strings.Join(parts[2:], " "), + Kind: tvKindLow, + }) + } + + return tokens + } + + // Request line: GET /path HTTP/1.1 + if len(parts) >= 3 { + return []tvToken{ + {Text: parts[0] + " ", Kind: tvKindMethod}, + {Text: parts[1], Kind: tvKindPath}, + {Text: " " + strings.Join(parts[2:], " "), Kind: tvKindVersion}, + } + } + + return []tvToken{{Text: line, Kind: tvKindPlain}} +} + +// tokeniseBody tokenises a body line according to its content type. +func tokeniseBody(line string, contentType string) []tvToken { + if contentType == "json" { + return tokeniseJSON(line) + } + return []tvToken{{Text: line, Kind: tvKindPlain}} +} + +// tokeniseJSON performs a single-pass tokenisation of a JSON line. +func tokeniseJSON(line string) []tvToken { + if line == "" { + return []tvToken{{Text: "", Kind: tvKindPlain}} + } + var tokens []tvToken + runes := []rune(line) + pos := 0 + for pos < len(runes) { + ch := runes[pos] + switch { + case ch == '"': + end := pos + 1 + for end < len(runes) { + if runes[end] == '"' { + // Count preceding backslashes — even count means quote is unescaped. + numBackslashes := 0 + for i := end - 1; i >= pos && runes[i] == '\\'; i-- { + numBackslashes++ + } + if numBackslashes%2 == 0 { + break + } + } + end++ + } + if end < len(runes) { + end++ + } + str := string(runes[pos:end]) + peek := end + for peek < len(runes) && runes[peek] == ' ' { + peek++ + } + if peek < len(runes) && runes[peek] == ':' { + tokens = append(tokens, tvToken{Text: str, Kind: tvKindJSONKey}) + } else { + tokens = append(tokens, tvToken{Text: str, Kind: tvKindJSONStr}) + } + pos = end + case ch == '{' || ch == '}' || ch == '[' || ch == ']' || ch == ',' || ch == ':': + tokens = append(tokens, tvToken{Text: string(ch), Kind: tvKindJSONPunct}) + pos++ + case ch == '-' || (ch >= '0' && ch <= '9'): + end := pos + 1 + for end < len(runes) && (runes[end] >= '0' && runes[end] <= '9' || runes[end] == '.' || runes[end] == 'e' || runes[end] == 'E' || runes[end] == '+' || runes[end] == '-') { + end++ + } + tokens = append(tokens, tvToken{Text: string(runes[pos:end]), Kind: tvKindJSONNum}) + pos = end + case ch == 't' || ch == 'f' || ch == 'n': + matched := false + for _, keyword := range []string{"true", "false", "null"} { + if strings.HasPrefix(string(runes[pos:]), keyword) { + tokens = append(tokens, tvToken{Text: keyword, Kind: tvKindJSONBool}) + pos += len([]rune(keyword)) + matched = true + break + } + } + if !matched { + tokens = append(tokens, tvToken{Text: string(ch), Kind: tvKindPlain}) + pos++ + } + case ch == ' ' || ch == '\t': + end := pos + 1 + for end < len(runes) && (runes[end] == ' ' || runes[end] == '\t') { + end++ + } + tokens = append(tokens, tvToken{Text: string(runes[pos:end]), Kind: tvKindPlain}) + pos = end + default: + tokens = append(tokens, tvToken{Text: string(ch), Kind: tvKindPlain}) + pos++ + } + } + return tokens +} + +// wrapTokens splits a logical line into visual lines of at most charsPerLine runes. +func wrapTokens(tokens []tvToken, rawLine string, charsPerLine int) []tvLine { + if charsPerLine <= 0 { + charsPerLine = 80 + } + + var result []tvLine + var current []tvToken + currentLen := 0 + + flush := func() { + if len(current) == 0 { + return + } + var raw strings.Builder + for _, t := range current { + raw.WriteString(t.Text) + } + result = append(result, tvLine{Tokens: current, Raw: raw.String()}) + current = nil + currentLen = 0 + } + + for _, token := range tokens { + tokenLen := len([]rune(token.Text)) + + // never split JSON strings + if token.Kind == tvKindJSONStr || token.Kind == tvKindJSONKey { + if currentLen+tokenLen > charsPerLine { + flush() + } + current = append(current, token) + currentLen += tokenLen + continue + } + + // split other tokens if needed + runes := []rune(token.Text) + for len(runes) > 0 { + spaceLeft := charsPerLine - currentLen + if spaceLeft <= 0 { + flush() + spaceLeft = charsPerLine + } + + n := min(len(runes), spaceLeft) + + current = append(current, tvToken{ + Text: string(runes[:n]), + Kind: token.Kind, + }) + currentLen += n + runes = runes[n:] + } + } + + flush() + return result +} diff --git a/internal/ui/widgets/tokeniser_test.go b/internal/ui/widgets/tokeniser_test.go new file mode 100644 index 0000000..11c4d63 --- /dev/null +++ b/internal/ui/widgets/tokeniser_test.go @@ -0,0 +1,213 @@ +package widgets + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── detectContentType ──────────────────────────────────────────────────────── + +func TestDetectContentType_JSON(t *testing.T) { + lines := []string{"POST /foo HTTP/1.1", "Content-Type: application/json", ""} + assert.Equal(t, "json", detectContentType(lines)) +} + +func TestDetectContentType_HTML(t *testing.T) { + lines := []string{"GET /foo HTTP/1.1", "Content-Type: text/html; charset=utf-8", ""} + assert.Equal(t, "html", detectContentType(lines)) +} + +func TestDetectContentType_Plain(t *testing.T) { + lines := []string{"GET /foo HTTP/1.1", "Content-Type: text/plain", ""} + assert.Equal(t, "text", detectContentType(lines)) +} + +func TestDetectContentType_Missing(t *testing.T) { + lines := []string{"GET /foo HTTP/1.1", "Accept: */*", ""} + assert.Equal(t, "text", detectContentType(lines)) +} + +func TestDetectContentType_AfterBlankLineIgnored(t *testing.T) { + lines := []string{"GET /foo HTTP/1.1", "", "Content-Type: application/json"} + assert.Equal(t, "text", detectContentType(lines)) +} + +func TestDetectContentType_CaseInsensitive(t *testing.T) { + lines := []string{"GET /foo HTTP/1.1", "content-type: APPLICATION/JSON", ""} + assert.Equal(t, "json", detectContentType(lines)) +} + +// ── tokeniseFirstLine ──────────────────────────────────────────────────────── + +func TestTokeniseFirstLine_Request(t *testing.T) { + tokens := tokeniseFirstLine("GET /search?q=foo HTTP/1.1") + + require.Len(t, tokens, 3) + assert.Equal(t, tvKindMethod, tokens[0].Kind) + assert.Equal(t, tvKindPath, tokens[1].Kind) + assert.Equal(t, tvKindVersion, tokens[2].Kind) +} + +func TestTokeniseFirstLine_ResponseStatusKinds(t *testing.T) { + tests := []struct { + line string + wantKind tvTokenKind + }{ + {"HTTP/1.1 200 OK", tvKindStatus2xx}, + {"HTTP/1.1 201 Created", tvKindStatus2xx}, + {"HTTP/1.1 301 Moved Permanently", tvKindStatus3xx}, + {"HTTP/1.1 404 Not Found", tvKindStatus4xx}, + {"HTTP/1.1 500 Internal Server Error", tvKindStatus5xx}, + } + + for _, tt := range tests { + t.Run(tt.line, func(t *testing.T) { + tokens := tokeniseFirstLine(tt.line) + require.GreaterOrEqual(t, len(tokens), 2) + assert.Equal(t, tt.wantKind, tokens[1].Kind) + }) + } +} + +// ── tokeniseHTTPMeta ───────────────────────────────────────────────────────── + +func TestTokeniseHTTPMeta_Header(t *testing.T) { + tokens := tokeniseHTTPMeta("Content-Type: application/json", false) + + require.Len(t, tokens, 3) + assert.Equal(t, tvKindHdrName, tokens[0].Kind) + assert.Equal(t, "Content-Type", tokens[0].Text) + assert.Equal(t, tvKindHdrColon, tokens[1].Kind) + assert.Equal(t, tvKindHdrValue, tokens[2].Kind) +} + +func TestTokeniseHTTPMeta_EmptyLine(t *testing.T) { + tokens := tokeniseHTTPMeta("", false) + + require.Len(t, tokens, 1) + assert.Equal(t, tvKindPlain, tokens[0].Kind) + assert.Equal(t, "", tokens[0].Text) +} + +func TestTokeniseHTTPMeta_NoColon(t *testing.T) { + tokens := tokeniseHTTPMeta("plaintext", false) + + require.Len(t, tokens, 1) + assert.Equal(t, tvKindPlain, tokens[0].Kind) +} + +// ── tokeniseJSON ───────────────────────────────────────────────────────────── + +func TestTokeniseJSON_KeyAndString(t *testing.T) { + tokens := tokeniseJSON(`"name": "alice"`) + + kinds := map[tvTokenKind]bool{} + for _, tok := range tokens { + kinds[tok.Kind] = true + } + assert.True(t, kinds[tvKindJSONKey], "expected tvKindJSONKey token") + assert.True(t, kinds[tvKindJSONStr], "expected tvKindJSONStr token") +} + +func TestTokeniseJSON_Number(t *testing.T) { + tokens := tokeniseJSON(`"count": 42`) + + found := false + for _, tok := range tokens { + if tok.Kind == tvKindJSONNum && tok.Text == "42" { + found = true + } + } + assert.True(t, found, "expected number token '42'") +} + +func TestTokeniseJSON_BoolAndNull(t *testing.T) { + for _, kw := range []string{"true", "false", "null"} { + t.Run(kw, func(t *testing.T) { + tokens := tokeniseJSON(`"flag": ` + kw) + found := false + for _, tok := range tokens { + if tok.Kind == tvKindJSONBool && tok.Text == kw { + found = true + } + } + assert.True(t, found, "expected tvKindJSONBool token %q", kw) + }) + } +} + +func TestTokeniseJSON_Punctuation(t *testing.T) { + for _, ch := range []string{"{", "}", "[", "]", ","} { + t.Run(ch, func(t *testing.T) { + tokens := tokeniseJSON(ch) + require.Len(t, tokens, 1) + assert.Equal(t, tvKindJSONPunct, tokens[0].Kind) + }) + } +} + +func TestTokeniseJSON_EscapedQuoteInsideString(t *testing.T) { + tokens := tokeniseJSON(`"say \"hi\""`) + + strCount := 0 + for _, tok := range tokens { + if tok.Kind == tvKindJSONStr || tok.Kind == tvKindJSONKey { + strCount++ + } + } + assert.Equal(t, 1, strCount, "escaped quote should produce exactly 1 string token") +} + +func TestTokeniseJSON_DoubleBackslashBeforeQuote(t *testing.T) { + tokens := tokeniseJSON(`"\\"`) + + strCount := 0 + for _, tok := range tokens { + if tok.Kind == tvKindJSONStr || tok.Kind == tvKindJSONKey { + strCount++ + } + } + assert.Equal(t, 1, strCount, "double backslash before closing quote should produce exactly 1 string token") +} + +func TestTokeniseJSON_Empty(t *testing.T) { + tokens := tokeniseJSON("") + + require.Len(t, tokens, 1) + assert.Equal(t, tvKindPlain, tokens[0].Kind) +} + +// ── wrapTokens ─────────────────────────────────────────────────────────────── + +func TestWrapTokens_ShortLineNoWrap(t *testing.T) { + tokens := []tvToken{{Text: "hello", Kind: tvKindPlain}} + lines := wrapTokens(tokens, "hello", 80) + + require.Len(t, lines, 1) + assert.Equal(t, "hello", lines[0].Raw) +} + +func TestWrapTokens_LongLineWraps(t *testing.T) { + long := strings.Repeat("x", 100) + tokens := []tvToken{{Text: long, Kind: tvKindPlain}} + lines := wrapTokens(tokens, long, 40) + + assert.Greater(t, len(lines), 1, "expected wrapping into multiple lines") + + total := 0 + for _, l := range lines { + total += len([]rune(l.Raw)) + } + assert.Equal(t, 100, total, "all characters must be preserved across wrapped lines") +} + +func TestWrapTokens_JSONStringNotSplit(t *testing.T) { + str := `"` + strings.Repeat("a", 50) + `"` + tokens := []tvToken{{Text: str, Kind: tvKindJSONStr}} + lines := wrapTokens(tokens, str, 20) + + assert.Len(t, lines, 1, "JSON string must not be split across lines") +} -- cgit v1.2.3