diff --git a/agent/worker.go b/agent/worker.go index 4965e3c..99125c9 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -159,7 +159,7 @@ func workerTrafficBlock(proxyAddr string) string { if proxyAddr == "" { return "" } - return "\n\n**流量工具**:\n- traffic_search / traffic_get:目标流量已被记录代理全量落库(本地文件树 data/traffic//<访问路径>/)。回看响应、找已访问过的资源,**先查流量、不要重复 curl 同一 URL**。traffic_search **必须指定 host**、默认只回 3 条极轻量索引(id/method/url/status/resp_len,无响应内容),需要更多显式调大 limit;要看某条原文用 traffic_get(id)。" + return "\n\n**流量工具**:\n- traffic_search / traffic_get / traffic_blob:回看响应、找已访问过的资源,**先查流量、不要重复 curl 同一 URL**。traffic_search **必须指定 host**、默认只回 3 条极轻量索引(id/method/url/status/resp_len,无响应内容),需要更多显式调大 limit;可用 body_contains 在请求/响应正文里做全文搜索(至少 3 字符,支持子串和中文,如找密码/密钥/报错/内网地址);要看某条原文用 traffic_get(id),其中超大正文显示为 @blob sha256:,用 traffic_blob(hash) 分段取全文。" } // artifactSpec is 段 [C]: the code-owned, non-editable tail appended to every diff --git a/server/server.go b/server/server.go index a1ca28b..f376df2 100644 --- a/server/server.go +++ b/server/server.go @@ -717,6 +717,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("DELETE /api/traffic", s.deleteTraffic) mux.HandleFunc("DELETE /api/traffic/hosts", s.deleteTrafficHosts) mux.HandleFunc("GET /api/traffic/exchange", s.getTrafficExchange) + mux.HandleFunc("GET /api/traffic/blob", s.getTrafficBlob) mux.HandleFunc("GET /api/commands", s.pgListCommands) mux.HandleFunc("GET /api/llm/records", s.pgListLLMRecords) mux.HandleFunc("DELETE /api/llm/records", s.pgDeleteLLMRecords) @@ -2934,6 +2935,35 @@ func (s *Server) getTrafficExchange(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"req": req, "resp": resp}) } +// getTrafficBlob streams one oversized body by its sha256. Bodies past the inline +// threshold are not carried by the exchange endpoint — it returns a preview plus +// an "@blob sha256:" pointer — so this is how the UI fetches them whole. +// Streamed rather than buffered: these are the bodies too large to hold in memory. +func (s *Server) getTrafficBlob(w http.ResponseWriter, r *http.Request) { + tr := s.m.Traffic() + if tr == nil { + writeErr(w, 404, "traffic disabled") + return + } + hash := strings.TrimSpace(r.URL.Query().Get("hash")) + if hash == "" { + writeErr(w, 400, "missing hash") + return + } + f, size, err := tr.Blob(hash) + if err != nil { + writeErr(w, 404, err.Error()) + return + } + defer f.Close() + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Content-Length", strconv.FormatInt(size, 10)) + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", hash+".bin")) + if _, err := io.Copy(w, f); err != nil { + log.Printf("[traffic] 下载 blob %s 中断:%v", hash, err) + } +} + // getSettings returns the runtime app settings the UI toggles. The Brave API key // is returned as a boolean presence flag (brave_key_set), never the value itself, // so the UI can show "configured" without echoing the secret back. diff --git a/traffic/traffic.go b/traffic/traffic.go index bf8f241..e827443 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -5,6 +5,7 @@ package traffic import ( + "bytes" "context" "crypto/sha256" "database/sql" @@ -12,6 +13,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "io/fs" "log" "net" @@ -20,10 +22,12 @@ import ( "os" "path/filepath" "regexp" + "sort" "strings" "sync" "sync/atomic" "time" + "unicode/utf8" "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/permission" @@ -49,9 +53,53 @@ CREATE TABLE IF NOT EXISTS exchanges ( CREATE INDEX IF NOT EXISTS idx_ex_host ON exchanges(host); CREATE INDEX IF NOT EXISTS idx_ex_tmpl ON exchanges(host, url_template); CREATE INDEX IF NOT EXISTS idx_ex_ts ON exchanges(ts); + +CREATE TABLE IF NOT EXISTS exchange_bodies ( + id TEXT PRIMARY KEY, + req_head TEXT NOT NULL DEFAULT '', + req_body BLOB, + req_blob TEXT, + resp_head TEXT NOT NULL DEFAULT '', + resp_body BLOB, + resp_blob TEXT +); + +CREATE TABLE IF NOT EXISTS blob_refs ( + hash TEXT NOT NULL, + exchange_id TEXT NOT NULL, + PRIMARY KEY (hash, exchange_id) +); +CREATE INDEX IF NOT EXISTS idx_blob_refs_ex ON blob_refs(exchange_id); ` -const maxInlineBody = 256 * 1024 +// ftsSchema is applied separately from indexSchema: a driver build without FTS5 +// must degrade to "no full-text search" rather than take the whole recorder down. +// trigram (not the default unicode61) is required for two reasons this subsystem +// depends on: it matches arbitrary substrings — "ssw0r" finds "P@ssw0rd" — and it +// handles CJK, which unicode61 does not tokenize. Contentless (content='') keeps +// only the index, since the text itself lives in exchange_bodies; contentless_delete +// lets rows be deleted without replaying the original text back in. +const ftsSchema = `CREATE VIRTUAL TABLE IF NOT EXISTS ex_fts USING fts5( + content, tokenize='trigram', content='', contentless_delete=1 +);` + +const ( + // maxInlineBody is the cutoff between a body stored inline (SQLite column) and + // one spilled to the content-addressed blob store. + maxInlineBody = 256 * 1024 + // blobPreview is how much of a spilled body stays inline so readers can tell + // what it is (JSON shape, SQL dump header, ZIP magic) without fetching it. + blobPreview = 8 * 1024 + // maxIndexBody caps a single body's contribution to the full-text index. Text + // is indexed in full below this; binary never reaches the index at all. + maxIndexBody = 4 * 1024 * 1024 + // minTrigram is the shortest term the trigram tokenizer can match; shorter + // queries fall back to metadata LIKE. + minTrigram = 3 + // maxBlobRead caps one traffic_blob read so paging through a large body never + // floods the agent's context. + maxBlobRead = 8 * 1024 +) // Traffic runs the recording proxy and owns the file tree + index. type Traffic struct { @@ -61,6 +109,13 @@ type Traffic struct { wmu sync.Mutex // serializes record() vs DeleteHost (incl. blob GC) seq atomic.Int64 proxy *mproxy.Proxy + // fts reports whether the full-text index is available. False on a driver + // build without FTS5: recording and metadata search still work, body search + // degrades to unsupported rather than erroring. + fts bool + // reaping tracks background reclamation of legacy trees, so shutdown and tests + // can wait for it instead of racing it. + reaping sync.WaitGroup // pass is the set of hosts whose MITM interception failed for a proxy/protocol // reason; connections to them are tunneled transparently (fail-open) so the // request still reaches the target — unrecorded — instead of being killed. @@ -89,6 +144,11 @@ func Open(dir, addr string) (*Traffic, error) { return nil, err } t := &Traffic{dir: dir, addr: addr, db: db} + if _, err := db.Exec(ftsSchema); err != nil { + log.Printf("[traffic] 全文索引不可用,正文搜索将被禁用(元数据搜索不受影响):%v", err) + } else { + t.fts = true + } p, err := mproxy.NewProxy(&mproxy.Options{ Addr: addr, @@ -138,7 +198,13 @@ func (t *Traffic) CACertPath() string { // Start runs the proxy (blocking); run in a goroutine. func (t *Traffic) Start() error { return t.proxy.Start() } -func (t *Traffic) Close() error { return t.db.Close() } +// Close waits for background tree reclamation to finish before closing the +// index, so shutdown never leaves a goroutine unlinking files out from under a +// removed data directory. +func (t *Traffic) Close() error { + t.reaping.Wait() + return t.db.Close() +} func (t *Traffic) DB() *sql.DB { return t.db } // sink is the go-mitmproxy addon that records completed exchanges. @@ -188,9 +254,12 @@ func proxyCausedErr(err error) bool { return false } +// record persists one exchange entirely in SQLite: metadata, bodies and the +// full-text index. Nothing is written to a per-request directory — only bodies +// above maxInlineBody spill to the content-addressed blob store. func (t *Traffic) record(f *mproxy.Flow) { - // The write lock covers blob + tree + index writes, so DeleteHost (and its - // blob GC) can run under the same lock without racing a concurrent record. + // The write lock covers blob + index writes, so DeleteHost (and its blob GC) + // can run under the same lock without racing a concurrent record. t.wmu.Lock() defer t.wmu.Unlock() host := f.Request.URL.Hostname() @@ -200,58 +269,193 @@ func (t *Traffic) record(f *mproxy.Flow) { now := time.Now() id := fmt.Sprintf("%d-%04d", now.Unix(), n%10000) - // dir mirrors the URL path as a browsable file tree: - // ///...///{request.http,response.http,meta.json} - // e.g. http://h/api/v1/login → traffic/h/api/v1/login/GET// - parts := []string{t.dir, sanitize(host)} - for _, seg := range strings.Split(strings.Trim(tmpl, "/"), "/") { - if seg != "" { - parts = append(parts, sanitize(seg)) - } + ct := f.Response.Header.Get("Content-Type") + url := f.Request.URL.String() + reqHead := fmt.Sprintf("%s %s %s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, requestHeaderLines(f.Request)) + respHead := fmt.Sprintf("HTTP %d\n%s", f.Response.StatusCode, headerLines(f.Response.Header)) + // Bodies are spilled before the transaction opens: blob writes are filesystem + // work and must not sit inside the SQLite write lock. + reqB := t.spill(f.Request.Body, f.Request.Header.Get("Content-Type")) + respB := t.spill(f.Response.Body, ct) + + tx, err := t.db.Begin() + if err != nil { + log.Printf("[traffic] 记录 %s 失败(开启事务):%v", url, err) + return } - parts = append(parts, method, id) - exDir := filepath.Join(parts...) - if err := os.MkdirAll(exDir, 0o755); err != nil { + defer tx.Rollback() //nolint:errcheck // no-op once committed + + // path stays empty for db-resident exchanges; a non-empty path marks a legacy + // row whose bodies still live in the old on-disk tree (see Get). + res, err := tx.Exec(`INSERT OR REPLACE INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,'')`, + id, now.Unix(), host, method, tmpl, url, f.Response.StatusCode, ct, + len(f.Request.Body), len(f.Response.Body)) + if err != nil { + log.Printf("[traffic] 记录 %s 失败(写索引):%v", url, err) + return + } + rowid, err := res.LastInsertId() + if err != nil { + log.Printf("[traffic] 记录 %s 失败(取 rowid):%v", url, err) return } - reqBody := t.bodyOrBlob(f.Request.Body) - respBody := t.bodyOrBlob(f.Response.Body) - ct := f.Response.Header.Get("Content-Type") + if _, err := tx.Exec(`INSERT OR REPLACE INTO exchange_bodies(id,req_head,req_body,req_blob,resp_head,resp_body,resp_blob) +VALUES(?,?,?,?,?,?,?)`, + id, reqHead, reqB.inline, nullIfEmpty(reqB.hash), respHead, respB.inline, nullIfEmpty(respB.hash)); err != nil { + log.Printf("[traffic] 记录 %s 失败(写正文):%v", url, err) + return + } - reqTxt := fmt.Sprintf("%s %s %s\n%s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, requestHeaderLines(f.Request), reqBody) - respTxt := fmt.Sprintf("HTTP %d\n%s\n%s", f.Response.StatusCode, headerLines(f.Response.Header), respBody) - _ = os.WriteFile(filepath.Join(exDir, "request.http"), []byte(reqTxt), 0o644) - _ = os.WriteFile(filepath.Join(exDir, "response.http"), []byte(respTxt), 0o644) - meta := fmt.Sprintf(`{"id":%q,"host":%q,"method":%q,"url":%q,"template":%q,"status":%d,"content_type":%q}`, - id, host, method, f.Request.URL.String(), tmpl, f.Response.StatusCode, ct) - _ = os.WriteFile(filepath.Join(exDir, "meta.json"), []byte(meta), 0o644) + for _, h := range []string{reqB.hash, respB.hash} { + if h == "" { + continue + } + if _, err := tx.Exec(`INSERT OR IGNORE INTO blob_refs(hash,exchange_id) VALUES(?,?)`, h, id); err != nil { + log.Printf("[traffic] 记录 %s 失败(登记 blob 引用):%v", url, err) + return + } + } - rel, _ := filepath.Rel(t.dir, exDir) - _, _ = t.db.Exec(`INSERT OR REPLACE INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) -VALUES(?,?,?,?,?,?,?,?,?,?,?)`, - id, now.Unix(), host, method, tmpl, f.Request.URL.String(), f.Response.StatusCode, ct, - len(f.Request.Body), len(f.Response.Body), rel) + if t.fts { + // The index is fed from memory, so a body that spilled to _blobs is still + // fully searchable even though only its preview is stored inline. + idx := strings.Join([]string{url, reqHead, reqB.index, respHead, respB.index}, "\n") + if _, err := tx.Exec(`INSERT INTO ex_fts(rowid,content) VALUES(?,?)`, rowid, idx); err != nil { + log.Printf("[traffic] 记录 %s 失败(写全文索引):%v", url, err) + return + } + } + + if err := tx.Commit(); err != nil { + log.Printf("[traffic] 记录 %s 失败(提交):%v", url, err) + } } -// bodyOrBlob returns the body inline if small, else stores it content-addressed -// and returns a human-readable pointer. -func (t *Traffic) bodyOrBlob(body []byte) string { +// storedBody is one body after the inline/spill decision: inline is what goes in +// the SQLite column (the whole body, or a preview when spilled), hash names the +// blob when it spilled, and index is the text handed to FTS (empty for binary). +type storedBody struct { + inline []byte + hash string + index string +} + +// spill decides where one body lives. Small bodies stay inline. Large ones are +// written to the content-addressed store and keep a readable preview inline so a +// reader can identify them without fetching the blob. Either way, text bodies +// are handed to the full-text index in full (up to maxIndexBody) — indexing is +// independent of where the bytes end up. +func (t *Traffic) spill(body []byte, contentType string) storedBody { if len(body) == 0 { - return "" + return storedBody{} + } + text := !isBinaryBody(contentType, body) + indexText := func() string { + if !text { + return "" + } + return string(clipBytes(body, maxIndexBody)) } if len(body) <= maxInlineBody { - return string(body) + return storedBody{inline: body, index: indexText()} } + sum := sha256.Sum256(body) h := hex.EncodeToString(sum[:]) - blobDir := filepath.Join(t.dir, "_blobs", "sha256", h[:2], h[2:4]) - _ = os.MkdirAll(blobDir, 0o755) + // One bucket level (256 buckets) is enough to keep any single directory small; + // the store only ever holds bodies above maxInlineBody, deduplicated by hash. + blobDir := filepath.Join(t.dir, "_blobs", "sha256", h[:2]) + if err := os.MkdirAll(blobDir, 0o755); err != nil { + log.Printf("[traffic] 创建 blob 目录失败:%v", err) + return storedBody{inline: clipBytes(body, blobPreview), index: indexText()} + } blobPath := filepath.Join(blobDir, h+".bin") if _, err := os.Stat(blobPath); os.IsNotExist(err) { - _ = os.WriteFile(blobPath, body, 0o644) + if err := os.WriteFile(blobPath, body, 0o644); err != nil { + log.Printf("[traffic] 写 blob %s 失败:%v", h, err) + return storedBody{inline: clipBytes(body, blobPreview), index: indexText()} + } } - return fmt.Sprintf("@blob sha256:%s (len=%d)", h, len(body)) + sb := storedBody{hash: h, index: indexText()} + if text { + sb.inline = []byte(truncateUTF8(body, blobPreview)) + } else { + sb.inline = []byte(binaryTag(contentType, body)) + } + return sb +} + +// binaryTypes are content-type prefixes whose bodies are never worth indexing or +// previewing as text. Anything not listed is treated as text (with a NUL-byte +// check as backstop), so unusual-but-searchable types like application/sql or a +// bare text/* are not silently dropped from the index. +var binaryTypes = []string{ + "image/", "audio/", "video/", "font/", + "application/octet-stream", "application/zip", "application/gzip", + "application/x-gzip", "application/x-tar", "application/x-7z-compressed", + "application/x-rar", "application/pdf", "application/x-msdownload", + "application/vnd.android.package-archive", "application/java-archive", + "application/wasm", "application/x-shockwave-flash", +} + +func isBinaryBody(contentType string, body []byte) bool { + ct := strings.ToLower(strings.TrimSpace(contentType)) + if i := strings.IndexByte(ct, ';'); i >= 0 { + ct = strings.TrimSpace(ct[:i]) + } + for _, p := range binaryTypes { + if strings.HasPrefix(ct, p) { + return true + } + } + // Backstop for mislabeled or absent content types: real text does not carry + // NUL bytes, so a NUL in the head of the body means binary regardless. + return bytes.IndexByte(clipBytes(body, 512), 0) >= 0 +} + +// binaryTag describes a spilled binary body in one line, including the leading +// magic bytes so a reader can recognize the format without downloading it. +func binaryTag(contentType string, body []byte) string { + ct := strings.TrimSpace(contentType) + if ct == "" { + ct = "application/octet-stream" + } + magic := hex.EncodeToString(clipBytes(body, 4)) + return fmt.Sprintf("[binary %s, %d bytes, magic=%s]", ct, len(body), magic) +} + +func clipBytes(b []byte, n int) []byte { + if len(b) <= n { + return b + } + return b[:n] +} + +// truncateUTF8 cuts b to at most n bytes without splitting a multi-byte rune, +// so a preview never ends in half a CJK character. +func truncateUTF8(b []byte, n int) string { + if len(b) <= n { + return string(b) + } + b = b[:n] + // A rune is at most 4 bytes, so backing off 3 bytes always finds the boundary + // (unless the input was already invalid UTF-8, in which case we keep the cut). + for i := 0; i < utf8.UTFMax-1 && len(b) > 0; i++ { + if r, size := utf8.DecodeLastRune(b); r != utf8.RuneError || size != 1 { + break + } + b = b[:len(b)-1] + } + return string(b) +} + +func nullIfEmpty(s string) any { + if s == "" { + return nil + } + return s } func headerLines(h map[string][]string) string { @@ -299,6 +503,75 @@ func sanitize(s string) string { return out } +var blobHashRe = regexp.MustCompile(`^[0-9a-f]{64}$`) + +// blobPath locates a stored blob. Current writes use one bucket level; blobs +// written before that change used two, so both layouts are probed rather than +// migrated. +func (t *Traffic) blobPath(hash string) (string, error) { + hash = strings.ToLower(strings.TrimSpace(hash)) + // Validated as pure hex before touching the filesystem, so a crafted hash can + // never traverse out of the blob directory. + if !blobHashRe.MatchString(hash) { + return "", fmt.Errorf("非法的 blob hash") + } + for _, p := range []string{ + filepath.Join(t.dir, "_blobs", "sha256", hash[:2], hash+".bin"), + filepath.Join(t.dir, "_blobs", "sha256", hash[:2], hash[2:4], hash+".bin"), + } { + if st, err := os.Stat(p); err == nil && !st.IsDir() { + return p, nil + } + } + return "", fmt.Errorf("blob %s 不存在", hash) +} + +// Blob opens a spilled body for streaming; the caller must close the file. +// Streaming matters here: the driver exposes no incremental BLOB API, so keeping +// large bodies on disk is what lets them be served without loading them whole. +func (t *Traffic) Blob(hash string) (*os.File, int64, error) { + p, err := t.blobPath(hash) + if err != nil { + return nil, 0, err + } + f, err := os.Open(p) + if err != nil { + return nil, 0, err + } + st, err := f.Stat() + if err != nil { + f.Close() + return nil, 0, err + } + return f, st.Size(), nil +} + +// BlobRange reads at most length bytes of a blob from offset, and reports the +// blob's total size. Used by the agent tool to page through a large body without +// pulling all of it into the model's context. +func (t *Traffic) BlobRange(hash string, offset, length int64) (data []byte, total int64, err error) { + f, size, err := t.Blob(hash) + if err != nil { + return nil, 0, err + } + defer f.Close() + if offset < 0 { + offset = 0 + } + if offset >= size { + return nil, size, nil + } + if length <= 0 || offset+length > size { + length = size - offset + } + buf := make([]byte, length) + n, err := f.ReadAt(buf, offset) + if err != nil && !errors.Is(err, io.EOF) { + return nil, size, err + } + return buf[:n], size, nil +} + // ExchangeMeta is one row of the index (returned by Search). type ExchangeMeta struct { ID string `json:"id"` @@ -342,11 +615,31 @@ func (t *Traffic) Search(host string, page, size int) ([]ExchangeMeta, error) { return out, rows.Err() } +// ftsFilter builds the SQL condition restricting rows to those whose indexed +// text matches term. ok is false when full-text search cannot serve the term — +// no FTS index on this build, or fewer than three characters, which the trigram +// tokenizer cannot match — leaving callers to fall back to metadata matching. +func (t *Traffic) ftsFilter(term string) (cond string, arg any, ok bool) { + term = strings.TrimSpace(term) + if !t.fts || utf8.RuneCountInString(term) < minTrigram { + return "", nil, false + } + return `rowid IN (SELECT rowid FROM ex_fts WHERE ex_fts MATCH ?)`, ftsQuote(term), true +} + +// ftsQuote wraps a term as an FTS5 string literal so that punctuation and query +// operators inside it (quotes, AND/OR/NEAR, *, ^) are matched literally instead +// of being parsed as query syntax. +func ftsQuote(s string) string { + return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` +} + // Page returns one page of exchange metadata filtered by an optional host -// substring, exact method, and a free-text query q that fuzzy-matches across all -// indexed metadata columns (host/url/method/content-type/status), together with -// the total number of rows matching that filter (for the UI's pagination). -// Newest first. Bodies are never included. +// substring, exact method, and a free-text query q that matches across the +// indexed metadata columns (host/url/method/content-type/status) and, when the +// term is long enough for the trigram index, across captured request/response +// text as well. Also returns the total number of rows matching that filter (for +// the UI's pagination). Newest first. Bodies are never included in the rows. func (t *Traffic) Page(host, method, q string, page, size int) (rows []ExchangeMeta, total int, err error) { if size <= 0 || size > 500 { size = 100 @@ -373,9 +666,14 @@ func (t *Traffic) Page(host, method, q string, page, size int) (rows []ExchangeM } if s := strings.TrimSpace(q); s != "" { like := "%" + s + "%" - // Fuzzy match across every indexed column (bodies aren't indexed). - add("(host LIKE ? OR url LIKE ? OR url_template LIKE ? OR method LIKE ? OR content_type LIKE ? OR CAST(status AS TEXT) LIKE ?)", - like, like, like, like, like, like) + const meta = "host LIKE ? OR url LIKE ? OR url_template LIKE ? OR method LIKE ? OR content_type LIKE ? OR CAST(status AS TEXT) LIKE ?" + // Metadata match OR full-text match: one search box, widest recall. Terms + // too short for trigram silently fall back to metadata only. + if cond, arg, ok := t.ftsFilter(s); ok { + add("(("+meta+") OR "+cond+")", like, like, like, like, like, like, arg) + } else { + add("("+meta+")", like, like, like, like, like, like) + } } if err = t.db.QueryRow(`SELECT COUNT(*) FROM exchanges`+where, args...).Scan(&total); err != nil { return nil, 0, err @@ -398,17 +696,51 @@ func (t *Traffic) Page(host, method, q string, page, size int) (rows []ExchangeM return rows, total, rs.Err() } -// Get returns the full request/response of one exchange (reads from the tree). +// Get returns the full request/response text of one exchange. Bodies come from +// the database; rows recorded before bodies moved into SQLite carry a non-empty +// path and are read from the legacy on-disk tree instead, so history stays +// readable without being migrated. func (t *Traffic) Get(id string) (req, resp string, err error) { + var reqHead, respHead string + var reqBody, respBody []byte + var reqBlob, respBlob sql.NullString + var reqLen, respLen int + err = t.db.QueryRow(`SELECT b.req_head,b.req_body,b.req_blob,b.resp_head,b.resp_body,b.resp_blob,e.req_len,e.resp_len +FROM exchange_bodies b JOIN exchanges e ON e.id=b.id WHERE b.id=?`, id). + Scan(&reqHead, &reqBody, &reqBlob, &respHead, &respBody, &respBlob, &reqLen, &respLen) + if err == nil { + return assembleRaw(reqHead, reqBody, reqBlob, reqLen), assembleRaw(respHead, respBody, respBlob, respLen), nil + } + if !errors.Is(err, sql.ErrNoRows) { + return "", "", err + } + var rel string if err = t.db.QueryRow(`SELECT path FROM exchanges WHERE id=?`, id).Scan(&rel); err != nil { return "", "", err } + if strings.TrimSpace(rel) == "" { + return "", "", fmt.Errorf("exchange %s 无正文记录", id) + } rb, _ := os.ReadFile(filepath.Join(t.dir, rel, "request.http")) pb, _ := os.ReadFile(filepath.Join(t.dir, rel, "response.http")) return string(rb), string(pb), nil } +// assembleRaw rebuilds one side's raw HTTP text: head, blank line, body. A body +// that spilled to the blob store shows its inline preview followed by the blob +// pointer, so the reader can see what it is and fetch the rest by hash. +func assembleRaw(head string, body []byte, blob sql.NullString, total int) string { + var b strings.Builder + b.WriteString(head) + b.WriteByte('\n') + b.Write(body) + if blob.Valid && blob.String != "" { + fmt.Fprintf(&b, "\n…[truncated] @blob sha256:%s (len=%d)", blob.String, total) + } + return b.String() +} + // HostCount is one distinct recorded host plus its exchange count. type HostCount struct { Host string `json:"host"` @@ -444,46 +776,177 @@ func (t *Traffic) Count() (int, error) { // DeleteHost removes every recorded exchange whose host contains the given // substring — mirroring the UI's host filter, so what you filtered is what gets -// deleted: the index rows plus each matched host's file tree (//…). -// Content-addressed blobs in _blobs are shared across hosts, so instead of -// deleting by host they are garbage-collected afterwards: any blob no longer -// referenced by a remaining exchange tree is removed (frees the data volume). -// Returns rows deleted. +// deleted. Everything lives in SQLite, so the deletion is one transaction across +// the index, bodies, full-text index and blob references. Content-addressed +// blobs are shared across hosts and so are garbage-collected afterwards rather +// than deleted by host. Returns rows deleted. func (t *Traffic) DeleteHost(host string) (int64, error) { like := "%" + host + "%" t.wmu.Lock() defer t.wmu.Unlock() - // Collect distinct matched hosts first (index is the source of truth for - // tree names) so the trees are removed even if the DELETE row count is 0. - rows, err := t.db.Query(`SELECT DISTINCT host FROM exchanges WHERE host LIKE ?`, like) + // Resolved before the rows go away: with a substring match, the index is the + // only record of which host directories the filter actually hit. + legacy, err := t.hostTrees(`host LIKE ?`, like) if err != nil { return 0, err } - var hosts []string + tx, err := t.db.Begin() + if err != nil { + return 0, err + } + defer tx.Rollback() //nolint:errcheck // no-op once committed + n, err := t.deleteWhere(tx, `host LIKE ?`, like) + if err != nil { + return 0, err + } + // Staged before the commit so a filesystem failure can still abort the whole + // deletion, and before gcBlobs so reference collection sees the correct live + // set — a staged tree is invisible to it. + stageDir, moves, err := t.stageTrees(legacy) + if err != nil { + return 0, errors.Join(err, restoreTrees(stageDir, moves)) + } + if err := tx.Commit(); err != nil { + return 0, errors.Join(fmt.Errorf("提交流量索引删除: %w", err), restoreTrees(stageDir, moves)) + } + t.reapStage(stageDir) + if n > 0 { + if err := t.gcBlobs(); err != nil { + return n, err + } + } + return n, nil +} + +// deleteWhere removes every trace of the exchanges matching the condition: +// full-text index rows, bodies, blob references, and finally the index rows. +// Order matters — each sub-select reads exchanges, so that table is emptied last. +func (t *Traffic) deleteWhere(tx *sql.Tx, where string, args ...any) (int64, error) { + if t.fts { + // ex_fts is contentless and addressed by rowid, hence the rowid sub-select. + if _, err := tx.Exec(`DELETE FROM ex_fts WHERE rowid IN (SELECT rowid FROM exchanges WHERE `+where+`)`, args...); err != nil { + return 0, fmt.Errorf("删除全文索引: %w", err) + } + } + if _, err := tx.Exec(`DELETE FROM exchange_bodies WHERE id IN (SELECT id FROM exchanges WHERE `+where+`)`, args...); err != nil { + return 0, fmt.Errorf("删除正文: %w", err) + } + if _, err := tx.Exec(`DELETE FROM blob_refs WHERE exchange_id IN (SELECT id FROM exchanges WHERE `+where+`)`, args...); err != nil { + return 0, fmt.Errorf("删除 blob 引用: %w", err) + } + res, err := tx.Exec(`DELETE FROM exchanges WHERE `+where, args...) + if err != nil { + return 0, fmt.Errorf("删除索引行: %w", err) + } + n, _ := res.RowsAffected() + return n, nil +} + +// hostTrees returns the on-disk directory of every host matching the condition. +// Exchanges recorded since bodies moved into SQLite have no directory at all, so +// most of these paths simply will not exist — staging skips those. The lookup is +// deliberately not restricted to rows with a path: a host whose index rows are +// already gone can still have an orphaned directory, and deleting the host should +// take that with it. +func (t *Traffic) hostTrees(where string, args ...any) ([]string, error) { + rows, err := t.db.Query(`SELECT DISTINCT host FROM exchanges WHERE `+where, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var dirs []string for rows.Next() { var h string if err := rows.Scan(&h); err != nil { - rows.Close() - return 0, err + return nil, err } - hosts = append(hosts, h) + dirs = append(dirs, filepath.Join(t.dir, sanitize(h))) } - rows.Close() - if err := rows.Err(); err != nil { - return 0, err + return dirs, rows.Err() +} + +type stagedTrafficPath struct { + source string + staged string +} + +// stageTrees moves host directories aside onto the same filesystem. The rename is +// atomic and instant, which buys two things the unlink cannot: the deletion stays +// reversible until the transaction commits, and the tree stops being visible to +// blob reference collection right away (legacyBlobRefs skips underscore-prefixed +// directories, so anything under _delete_staging is already out of the live set). +// An empty stageDir is returned when there was nothing to stage. +func (t *Traffic) stageTrees(dirs []string) (stageDir string, moves []stagedTrafficPath, err error) { + seen := make(map[string]struct{}, len(dirs)) + for _, source := range dirs { + if _, dup := seen[source]; dup { + continue + } + seen[source] = struct{}{} + if _, err := os.Lstat(source); err != nil { + if os.IsNotExist(err) { + continue + } + return stageDir, moves, fmt.Errorf("检查历史流量目录 %s: %w", source, err) + } + if stageDir == "" { + parent := filepath.Join(t.dir, "_delete_staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return stageDir, moves, fmt.Errorf("创建流量暂存目录: %w", err) + } + if stageDir, err = os.MkdirTemp(parent, "hosts-"); err != nil { + return stageDir, moves, fmt.Errorf("创建流量暂存目录: %w", err) + } + } + staged := filepath.Join(stageDir, fmt.Sprintf("%d-%s", len(moves), filepath.Base(source))) + if err := os.Rename(source, staged); err != nil { + return stageDir, moves, fmt.Errorf("移出历史流量目录 %s: %w", source, err) + } + moves = append(moves, stagedTrafficPath{source: source, staged: staged}) } - res, err := t.db.Exec(`DELETE FROM exchanges WHERE host LIKE ?`, like) - if err != nil { - return 0, err + return stageDir, moves, nil +} + +// restoreTrees puts staged directories back where they came from, newest move +// first, and drops the staging directory once everything is home. +func restoreTrees(stageDir string, moves []stagedTrafficPath) error { + var errs []error + for i := len(moves) - 1; i >= 0; i-- { + move := moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } } - n, _ := res.RowsAffected() - for _, h := range hosts { - os.RemoveAll(filepath.Join(t.dir, sanitize(h))) + if len(errs) == 0 && stageDir != "" { + if err := os.RemoveAll(stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) + } } - if n > 0 { - _ = t.gcBlobs() + return errors.Join(errs...) +} + +// reapStage unlinks a committed staging directory in the background. This is the +// step that used to stall the recorder: a legacy tree mirrors the URL path per +// request and can hold hundreds of thousands of small files, and the whole unlink +// ran while the write lock was held. The rows are already gone by now, so losing +// this to a shutdown leaves garbage under _delete_staging, never inconsistent +// state. +func (t *Traffic) reapStage(stageDir string) { + if stageDir == "" { + return } - return n, nil + t.reaping.Go(func() { + if err := os.RemoveAll(stageDir); err != nil { + log.Printf("[traffic] 清理历史流量目录 %s 失败:%v", stageDir, err) + } + }) } // DeleteHostsExact removes recorded exchanges for a set of EXACT hosts (the @@ -503,15 +966,12 @@ func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) { return deleted, nil } -type stagedTrafficPath struct { - source string - staged string -} - -// HostDeleteStage keeps the SQLite delete transaction open and moves matching -// host trees to a same-filesystem staging directory. The Traffic write lock is -// held until Commit or Rollback, so no recorder can recreate a staged tree or a -// blob reference while a task deletion is being coordinated with PostgreSQL. +// HostDeleteStage keeps the SQLite delete transaction open so a task deletion +// can be coordinated with PostgreSQL: the traffic side is staged here and only +// committed once the caller's own transaction has succeeded. The Traffic write +// lock is held until Commit or Rollback, so no recorder can add rows or a blob +// reference for a host that is mid-deletion. Rolling back is just a transaction +// rollback now that nothing is moved on disk. type HostDeleteStage struct { traffic *Traffic tx *sql.Tx @@ -536,10 +996,15 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error stage := &HostDeleteStage{traffic: t} fail := func(cause error) (*HostDeleteStage, error) { if rollbackErr := stage.rollbackLocked(); rollbackErr != nil { - return nil, errors.Join(cause, fmt.Errorf("restore staged traffic: %w", rollbackErr)) + return nil, errors.Join(cause, fmt.Errorf("回滚流量删除: %w", rollbackErr)) } return nil, cause } + abort := func(cause error) (*HostDeleteStage, error) { + t.wmu.Unlock() + stage.done = true + return nil, cause + } unique := make([]string, 0, len(hosts)) seen := make(map[string]struct{}, len(hosts)) @@ -551,55 +1016,35 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error unique = append(unique, host) } + // Exact hosts are known up front, so their directories are derived directly + // rather than looked up: a host whose index rows are already gone can still + // own an orphaned directory that this deletion should take with it. + legacy := make([]string, 0, len(unique)) + for _, h := range unique { + legacy = append(legacy, filepath.Join(t.dir, sanitize(h))) + } + tx, err := t.db.Begin() if err != nil { - t.wmu.Unlock() - stage.done = true - return nil, err + return abort(err) } stage.tx = tx for _, h := range unique { - res, err := tx.Exec(`DELETE FROM exchanges WHERE host=?`, h) + n, err := t.deleteWhere(tx, `host=?`, h) if err != nil { return fail(err) } - n, _ := res.RowsAffected() stage.deleted += n } - stagedSources := make(map[string]struct{}, len(unique)) - for _, host := range unique { - source := filepath.Join(t.dir, sanitize(host)) - if _, duplicatePath := stagedSources[source]; duplicatePath { - continue - } - stagedSources[source] = struct{}{} - if _, err := os.Lstat(source); err != nil { - if os.IsNotExist(err) { - continue - } - return fail(fmt.Errorf("inspect traffic tree for %q: %w", host, err)) - } - if stage.stageDir == "" { - parent := filepath.Join(t.dir, "_delete_staging") - if err := os.MkdirAll(parent, 0o700); err != nil { - return fail(fmt.Errorf("create traffic staging directory: %w", err)) - } - stage.stageDir, err = os.MkdirTemp(parent, "hosts-") - if err != nil { - return fail(fmt.Errorf("create traffic stage: %w", err)) - } - } - staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source))) - if err := os.Rename(source, staged); err != nil { - return fail(fmt.Errorf("stage traffic tree for %q: %w", host, err)) - } - stage.moves = append(stage.moves, stagedTrafficPath{source: source, staged: staged}) + stage.stageDir, stage.moves, err = t.stageTrees(legacy) + if err != nil { + return fail(err) } return stage, nil } -// Rollback restores both the SQLite rows and every staged host tree. +// Rollback discards the staged deletion, leaving the traffic store untouched. func (s *HostDeleteStage) Rollback() error { if s == nil || s.done { return nil @@ -614,26 +1059,11 @@ func (s *HostDeleteStage) rollbackLocked() error { var errs []error if s.tx != nil { if err := s.tx.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) { - errs = append(errs, fmt.Errorf("rollback traffic index: %w", err)) + errs = append(errs, fmt.Errorf("回滚流量索引: %w", err)) } } - for i := len(s.moves) - 1; i >= 0; i-- { - move := s.moves[i] - if _, err := os.Lstat(move.source); err == nil { - errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) - continue - } else if !os.IsNotExist(err) { - errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) - continue - } - if err := os.Rename(move.staged, move.source); err != nil { - errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) - } - } - if len(errs) == 0 && s.stageDir != "" { - if err := os.RemoveAll(s.stageDir); err != nil { - errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) - } + if err := restoreTrees(s.stageDir, s.moves); err != nil { + errs = append(errs, err) } s.done = true s.traffic.wmu.Unlock() @@ -649,57 +1079,98 @@ func (s *HostDeleteStage) Commit() error { if err := s.tx.Commit(); err != nil { // A failed SQLite commit normally rolls the transaction back. Restore the // trees so the traffic store remains internally consistent and recoverable. - restoreErr := s.restoreTreesLocked() + restoreErr := restoreTrees(s.stageDir, s.moves) s.done = true s.traffic.wmu.Unlock() - return errors.Join(fmt.Errorf("commit traffic index deletion: %w", err), restoreErr) + return errors.Join(fmt.Errorf("提交流量索引删除: %w", err), restoreErr) } + // Unlinking the staged trees is what used to hold the write lock for hours; + // it now runs in the background, while collection below only needs them to be + // out of the live tree — which the staging rename already guaranteed. + s.traffic.reapStage(s.stageDir) var errs []error - if s.stageDir != "" { - if err := os.RemoveAll(s.stageDir); err != nil { - errs = append(errs, fmt.Errorf("remove staged traffic trees: %w", err)) - } - } if err := s.traffic.gcBlobs(); err != nil { - errs = append(errs, fmt.Errorf("garbage collect traffic blobs: %w", err)) + errs = append(errs, fmt.Errorf("回收流量 blob: %w", err)) } s.done = true s.traffic.wmu.Unlock() return errors.Join(errs...) } -func (s *HostDeleteStage) restoreTreesLocked() error { - var errs []error - for i := len(s.moves) - 1; i >= 0; i-- { - move := s.moves[i] - if _, err := os.Lstat(move.source); err == nil { - errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) - continue - } else if !os.IsNotExist(err) { - errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) - continue - } - if err := os.Rename(move.staged, move.source); err != nil { - errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) - } - } - if len(errs) == 0 && s.stageDir != "" { - if err := os.RemoveAll(s.stageDir); err != nil { - errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) - } - } - return errors.Join(errs...) -} - -// gcBlobs removes content-addressed blobs in _blobs that no remaining exchange -// tree references. References appear as "@blob sha256:" lines inside the -// tree's request.http/response.http files (see bodyOrBlob). Best-effort: walk -// errors only skip cleanup of the affected files. Callers must hold wmu so this -// never races a concurrent record() writing a fresh blob + tree. +// gcBlobs removes blobs that no remaining exchange references. Live references +// come from blob_refs, so collection is one query plus a sweep of the blob +// directory — no exchange body is ever read. Emptied bucket directories are +// removed as well; the previous file-scanning collector deleted only files and +// left the buckets behind permanently. Best-effort: walk errors only skip the +// affected entries. Callers must hold wmu so this never races a concurrent +// record() writing a fresh blob + reference. func (t *Traffic) gcBlobs() error { refs := make(map[string]struct{}) + rows, err := t.db.Query(`SELECT DISTINCT hash FROM blob_refs`) + if err != nil { + return err + } + for rows.Next() { + var h string + if err := rows.Scan(&h); err != nil { + rows.Close() + return err + } + refs[h] = struct{}{} + } + rows.Close() + if err := rows.Err(); err != nil { + return err + } + if err := t.legacyBlobRefs(refs); err != nil { + return err + } + + root := filepath.Join(t.dir, "_blobs", "sha256") + var buckets []string + if err := filepath.WalkDir(root, func(p string, d fs.DirEntry, err error) error { + if err != nil || d == nil { + return nil + } + if d.IsDir() { + if p != root { + buckets = append(buckets, p) + } + return nil + } + h := strings.TrimSuffix(d.Name(), ".bin") + if _, ok := refs[h]; !ok { + os.Remove(p) + } + return nil + }); err != nil { + return err + } + // Deepest first, so a two-level bucket left by the old layout collapses fully. + // Remove fails harmlessly on a non-empty directory — exactly the guard needed + // to avoid deleting a bucket that still holds live blobs. + sort.Slice(buckets, func(i, j int) bool { return len(buckets[i]) > len(buckets[j]) }) + for _, b := range buckets { + os.Remove(b) + } + return nil +} + +// legacyBlobRefs adds hashes referenced by pre-SQLite exchanges, whose bodies are +// still .http files on disk holding "@blob sha256:" pointers. Without this +// the first collection after the upgrade would delete blobs that history still +// points at. Skipped entirely once no legacy rows remain — the steady state — so +// the tree walk is transitional rather than a permanent cost. +func (t *Traffic) legacyBlobRefs(refs map[string]struct{}) error { + var n int + if err := t.db.QueryRow(`SELECT COUNT(*) FROM exchanges WHERE path<>''`).Scan(&n); err != nil { + return err + } + if n == 0 { + return nil + } blobRe := regexp.MustCompile(`@blob sha256:([0-9a-f]{64})`) - err := filepath.WalkDir(t.dir, func(p string, d fs.DirEntry, err error) error { + return filepath.WalkDir(t.dir, func(p string, d fs.DirEntry, err error) error { if err != nil || d == nil { return nil // skip unreadable entries } @@ -710,7 +1181,7 @@ func (t *Traffic) gcBlobs() error { } return nil } - if n := d.Name(); n != "request.http" && n != "response.http" { + if nm := d.Name(); nm != "request.http" && nm != "response.http" { return nil } b, err := os.ReadFile(p) @@ -722,25 +1193,12 @@ func (t *Traffic) gcBlobs() error { } return nil }) - if err != nil { - return err - } - return filepath.WalkDir(filepath.Join(t.dir, "_blobs", "sha256"), func(p string, d fs.DirEntry, err error) error { - if err != nil || d == nil || d.IsDir() { - return nil - } - h := strings.TrimSuffix(d.Name(), ".bin") - if _, ok := refs[h]; !ok { - os.Remove(p) - } - return nil - }) } // query returns one page of exchange metadata filtered by host and/or a url // substring. Default page size is intentionally small (3) to keep tool results // lightweight and capped at 10; page is 0-based (page*limit offset). -func (t *Traffic) query(host, contains string, page, limit int) ([]ExchangeMeta, error) { +func (t *Traffic) query(host, contains, bodyContains string, page, limit int) ([]ExchangeMeta, error) { if limit <= 0 { limit = 3 } @@ -760,6 +1218,17 @@ func (t *Traffic) query(host, contains string, page, limit int) ([]ExchangeMeta, q += ` AND (url LIKE ? OR url_template LIKE ?)` args = append(args, "%"+contains+"%", "%"+contains+"%") } + if b := strings.TrimSpace(bodyContains); b != "" { + cond, arg, ok := t.ftsFilter(b) + if !ok { + if !t.fts { + return nil, fmt.Errorf("当前实例未启用全文索引,无法按正文搜索") + } + return nil, fmt.Errorf("正文搜索关键词至少需要 %d 个字符(当前 %d 个)", minTrigram, utf8.RuneCountInString(b)) + } + q += ` AND ` + cond + args = append(args, arg) + } q += ` ORDER BY ts DESC LIMIT ? OFFSET ?` args = append(args, limit, page*limit) rows, err := t.db.Query(q, args...) @@ -788,14 +1257,15 @@ func (t *Traffic) Tools() []actool.CoreTool { search := actool.Build(actool.Spec{ Name: "traffic_search", - Description: "查询记录代理已抓取的目标流量(必须指定 host,可再按 URL 子串过滤)。仅返回极轻量索引(id/method/url/status/resp_len),不含任何响应内容。默认只返回 3 条、每页最多 10 条;结果多时用 page 翻页(page=0 起);要看某条的请求/响应原文用 traffic_get(id)。回看已访问资源、找端点先用它,避免重复 curl 同一 URL。", + Description: "查询记录代理已抓取的目标流量(必须指定 host,可再按 URL 子串或正文关键词过滤)。body_contains 会在已抓取的请求/响应头与正文中做全文搜索,支持任意子串和中文(至少 3 个字符),可用来找响应里的密码、密钥、报错、内网地址等。仅返回极轻量索引(id/method/url/status/resp_len),不含任何响应内容。默认只返回 3 条、每页最多 10 条;结果多时用 page 翻页(page=0 起);要看某条的请求/响应原文用 traffic_get(id)。回看已访问资源、找端点先用它,避免重复 curl 同一 URL。", Schema: map[string]any{ "type": "object", "properties": map[string]any{ - "host": map[string]any{"type": "string", "description": "按主机过滤(必填,如 '107.172.96.177:8082')"}, - "contains": map[string]any{"type": "string", "description": "URL 子串过滤(可选,如 'api' / 'login')"}, - "limit": map[string]any{"type": "integer", "description": "每页条数,默认 3,最大 10"}, - "page": map[string]any{"type": "integer", "description": "页码,从 0 开始,默认 0(按 ts 倒序分页)"}, + "host": map[string]any{"type": "string", "description": "按主机过滤(必填,如 '107.172.96.177:8082')"}, + "contains": map[string]any{"type": "string", "description": "URL 子串过滤(可选,如 'api' / 'login')"}, + "body_contains": map[string]any{"type": "string", "description": "正文全文搜索(可选,至少 3 个字符),匹配请求/响应的头与正文,如 'password' / 'root:x:0' / '内网测试'"}, + "limit": map[string]any{"type": "integer", "description": "每页条数,默认 3,最大 10"}, + "page": map[string]any{"type": "integer", "description": "页码,从 0 开始,默认 0(按 ts 倒序分页)"}, }, "required": []any{"host"}, }, @@ -804,6 +1274,7 @@ func (t *Traffic) Tools() []actool.CoreTool { Run: func(_ context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) { var a struct { Host, Contains string + BodyContains string `json:"body_contains"` Limit int Page int } @@ -811,7 +1282,7 @@ func (t *Traffic) Tools() []actool.CoreTool { if strings.TrimSpace(a.Host) == "" { return actool.Errorf("host 为必填参数:请指定要查询的主机(如 '107.172.96.177:8082'),避免全库扫描。"), nil } - rows, err := t.query(a.Host, a.Contains, a.Page, a.Limit) + rows, err := t.query(a.Host, a.Contains, a.BodyContains, a.Page, a.Limit) if err != nil { return actool.Errorf(err.Error()), nil } @@ -855,7 +1326,47 @@ func (t *Traffic) Tools() []actool.CoreTool { return actool.Text("=== REQUEST ===\n" + clip(req, 2500) + "\n\n=== RESPONSE ===\n" + clip(resp, 4000)), nil }, }) - return []actool.CoreTool{search, get} + + blob := actool.Build(actool.Spec{ + Name: "traffic_blob", + Description: "分段读取超大请求/响应体的原文。traffic_get 里显示为 '…[truncated] @blob sha256:' 的部分即存放于此,把该 hash 传进来即可取完整内容。单次最多返回 8KB,用 offset 继续往后读(返回结果会给出总长度)。适合翻阅备份文件、源码泄露、大 JSON 导出等超过内联阈值的响应。", + Schema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "hash": map[string]any{"type": "string", "description": "traffic_get 中 @blob sha256: 后面的 64 位十六进制值"}, + "offset": map[string]any{"type": "integer", "description": "起始字节偏移,默认 0"}, + "length": map[string]any{"type": "integer", "description": "本次读取字节数,默认且最大 8192"}, + }, + "required": []any{"hash"}, + }, + ReadOnly: ro, + Permissions: allow, + Run: func(_ context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) { + var a struct { + Hash string + Offset int64 + Length int64 + } + _ = json.Unmarshal(in, &a) + if a.Length <= 0 || a.Length > maxBlobRead { + a.Length = maxBlobRead + } + data, total, err := t.BlobRange(a.Hash, a.Offset, a.Length) + if err != nil { + return actool.Errorf(err.Error()), nil + } + if len(data) == 0 { + return actool.Text(fmt.Sprintf("偏移 %d 已超出内容长度(总长 %d 字节)。", a.Offset, total)), nil + } + head := fmt.Sprintf("[offset=%d 本次=%d 总长=%d]\n", a.Offset, len(data), total) + if isBinaryBody("", data) { + return actool.Text(head + "二进制内容,以十六进制展示前 512 字节:\n" + hex.EncodeToString(clipBytes(data, 512))), nil + } + return actool.Text(head + truncateUTF8(data, len(data))), nil + }, + }) + + return []actool.CoreTool{search, get, blob} } // SeedToolMetas returns the traffic tools built on a ZERO receiver, for seeding the diff --git a/traffic/traffic_store_test.go b/traffic/traffic_store_test.go new file mode 100644 index 0000000..2d744c6 --- /dev/null +++ b/traffic/traffic_store_test.go @@ -0,0 +1,440 @@ +package traffic + +import ( + "fmt" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + mproxy "github.com/lqqyt2423/go-mitmproxy/proxy" +) + +// flowOpt tweaks the synthetic flow built by newFlow. +type flowOpt func(*mproxy.Flow) + +func withRespType(ct string) flowOpt { + return func(f *mproxy.Flow) { f.Response.Header.Set("Content-Type", ct) } +} + +// newFlow builds the minimal flow record() needs: a request with a URL, method +// and body, plus a response with a status and body. +func newFlow(host, method, path string, reqBody, respBody []byte, opts ...flowOpt) *mproxy.Flow { + u, err := url.Parse("http://" + host + path) + if err != nil { + panic(err) + } + f := &mproxy.Flow{ + Request: &mproxy.Request{ + Method: method, + URL: u, + Proto: "HTTP/1.1", + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: reqBody, + }, + Response: &mproxy.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: respBody, + }, + } + for _, o := range opts { + o(f) + } + return f +} + +func openTraffic(t *testing.T) (*Traffic, string) { + t.Helper() + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { tr.Close() }) + return tr, dir +} + +func onlyExchangeID(t *testing.T, tr *Traffic) string { + t.Helper() + var id string + if err := tr.DB().QueryRow(`SELECT id FROM exchanges`).Scan(&id); err != nil { + t.Fatalf("读取 exchange id: %v", err) + } + return id +} + +// TestRecordKeepsBodiesInIndex is the core of the storage change: a recorded +// exchange produces no per-request directory at all, and its bodies are served +// back out of SQLite. +func TestRecordKeepsBodiesInIndex(t *testing.T) { + tr, dir := openTraffic(t) + + tr.record(newFlow("api.example.com", "POST", "/v1/login", + []byte(`{"user":"admin","password":"P@ssw0rd"}`), + []byte(`{"token":"abc123","note":"内网测试账号"}`))) + + // The URL-mirroring tree is gone: no host directory, no nested path segments. + if _, err := os.Stat(filepath.Join(dir, "api.example.com")); !os.IsNotExist(err) { + t.Fatalf("record 仍在磁盘上创建 host 目录(stat err=%v)", err) + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + for _, e := range entries { + if !strings.HasPrefix(e.Name(), "_") { + t.Fatalf("data 目录下出现非内部目录 %q,说明仍在写文件树", e.Name()) + } + } + + id := onlyExchangeID(t, tr) + req, resp, err := tr.Get(id) + if err != nil { + t.Fatal(err) + } + for _, want := range []string{"POST /v1/login HTTP/1.1", "Host: api.example.com", `"password":"P@ssw0rd"`} { + if !strings.Contains(req, want) { + t.Fatalf("请求原文缺少 %q,实际:\n%s", want, req) + } + } + for _, want := range []string{"HTTP 200", `"token":"abc123"`, "内网测试账号"} { + if !strings.Contains(resp, want) { + t.Fatalf("响应原文缺少 %q,实际:\n%s", want, resp) + } + } +} + +// TestFullTextSearchMatchesBodies covers what the trigram index buys over the +// previous URL-only search: arbitrary substrings and CJK, across request and +// response bodies. +func TestFullTextSearchMatchesBodies(t *testing.T) { + tr, _ := openTraffic(t) + if !tr.fts { + t.Skip("驱动未启用 FTS5") + } + const host = "api.example.com" + tr.record(newFlow(host, "POST", "/v1/login", + []byte(`{"user":"admin","password":"P@ssw0rd"}`), + []byte(`{"token":"abc123","note":"内网测试账号"}`))) + tr.record(newFlow(host, "GET", "/v1/health", nil, []byte(`{"status":"ok"}`))) + + hits := func(term string) int { + t.Helper() + rows, err := tr.query(host, "", term, 0, 10) + if err != nil { + t.Fatalf("按正文搜索 %q 出错:%v", term, err) + } + return len(rows) + } + if n := hits("password"); n != 1 { + t.Fatalf("搜 password 命中 %d 条,应为 1", n) + } + // Substring inside a token — the default unicode61 tokenizer cannot do this. + if n := hits("ssw0r"); n != 1 { + t.Fatalf("搜子串 ssw0r 命中 %d 条,应为 1", n) + } + if n := hits("内网测试"); n != 1 { + t.Fatalf("搜中文命中 %d 条,应为 1", n) + } + if n := hits("nonexistent-marker"); n != 0 { + t.Fatalf("无关关键词命中 %d 条,应为 0", n) + } + + // Too-short terms are reported, not silently treated as "no match". + if _, err := tr.query(host, "", "ab", 0, 10); err == nil { + t.Fatal("两字符正文关键词应返回明确错误") + } +} + +// TestLargeBodySpillsButStaysSearchable is the case that motivated indexing from +// memory: the body lives in the blob store, only a preview is inline, and the +// part past the preview is still findable. +func TestLargeBodySpillsButStaysSearchable(t *testing.T) { + tr, dir := openTraffic(t) + if !tr.fts { + t.Skip("驱动未启用 FTS5") + } + const host = "dump.example.com" + const marker = "DB_PASSWORD=hunter2" + // Marker sits far past blobPreview, so only the full-text index can find it. + big := []byte(strings.Repeat("-- MySQL dump\n", maxInlineBody/14+2000) + marker) + if len(big) <= maxInlineBody+blobPreview { + t.Fatalf("测试数据不够大:%d 字节", len(big)) + } + tr.record(newFlow(host, "GET", "/backup.sql", nil, big, withRespType("application/sql"))) + + // Stored under a single bucket level, named by hash. + var hash string + if err := tr.DB().QueryRow(`SELECT resp_blob FROM exchange_bodies`).Scan(&hash); err != nil { + t.Fatal(err) + } + if len(hash) != 64 { + t.Fatalf("resp_blob=%q,应为 64 位 sha256", hash) + } + blob := filepath.Join(dir, "_blobs", "sha256", hash[:2], hash+".bin") + st, err := os.Stat(blob) + if err != nil { + t.Fatalf("blob 未落盘到单层桶 %s:%v", blob, err) + } + if st.Size() != int64(len(big)) { + t.Fatalf("blob 大小 %d,应为 %d", st.Size(), len(big)) + } + + // The reference is registered, which is what GC consults. + var refs int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM blob_refs WHERE hash=?`, hash).Scan(&refs); err != nil { + t.Fatal(err) + } + if refs != 1 { + t.Fatalf("blob_refs 行数 %d,应为 1", refs) + } + + // Inline: a readable preview plus the pointer, not the whole body. + _, resp, err := tr.Get(onlyExchangeID(t, tr)) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(resp, "-- MySQL dump") { + t.Fatalf("响应缺少头部预览:\n%s", clip(resp, 300)) + } + if !strings.Contains(resp, "@blob sha256:"+hash) { + t.Fatalf("响应缺少 blob 指针:\n%s", clip(resp, 300)) + } + if strings.Contains(resp, marker) { + t.Fatal("预览不应包含超出 blobPreview 的内容") + } + if len(resp) > blobPreview*2 { + t.Fatalf("内联内容 %d 字节,远超预览上限", len(resp)) + } + + // Searchable despite living on disk — the index was fed from memory. + rows, err := tr.query(host, "", marker, 0, 10) + if err != nil { + t.Fatal(err) + } + if len(rows) != 1 { + t.Fatalf("超大正文中的关键词命中 %d 条,应为 1", len(rows)) + } + + // And retrievable in pages. + data, total, err := tr.BlobRange(hash, int64(len(big)-len(marker)), 100) + if err != nil { + t.Fatal(err) + } + if total != int64(len(big)) { + t.Fatalf("BlobRange total=%d,应为 %d", total, len(big)) + } + if string(data) != marker { + t.Fatalf("BlobRange 读到 %q,应为 %q", data, marker) + } + if _, _, err := tr.BlobRange("../../etc/passwd", 0, 10); err == nil { + t.Fatal("非法 hash 应被拒绝") + } +} + +// TestBinaryBodyStaysOutOfIndex keeps the index spend on things worth searching: +// binary payloads contribute nothing but a type tag. +func TestBinaryBodyStaysOutOfIndex(t *testing.T) { + tr, _ := openTraffic(t) + if !tr.fts { + t.Skip("驱动未启用 FTS5") + } + const host = "cdn.example.com" + const marker = "SECRETINIMAGE" + png := append([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00}, + []byte(strings.Repeat("x", maxInlineBody)+marker)...) + tr.record(newFlow(host, "GET", "/logo.png", nil, png, withRespType("image/png"))) + + rows, err := tr.query(host, "", marker, 0, 10) + if err != nil { + t.Fatal(err) + } + if len(rows) != 0 { + t.Fatalf("二进制正文不应进入全文索引,却命中 %d 条", len(rows)) + } + _, resp, err := tr.Get(onlyExchangeID(t, tr)) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(resp, "[binary image/png") || !strings.Contains(resp, "magic=89504e47") { + t.Fatalf("二进制正文应展示类型与魔数,实际:\n%s", clip(resp, 300)) + } +} + +// TestGetFallsBackToLegacyTree keeps pre-migration captures readable: their rows +// carry a path and their bodies are still .http files on disk. +func TestGetFallsBackToLegacyTree(t *testing.T) { + tr, dir := openTraffic(t) + const host = "old.example.com" + const id = "1-0001" + rel := filepath.Join(host, "GET", id) + exDir := filepath.Join(dir, rel) + if err := os.MkdirAll(exDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(exDir, "request.http"), []byte("GET / HTTP/1.1\nHost: old.example.com\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(exDir, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 11, rel); err != nil { + t.Fatal(err) + } + + req, resp, err := tr.Get(id) + if err != nil { + t.Fatalf("历史记录应仍可读取:%v", err) + } + if !strings.Contains(req, "Host: old.example.com") { + t.Fatalf("历史请求原文错误:%q", req) + } + if !strings.Contains(resp, "legacy body") { + t.Fatalf("历史响应原文错误:%q", resp) + } +} + +// TestGCCollectsBlobsAndEmptyBuckets covers both halves of the collector: the +// reference lookup now comes from blob_refs, and emptied buckets are removed +// instead of accumulating forever. +func TestGCCollectsBlobsAndEmptyBuckets(t *testing.T) { + tr, dir := openTraffic(t) + const host = "dump.example.com" + big := []byte(strings.Repeat("A", maxInlineBody+1024)) + tr.record(newFlow(host, "GET", "/big.bin", nil, big, withRespType("application/sql"))) + + var hash string + if err := tr.DB().QueryRow(`SELECT resp_blob FROM exchange_bodies`).Scan(&hash); err != nil { + t.Fatal(err) + } + bucket := filepath.Join(dir, "_blobs", "sha256", hash[:2]) + if _, err := os.Stat(filepath.Join(bucket, hash+".bin")); err != nil { + t.Fatal(err) + } + + if n, err := tr.DeleteHostsExact([]string{host}); err != nil || n != 1 { + t.Fatalf("DeleteHostsExact=(%d,%v),应为 (1,nil)", n, err) + } + if _, err := os.Stat(filepath.Join(bucket, hash+".bin")); !os.IsNotExist(err) { + t.Fatalf("失去引用的 blob 未被回收:%v", err) + } + if _, err := os.Stat(bucket); !os.IsNotExist(err) { + t.Fatalf("空桶目录未被清理:%v", err) + } + // Bodies and full-text rows go with the exchange. + for _, q := range []string{ + `SELECT COUNT(*) FROM exchange_bodies`, + `SELECT COUNT(*) FROM blob_refs`, + } { + var c int + if err := tr.DB().QueryRow(q).Scan(&c); err != nil { + t.Fatal(err) + } + if c != 0 { + t.Fatalf("%s = %d,应为 0", q, c) + } + } + if tr.fts { + var c int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM ex_fts WHERE ex_fts MATCH ?`, ftsQuote("AAAA")).Scan(&c); err != nil { + t.Fatal(err) + } + if c != 0 { + t.Fatalf("全文索引残留 %d 条", c) + } + } +} + +// TestPageSearchesBodies checks the UI-facing search box picks up the full-text +// index too, not just metadata columns. +func TestPageSearchesBodies(t *testing.T) { + tr, _ := openTraffic(t) + if !tr.fts { + t.Skip("驱动未启用 FTS5") + } + tr.record(newFlow("api.example.com", "POST", "/v1/login", nil, []byte(`{"error":"invalid credentials"}`))) + tr.record(newFlow("api.example.com", "GET", "/v1/health", nil, []byte(`{"status":"ok"}`))) + + rows, total, err := tr.Page("", "", "invalid credentials", 0, 100) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(rows) != 1 { + t.Fatalf("正文关键词命中 total=%d rows=%d,应为 1/1", total, len(rows)) + } + // Metadata matching still works alongside it. + if _, total, err := tr.Page("", "", "health", 0, 100); err != nil || total != 1 { + t.Fatalf("URL 关键词 total=%d err=%v,应为 1", total, err) + } +} + +// TestTruncateUTF8 guards the preview cut: never split a multi-byte rune. +func TestTruncateUTF8(t *testing.T) { + s := "内网测试账号" + for n := 0; n <= len(s); n++ { + got := truncateUTF8([]byte(s), n) + if !strings.HasPrefix(s, got) { + t.Fatalf("n=%d 截断结果 %q 不是原串前缀", n, got) + } + if len(got) > n { + t.Fatalf("n=%d 截断后 %d 字节,超出上限", n, len(got)) + } + } + if got := truncateUTF8([]byte("abc"), 10); got != "abc" { + t.Fatalf("短于上限时应原样返回,得到 %q", got) + } +} + +// TestIsBinaryBody documents the classification: declared binary types, and the +// NUL backstop for anything mislabeled. +func TestIsBinaryBody(t *testing.T) { + cases := []struct { + ct string + body string + want bool + }{ + {"application/json", `{"a":1}`, false}, + {"text/html; charset=utf-8", "", false}, + {"application/sql", "-- dump", false}, + {"", "plain text", false}, + {"image/png", "whatever", true}, + {"application/zip", "PK", true}, + {"APPLICATION/PDF", "%PDF", true}, + {"text/plain", "has\x00nul", true}, + } + for _, c := range cases { + if got := isBinaryBody(c.ct, []byte(c.body)); got != c.want { + t.Errorf("isBinaryBody(%q, %q)=%v,应为 %v", c.ct, c.body, got, c.want) + } + } +} + +// TestRecordConcurrent exercises the write path under contention: ids stay +// unique and every exchange lands in all three tables. +func TestRecordConcurrent(t *testing.T) { + tr, _ := openTraffic(t) + const n = 50 + var wg sync.WaitGroup + for i := range n { + wg.Go(func() { + tr.record(newFlow("api.example.com", "GET", fmt.Sprintf("/item/%d", i), + nil, fmt.Appendf(nil, `{"id":%d}`, i))) + }) + } + wg.Wait() + var exchanges, bodies int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges`).Scan(&exchanges); err != nil { + t.Fatal(err) + } + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchange_bodies`).Scan(&bodies); err != nil { + t.Fatal(err) + } + if exchanges != n || bodies != n { + t.Fatalf("并发写入后 exchanges=%d bodies=%d,应各为 %d", exchanges, bodies, n) + } +}