diff --git a/.actrc b/.actrc new file mode 100644 index 00000000..820fe761 --- /dev/null +++ b/.actrc @@ -0,0 +1,4 @@ +--container-architecture=linux/amd64 +-P ubuntu-22.04=ghcr.io/catthehacker/ubuntu:act-22.04 +--pull=false +--container-options=--init diff --git a/aop/traffic/body.go b/aop/traffic/body.go new file mode 100644 index 00000000..0570c3c9 --- /dev/null +++ b/aop/traffic/body.go @@ -0,0 +1,200 @@ +package traffic + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "hash" + "io" + "os" + "path/filepath" + "sync" +) + +// BodyRef describes a body kept outside the hot Exchange value. Body contains +// only the configured preview; callers that need the complete payload can +// hydrate it from Path. +type BodyRef struct { + Path string `json:"path,omitempty"` + Size int64 `json:"size"` + SHA256 string `json:"sha256,omitempty"` + Complete bool `json:"complete"` + Truncated bool `json:"truncated,omitempty"` +} + +// BodySink writes a body to a temporary file while retaining a small preview. +// It is safe for a single reader goroutine and Close is idempotent. +type BodySink struct { + mu sync.Mutex + partPath string + finalPath string + file *os.File + hash hash.Hash + size int64 + preview []byte + previewMax int + err error + closed bool + complete bool +} + +// NewBodySink creates .part in dir and atomically publishes it as name +// when Close(true) succeeds. The directory is created when needed. +func NewBodySink(dir, name string, previewMax int) (*BodySink, error) { + if dir == "" { + return nil, fmt.Errorf("traffic: body directory is empty") + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("traffic: create body directory: %w", err) + } + finalPath := filepath.Join(dir, name) + partPath := finalPath + ".part" + f, err := os.OpenFile(partPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600) + if err != nil { + return nil, fmt.Errorf("traffic: create body file: %w", err) + } + if previewMax < 0 { + previewMax = 0 + } + h := sha256.New() + return &BodySink{ + partPath: partPath, + finalPath: finalPath, + file: f, + hash: h, + previewMax: previewMax, + }, nil +} + +// Write appends p to the body file and updates its preview and digest. +func (s *BodySink) Write(p []byte) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + if s.err != nil { + return 0, s.err + } + return 0, io.ErrClosedPipe + } + if s.err != nil { + return 0, s.err + } + n, err := s.file.Write(p) + if n > 0 { + s.size += int64(n) + _, _ = s.hash.Write(p[:n]) + if len(s.preview) < s.previewMax { + end := len(p) + if remaining := s.previewMax - len(s.preview); end > remaining { + end = remaining + } + s.preview = append(s.preview, p[:end]...) + } + } + if err != nil { + s.err = err + } + return n, err +} + +// Reader wraps r so body bytes are captured as they pass through the proxy. +func (s *BodySink) Reader(r io.Reader) io.Reader { + return &bodyReader{src: r, sink: s} +} + +type bodyReader struct { + src io.Reader + sink *BodySink +} + +func (r *bodyReader) Read(p []byte) (int, error) { + n, err := r.src.Read(p) + if n > 0 { + // Capture failures are recorded by BodySink but deliberately do not + // change the upstream reader result: observation must not alter the + // request/response being proxied. + _, _ = r.sink.Write(p[:n]) + } + return n, err +} + +// Close finishes the body. A failed or incomplete capture keeps its .part +// file for diagnosis and reports Complete=false in the returned reference. +func (s *BodySink) Close(complete bool) (BodyRef, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return s.refLocked(), s.err + } + s.closed = true + s.complete = complete && s.err == nil + if s.file != nil { + if err := s.file.Sync(); err != nil && s.err == nil { + s.err = err + s.complete = false + } + if err := s.file.Close(); err != nil && s.err == nil { + s.err = err + s.complete = false + } + } + if s.complete { + if err := os.Rename(s.partPath, s.finalPath); err != nil { + s.err = err + s.complete = false + } + } + return s.refLocked(), s.err +} + +// Preview returns a copy of the bytes retained for list/detail summaries. +func (s *BodySink) Preview() []byte { + s.mu.Lock() + defer s.mu.Unlock() + return append([]byte(nil), s.preview...) +} + +// Discard closes and removes a capture that was filtered out before it became +// a visible Exchange. +func (s *BodySink) Discard() error { + s.mu.Lock() + defer s.mu.Unlock() + if !s.closed { + s.closed = true + _ = s.file.Close() + } + if err := os.Remove(s.partPath); err != nil && !os.IsNotExist(err) { + return err + } + if err := os.Remove(s.finalPath); err != nil && !os.IsNotExist(err) { + return err + } + return nil +} + +func (s *BodySink) refLocked() BodyRef { + digest := "" + if s.hash != nil { + digest = hex.EncodeToString(s.hash.Sum(nil)) + } + path := s.partPath + if s.complete { + path = s.finalPath + } + return BodyRef{ + Path: path, + Size: s.size, + SHA256: digest, + Complete: s.complete, + Truncated: false, + } +} + +// ReadBody loads a body reference on demand. It intentionally does not cache +// the bytes in the reference so callers control the resulting allocation. +func ReadBody(ref *BodyRef) ([]byte, error) { + if ref == nil || ref.Path == "" { + return nil, nil + } + return os.ReadFile(ref.Path) +} diff --git a/aop/traffic/body_test.go b/aop/traffic/body_test.go new file mode 100644 index 00000000..b28f800a --- /dev/null +++ b/aop/traffic/body_test.go @@ -0,0 +1,29 @@ +package traffic + +import ( + "io" + "strings" + "testing" +) + +func TestBodySinkStreamsAndHydrates(t *testing.T) { + dir := t.TempDir() + sink, err := NewBodySink(dir, "response", 4) + if err != nil { + t.Fatal(err) + } + if _, err := io.Copy(sink, strings.NewReader("abcdefgh")); err != nil { + t.Fatal(err) + } + ref, err := sink.Close(true) + if err != nil { + t.Fatal(err) + } + if ref.Size != 8 || string(sink.Preview()) != "abcd" || !ref.Complete { + t.Fatalf("unexpected ref: %+v preview=%q", ref, sink.Preview()) + } + body, err := ReadBody(&ref) + if err != nil || string(body) != "abcdefgh" { + t.Fatalf("body=%q err=%v", body, err) + } +} diff --git a/aop/traffic/exchange.go b/aop/traffic/exchange.go index f625ad03..e788f9d5 100644 --- a/aop/traffic/exchange.go +++ b/aop/traffic/exchange.go @@ -2,7 +2,11 @@ package traffic import ( "encoding/json" + "net/http" "sort" + "strconv" + "strings" + "time" ) // Pair is one HTTP header line: flat, ordered, duplicates preserved. It is the @@ -20,6 +24,7 @@ type Request struct { Protocol string Headers []Pair Body []byte + BodyRef *BodyRef `json:"-"` } // Response is the response half of an exchange. It is optional on Exchange: a @@ -30,6 +35,32 @@ type Response struct { ReasonPhrase string Headers []Pair Body []byte + BodyRef *BodyRef `json:"-"` +} + +// HydrateBodies loads file-backed request/response bodies into Body. It is +// intentionally explicit so list/query paths do not allocate large payloads. +func (e *Exchange) HydrateBodies() error { + if e == nil { + return nil + } + if e.Request.BodyRef != nil { + body, err := ReadBody(e.Request.BodyRef) + if err != nil { + return err + } + e.Request.Body = body + e.Request.BodyRef = nil + } + if e.Response != nil && e.Response.BodyRef != nil { + body, err := ReadBody(e.Response.BodyRef) + if err != nil { + return err + } + e.Response.Body = body + e.Response.BodyRef = nil + } + return nil } // Exchange is the canonical in-memory form of one captured HTTP exchange, @@ -49,6 +80,90 @@ type Exchange struct { Complete bool } +// Clone returns an independent exchange value, including response metadata and +// body references. The proxy hot store uses it before hydrating a body so a +// query or subscriber never mutates the retained preview under a read lock. +func (e Exchange) Clone() Exchange { + out := e + out.Request.Headers = append([]Pair(nil), e.Request.Headers...) + out.Request.Body = append([]byte(nil), e.Request.Body...) + if e.Request.BodyRef != nil { + ref := *e.Request.BodyRef + out.Request.BodyRef = &ref + } + if e.Response != nil { + resp := *e.Response + resp.Headers = append([]Pair(nil), e.Response.Headers...) + resp.Body = append([]byte(nil), e.Response.Body...) + if e.Response.BodyRef != nil { + ref := *e.Response.BodyRef + resp.BodyRef = &ref + } + out.Response = &resp + } + return out +} + +// ExchangeFromHTTP converts the standard library's request/response pair into +// the canonical HTTP observation model. Callers provide body bytes explicitly +// because the http bodies are streaming and may already have been consumed by +// the caller (for example, by a file-backed recorder). +func ExchangeFromHTTP(req *http.Request, resp *http.Response, requestBody, responseBody []byte) *Exchange { + e := &Exchange{} + if req != nil { + urlString := "" + if req.URL != nil { + urlString = req.URL.String() + } + e.Request = Request{ + Method: req.Method, + URL: urlString, + Protocol: req.Proto, + Headers: PairsFromHTTP(req.Header), + Body: requestBody, + } + } + if resp != nil { + reason := resp.Status + if prefix := strconv.Itoa(resp.StatusCode) + " "; strings.HasPrefix(reason, prefix) { + reason = strings.TrimPrefix(reason, prefix) + } + e.Response = &Response{ + StatusCode: resp.StatusCode, + ReasonPhrase: reason, + Headers: PairsFromHTTP(resp.Header), + Body: responseBody, + } + e.Complete = true + } + return e +} + +// WebSocketMessage is a single message observed after an HTTP WebSocket +// handshake. WebSocket traffic is deliberately modeled separately from an +// HTTP Exchange while sharing the same header pair representation. +type WebSocketMessage struct { + Direction string + Type string + Body []byte + Timestamp time.Time +} + +// WebSocketExchange contains the handshake metadata and message stream for a +// WebSocket connection. The HTTP handshake itself can still be represented by +// Exchange; this type is for the bidirectional messages that follow it. +type WebSocketExchange struct { + ID string + URL string + Protocol string + Headers []Pair + Messages []WebSocketMessage + StartTime time.Time + EndTime time.Time + Complete bool + Error string +} + // exchangeJSON is the persisted shape: identical field names and order to the // http.exchange.v1 flow element, headers as a name→values map. type exchangeJSON struct { @@ -252,6 +367,26 @@ func pairsFromProto(headers []*Header) []Pair { return out } +// PairsFromHTTP converts net/http headers into the canonical deterministic +// pair sequence used by Exchange and Flow. +func PairsFromHTTP(headers http.Header) []Pair { + if len(headers) == 0 { + return nil + } + names := make([]string, 0, len(headers)) + for name := range headers { + names = append(names, name) + } + sort.Strings(names) + out := make([]Pair, 0, len(headers)) + for _, name := range names { + for _, value := range headers[name] { + out = append(out, Pair{Name: name, Value: value}) + } + } + return out +} + func pairsToProto(pairs []Pair) []*Header { if len(pairs) == 0 { return nil diff --git a/aop/traffic/exchange_test.go b/aop/traffic/exchange_test.go index 978078ce..6b64c022 100644 --- a/aop/traffic/exchange_test.go +++ b/aop/traffic/exchange_test.go @@ -2,9 +2,24 @@ package traffic import ( "encoding/json" + "net/http" + "net/url" "testing" ) +func TestExchangeFromHTTPUsesCanonicalPairs(t *testing.T) { + u, _ := url.Parse("https://example.test/a") + req := &http.Request{Method: "POST", URL: u, Proto: "HTTP/1.1", Header: http.Header{"X-Test": {"a", "b"}}} + resp := &http.Response{StatusCode: 201, Status: "201 Created", Header: http.Header{"Content-Type": {"application/json"}}} + e := ExchangeFromHTTP(req, resp, []byte("req"), []byte("resp")) + if e.Request.Method != "POST" || e.Request.URL != u.String() || !e.Complete { + t.Fatalf("unexpected exchange: %+v", e) + } + if len(e.Request.Headers) != 2 || e.Response.StatusCode != 201 || string(e.Response.Body) != "resp" { + t.Fatalf("unexpected canonical exchange: %+v", e) + } +} + func TestFlowExchangeRoundTrip(t *testing.T) { flow := &Flow{ Id: "flow-1", diff --git a/tools/curl/client.go b/tools/curl/client.go index 512f1a1f..a89952d3 100644 --- a/tools/curl/client.go +++ b/tools/curl/client.go @@ -25,6 +25,7 @@ import ( "time" toolpb "github.com/chainreactors/aiscan/aop/tool" + traffic "github.com/chainreactors/aiscan/aop/traffic" ) // A single stable, modern Chrome identity. Keeping one fingerprint per process @@ -208,7 +209,7 @@ func (c *Command) do(ctx context.Context, req *Request, env map[string]string, w fmt.Fprint(stdout, expandWriteOut(req.WriteOut, resp, written)) } - c.emitArtifact(ctx, resp, written) + c.emitArtifact(ctx, traffic.ExchangeFromHTTP(resp.Request, resp, nil, nil), written) return nil } @@ -219,7 +220,7 @@ func (c *Command) failResponse(ctx context.Context, client *http.Client, req *Re if req.WriteOut != "" { fmt.Fprint(stdout, expandWriteOut(req.WriteOut, resp, 0)) } - c.emitArtifact(ctx, resp, 0) + c.emitArtifact(ctx, traffic.ExchangeFromHTTP(resp.Request, resp, nil, nil), 0) return fmt.Errorf("curl: (22) The requested URL returned error: %s", resp.Status) } @@ -861,8 +862,8 @@ func resolveKey(host, port string) string { return host + ":" + strings.TrimSpace(port) } -func (c *Command) emitArtifact(ctx context.Context, resp *http.Response, size int64) { - if c.Events == nil || resp.Request == nil { +func (c *Command) emitArtifact(ctx context.Context, exchange *traffic.Exchange, size int64) { + if c.Events == nil || exchange == nil || exchange.Response == nil { return } summary := struct { @@ -871,14 +872,23 @@ func (c *Command) emitArtifact(ctx context.Context, resp *http.Response, size in ContentType string `json:"content_type,omitempty"` Size int64 `json:"size"` }{ - URL: resp.Request.URL.String(), - Status: resp.StatusCode, - ContentType: resp.Header.Get("Content-Type"), + URL: exchange.Request.URL, + Status: exchange.Response.StatusCode, + ContentType: headerValue(exchange.Response.Headers, "Content-Type"), Size: size, } c.EmitArtifactCtx(ctx, "curl", toolpb.ArtifactKindWeb, summary.URL, summary) } +func headerValue(headers []traffic.Pair, name string) string { + for _, h := range headers { + if strings.EqualFold(h.Name, name) { + return h.Value + } + } + return "" +} + // resolvePath anchors a relative file path at the tool's working directory, // matching how the other scanners treat file arguments. func resolvePath(workDir, path string) string { diff --git a/tools/proxy/hub.go b/tools/proxy/hub.go index 9413f282..0d685f53 100644 --- a/tools/proxy/hub.go +++ b/tools/proxy/hub.go @@ -8,6 +8,7 @@ import ( "net/http" "os" "path/filepath" + "strings" "sync" "sync/atomic" "time" @@ -45,13 +46,17 @@ type ProxyHub struct { // connections while in-flight children are undisturbed. recording atomic.Bool decrypt atomic.Bool + filterMu sync.RWMutex + filter QueryOpts subsMu sync.Mutex - subs map[int]chan *traffic.Flow + subs map[int]*flowSubscriber nextSub int } -const hubStreamLargeBodies = 10 * 1024 * 1024 +// Keep proxy-side buffering bounded. Bodies at or above this threshold are +// captured through the recorder reader and written to disk incrementally. +const hubStreamLargeBodies = 64 * 1024 // NewProxyHub builds the hub around an existing State (egress source of truth) // and FlowStore (capture sink). Both are owned by the caller so the mitm query @@ -68,7 +73,7 @@ func NewProxyHub(state *State, store *FlowStore, caRootPath string, capture bool if store == nil { store = NewFlowStore(10000) } - h := &ProxyHub{state: state, store: store, subs: make(map[int]chan *traffic.Flow)} + h := &ProxyHub{state: state, store: store, subs: make(map[int]*flowSubscriber)} // The CA path is always prepared so capture can be toggled on at runtime; // CAPath only advertises it to children while interception is actually on. h.caPath = filepath.Join(caRootPath, "mitmproxy-ca-cert.pem") @@ -90,6 +95,50 @@ func (h *ProxyHub) SetCapture(record, decryptHTTPS bool) { h.decrypt.Store(decryptHTTPS) } +// SetCaptureFilter applies the existing traffic FlowFilter before a flow is +// stored or published. It deliberately lives on the hub so filtering reduces +// both memory/disk work and subscriber traffic. +func (h *ProxyHub) SetCaptureFilter(filter *traffic.FlowFilter) { + h.filterMu.Lock() + defer h.filterMu.Unlock() + if filter == nil { + h.filter = QueryOpts{} + return + } + h.filter = QueryOpts{Host: filter.GetHost(), Status: filter.GetStatus(), CType: filter.GetType()} +} + +func (h *ProxyHub) captureMatches(flow Flow) bool { + h.filterMu.RLock() + f := h.filter + h.filterMu.RUnlock() + if f.Host != "" && !strings.Contains(strings.ToLower(flow.Host), strings.ToLower(f.Host)) { + return false + } + if f.Status != "" && (flow.Response == nil || !matchStatus(flow.Response.StatusCode, f.Status)) { + return false + } + if f.CType != "" && !strings.Contains(strings.ToLower(flow.ContentType), strings.ToLower(f.CType)) { + return false + } + return true +} + +func (h *ProxyHub) captureHostAllowed(host string) bool { + h.filterMu.RLock() + hostFilter := h.filter.Host + h.filterMu.RUnlock() + return hostFilter == "" || strings.Contains(strings.ToLower(host), strings.ToLower(hostFilter)) +} + +func (h *ProxyHub) captureResponseAllowed(status int, contentType string) bool { + h.filterMu.RLock() + f := h.filter + h.filterMu.RUnlock() + return (f.Status == "" || matchStatus(status, f.Status)) && + (f.CType == "" || strings.Contains(strings.ToLower(contentType), strings.ToLower(f.CType))) +} + // Start brings up the MITM listener on an ephemeral loopback port and exports // the CA certificate so external processes can trust intercepted HTTPS. It is // idempotent: repeated calls return the first outcome. @@ -105,6 +154,11 @@ func (h *ProxyHub) start(caRootPath string) error { if err := os.MkdirAll(caRootPath, 0o755); err != nil { return fmt.Errorf("proxy hub: create CA dir: %w", err) } + if h.store != nil { + if err := h.store.SetBodyDir(filepath.Join(caRootPath, "capture")); err != nil { + return fmt.Errorf("proxy hub: create capture dir: %w", err) + } + } } server, err := mitmproxy.NewProxy(&mitmproxy.Options{ Addr: "127.0.0.1:0", @@ -200,11 +254,15 @@ func (h *ProxyHub) CAPath() string { // Shutdown stops the listener. Safe to call on a never-started hub. func (h *ProxyHub) Shutdown(ctx context.Context) { + h.closeSubscribers() h.mu.Lock() server := h.server h.server = nil h.mu.Unlock() if server == nil { + if h.store != nil { + _ = h.store.Close() + } return } if ctx == nil { @@ -213,4 +271,17 @@ func (h *ProxyHub) Shutdown(ctx context.Context) { defer cancel() } _ = server.Shutdown(ctx) + if h.store != nil { + _ = h.store.Close() + } +} + +func (h *ProxyHub) closeSubscribers() { + h.subsMu.Lock() + subs := h.subs + h.subs = make(map[int]*flowSubscriber) + for _, subscriber := range subs { + close(subscriber.done) + } + h.subsMu.Unlock() } diff --git a/tools/proxy/hub_traffic.go b/tools/proxy/hub_traffic.go index b332ddfe..1824f096 100644 --- a/tools/proxy/hub_traffic.go +++ b/tools/proxy/hub_traffic.go @@ -1,6 +1,7 @@ package proxy import ( + "strconv" "sync" traffic "github.com/chainreactors/aiscan/aop/traffic" @@ -13,6 +14,9 @@ func (h *ProxyHub) ingest(flow Flow) { if !h.recording.Load() { return } + if !h.captureMatches(flow) { + return + } stored := h.store.Add(flow) h.publish(&stored) } @@ -22,33 +26,46 @@ func (h *ProxyHub) publish(flow *Flow) { return } h.subsMu.Lock() - if len(h.subs) == 0 { - h.subsMu.Unlock() - return - } - message := flowToProto(flow) for _, subscriber := range h.subs { - select { - case subscriber <- message: - default: - } + // The capture path never sends into a subscriber's bounded output + // channel. A single wake-up is enough: the subscriber owns a cursor + // and drains every FlowStore entry after it, in order. This keeps a + // slow Cairn connection from silently dropping observations or + // blocking the proxy response path. + subscriber.signal() } h.subsMu.Unlock() } func (h *ProxyHub) Subscribe(buffer int) (<-chan *traffic.Flow, func()) { + return h.SubscribeFrom(h.store.Sequence(), buffer) +} + +// SubscribeFrom starts a reliable FlowStore-backed subscription after the +// supplied numeric flow id. It is useful for reconnecting consumers that have +// persisted their last seen id. The normal Subscribe path starts at the +// current tail and observes only new flows. +func (h *ProxyHub) SubscribeFrom(after int, buffer int) (<-chan *traffic.Flow, func()) { if buffer <= 0 { buffer = 256 } channel := make(chan *traffic.Flow, buffer) + subscriber := &flowSubscriber{ + hub: h, + out: channel, + wake: make(chan struct{}, 1), + done: make(chan struct{}), + cursor: after, + } h.subsMu.Lock() if h.subs == nil { - h.subs = make(map[int]chan *traffic.Flow) + h.subs = make(map[int]*flowSubscriber) } id := h.nextSub h.nextSub++ - h.subs[id] = channel + h.subs[id] = subscriber h.subsMu.Unlock() + go subscriber.run() var once sync.Once cancel := func() { @@ -56,7 +73,7 @@ func (h *ProxyHub) Subscribe(buffer int) (<-chan *traffic.Flow, func()) { h.subsMu.Lock() if existing, ok := h.subs[id]; ok { delete(h.subs, id) - close(existing) + close(existing.done) } h.subsMu.Unlock() }) @@ -64,6 +81,54 @@ func (h *ProxyHub) Subscribe(buffer int) (<-chan *traffic.Flow, func()) { return channel, cancel } +// flowSubscriber turns a store cursor into the historical channel API. The +// output remains bounded; back-pressure is isolated to this worker and never +// reaches the MITM request/response callbacks. +type flowSubscriber struct { + hub *ProxyHub + out chan *traffic.Flow + wake chan struct{} + done chan struct{} + cursor int +} + +func (s *flowSubscriber) signal() { + select { + case s.wake <- struct{}{}: + default: + } +} + +func (s *flowSubscriber) run() { + defer close(s.out) + for { + flows := s.hub.store.after(s.cursor) + for i := range flows { + flow := flows[i] + message := flowToProto(&flow) + select { + case s.out <- message: + s.cursor = flowSequence(flow.ID) + case <-s.done: + return + } + } + select { + case <-s.done: + return + case <-s.wake: + } + } +} + +func flowSequence(id string) int { + seq, err := strconv.Atoi(id) + if err != nil || seq < 0 { + return 0 + } + return seq +} + // flowToProto renders a stored flow as a wire Flow: the exchange semantics go // through the canonical Exchange, attribution (tool id, timestamp) is stamped // on top. @@ -71,7 +136,12 @@ func flowToProto(flow *Flow) *traffic.Flow { if flow == nil { return nil } - message := flow.Proto() + // The hot store keeps only a preview and a file reference. A wire Flow + // retains the historical bytes field, so hydrate only at this boundary. + copy := *flow + copy.Exchange = flow.Clone() + _ = copy.HydrateBodies() + message := copy.Proto() message.ToolId = flow.ToolID if !flow.Timestamp.IsZero() { message.Timestamp = timestamppb.New(flow.Timestamp) diff --git a/tools/proxy/hub_traffic_test.go b/tools/proxy/hub_traffic_test.go index ffa82755..06142043 100644 --- a/tools/proxy/hub_traffic_test.go +++ b/tools/proxy/hub_traffic_test.go @@ -2,11 +2,15 @@ package proxy import ( "context" + "fmt" "io" "net/http" "net/url" + "strconv" "testing" "time" + + traffic "github.com/chainreactors/aiscan/aop/traffic" ) // hubClient builds an HTTP client that routes through the hub with callID as the @@ -122,3 +126,84 @@ func TestHubSubscribe(t *testing.T) { t.Fatal("no flow received on subscription") } } + +func TestHubSubscribeDoesNotDropWhenConsumerIsSlow(t *testing.T) { + hub := NewProxyHub(NewState(""), NewFlowStore(128), "", true) + ch, cancel := hub.Subscribe(1) + defer cancel() + + // Do not read while publishing. The old implementation filled the channel + // and silently discarded every flow after the first one; the store-backed + // subscriber only records a wake-up and drains its cursor in order. + for i := 1; i <= 64; i++ { + hub.ingest(Flow{Exchange: traffic.Exchange{ + ID: fmt.Sprintf("raw-%d", i), + Request: traffic.Request{Method: "GET", URL: "https://example.test/"}, + Response: &traffic.Response{StatusCode: 200}, + }, ToolID: "tool"}) + } + + for i := 1; i <= 64; i++ { + select { + case got := <-ch: + if got == nil || got.GetId() != strconv.Itoa(i) { + t.Fatalf("flow %d = %#v, want sequential id %d", i, got, i) + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for flow %d", i) + } + } +} + +func TestHubSubscribeFromReplaysRetainedFlows(t *testing.T) { + hub := NewProxyHub(NewState(""), NewFlowStore(8), "", true) + for i := 1; i <= 3; i++ { + hub.ingest(Flow{Exchange: traffic.Exchange{ + Request: traffic.Request{Method: "GET", URL: "https://example.test/"}, + Response: &traffic.Response{StatusCode: 200}, + }}) + } + ch, cancel := hub.SubscribeFrom(1, 2) + defer cancel() + for want := 2; want <= 3; want++ { + select { + case got := <-ch: + if got == nil || got.GetId() != strconv.Itoa(want) { + t.Fatalf("replayed flow = %#v, want id %d", got, want) + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for replayed flow %d", want) + } + } +} + +func TestFlowStoreReloadsMetadataIndexWithoutHydratingBodies(t *testing.T) { + dir := t.TempDir() + first := NewFlowStore(8) + if err := first.SetBodyDir(dir); err != nil { + t.Fatal(err) + } + first.Add(Flow{ + ToolID: "call-1", Host: "example.test", ContentType: "text/plain", + Exchange: traffic.Exchange{ + Request: traffic.Request{Method: "GET", URL: "https://example.test/"}, + Response: &traffic.Response{StatusCode: 200, BodyRef: &traffic.BodyRef{Path: "body/1.resp", Size: 10, Complete: true}}, + }, + }) + if err := first.Close(); err != nil { + t.Fatal(err) + } + + second := NewFlowStore(8) + if err := second.SetBodyDir(dir); err != nil { + t.Fatal(err) + } + defer second.Close() + flows := second.Query(QueryOpts{}) + if len(flows) != 1 || flows[0].ToolID != "call-1" { + t.Fatalf("reloaded flows = %#v", flows) + } + if flows[0].Response == nil || flows[0].Response.BodyRef == nil || len(flows[0].Response.Body) != 0 { + t.Fatalf("reloaded body metadata = %#v", flows[0].Response) + } +} diff --git a/tools/proxy/mitm.go b/tools/proxy/mitm.go index 5ea2b229..7e9bfdb2 100644 --- a/tools/proxy/mitm.go +++ b/tools/proxy/mitm.go @@ -2,9 +2,12 @@ package proxy import ( "context" + "encoding/json" "fmt" + "io" "net/http" - "sort" + "os" + "path/filepath" "strconv" "strings" "sync" @@ -173,7 +176,7 @@ const maxBodySnip = 4096 type captureAddon struct { mitmproxy.BaseAddon hub *ProxyHub - pending sync.Map + pending sync.Map // map[proxy flow id]*captureState } // toolIDOf returns the AOP tool-call id that opened this flow's connection, read @@ -187,99 +190,381 @@ func toolIDOf(f *mitmproxy.Flow) string { } func (a *captureAddon) Requestheaders(f *mitmproxy.Flow) { - a.pending.Store(f.Id.String(), time.Now()) + if a.hub == nil || !a.hub.recording.Load() || f == nil || f.Request == nil { + return + } + if f.Request.URL != nil && !a.hub.captureHostAllowed(f.Request.URL.Hostname()) { + return + } + // Successful CONNECT and WebSocket handshakes are connection-level + // lifecycles. Inner HTTPS requests and the separate WebSocket recorder are + // responsible for their own records; retaining this outer request would + // otherwise leak a pending capture until the process exits. + if strings.EqualFold(f.Request.Method, http.MethodConnect) || + strings.EqualFold(f.Request.Header.Get("Upgrade"), "websocket") { + return + } + state := newCaptureState(a.hub, f) + state.owner = a + a.pending.Store(f.Id.String(), state) } -func (a *captureAddon) Response(f *mitmproxy.Flow) { - var dur time.Duration - if start, ok := a.pending.LoadAndDelete(f.Id.String()); ok { - if t, ok := start.(time.Time); ok { - dur = time.Since(t) - } - } - flow := Flow{ - Exchange: traffic.Exchange{ - Request: traffic.Request{ - Method: f.Request.Method, - URL: f.Request.URL.String(), - Protocol: f.Request.Proto, - Headers: pairsFromHTTP(f.Request.Header), - }, - }, - Timestamp: f.StartTime, - ToolID: toolIDOf(f), - Host: f.Request.URL.Hostname(), - Duration: dur, - TLS: f.ConnContext.ClientConn.Tls, +func (a *captureAddon) Request(f *mitmproxy.Flow) { + if state := a.state(f); state != nil && f.Request != nil { + state.setRequestBody(f.Request.Body) } - if len(f.Request.Body) > 0 { - flow.Request.Body = snip(f.Request.Body, maxBodySnip) +} + +func (a *captureAddon) Responseheaders(f *mitmproxy.Flow) { + if state := a.state(f); state != nil && f.Response != nil { + if !a.hub.captureResponseAllowed(f.Response.StatusCode, f.Response.Header.Get("Content-Type")) { + state.discard() + return + } + state.setResponseMeta(f.Response.StatusCode, f.Response.Header) } - if f.Response != nil { - flow.Response = &traffic.Response{ - StatusCode: f.Response.StatusCode, - Headers: pairsFromHTTP(f.Response.Header), +} + +func (a *captureAddon) Response(f *mitmproxy.Flow) { + if state := a.state(f); state != nil { + if f.Request != nil { + state.setRequestBody(f.Request.Body) } - flow.ContentType = f.Response.Header.Get("Content-Type") - if len(f.Response.Body) > 0 { - flow.Response.Body = snip(f.Response.Body, maxBodySnip) + if f.Response != nil { + state.setResponseMeta(f.Response.StatusCode, f.Response.Header) + state.setResponseBody(f.Response.Body) } - flow.Complete = f.Response.StatusCode != 0 + state.finish(nil) } - a.hub.ingest(flow) +} + +func (a *captureAddon) StreamRequestModifier(f *mitmproxy.Flow, in io.Reader) io.Reader { + if state := a.state(f); state != nil { + return state.requestReader(in) + } + return in +} + +func (a *captureAddon) StreamResponseModifier(f *mitmproxy.Flow, in io.Reader) io.Reader { + if state := a.state(f); state != nil { + return state.responseReader(in) + } + return in } func (a *captureAddon) RequestError(f *mitmproxy.Flow, err error) { - var dur time.Duration - if start, ok := a.pending.LoadAndDelete(f.Id.String()); ok { - if t, ok := start.(time.Time); ok { - dur = time.Since(t) - } - } - a.hub.ingest(Flow{ - Exchange: traffic.Exchange{ - Request: traffic.Request{ - Method: f.Request.Method, - URL: f.Request.URL.String(), - Protocol: f.Request.Proto, - }, - Error: err.Error(), - }, - Timestamp: f.StartTime, - ToolID: toolIDOf(f), - Host: f.Request.URL.Hostname(), - Duration: dur, - }) -} - -func snip(b []byte, max int) []byte { - if len(b) > max { - b = b[:max] - } - out := make([]byte, len(b)) - copy(out, b) - return out -} - -// pairsFromHTTP flattens an http.Header into the canonical pair sequence. The -// wire order is already lost inside net/http, so names are sorted to keep the -// stored form deterministic. -func pairsFromHTTP(headers http.Header) []traffic.Pair { - if len(headers) == 0 { + if state := a.state(f); state != nil { + state.finish(err) + } +} + +func (a *captureAddon) HTTPConnectError(f *mitmproxy.Flow, err error) { + if state := a.state(f); state != nil { + state.finish(err) + } else if f != nil { + // CONNECT failures can occur before the normal HTTP exchange starts. + state := newCaptureState(a.hub, f) + state.owner = a + state.finish(err) + } +} + +func (a *captureAddon) SSEEnd(f *mitmproxy.Flow) { + if state := a.state(f); state != nil { + state.finish(nil) + } +} + +func (a *captureAddon) WebSocketEnd(f *mitmproxy.Flow) { + // WebSocket messages have a separate traffic.WebSocketExchange model. The + // HTTP capture state must still be released when the upgraded connection + // ends, otherwise a long-lived socket leaks its pending entry. + if f != nil { + a.pending.Delete(f.Id.String()) + } +} + +func (a *captureAddon) state(f *mitmproxy.Flow) *captureState { + if f == nil { return nil } - names := make([]string, 0, len(headers)) - for name := range headers { - names = append(names, name) + if value, ok := a.pending.Load(f.Id.String()); ok { + state, ok := value.(*captureState) + if ok { + return state + } + } + return nil +} + +type captureState struct { + owner *captureAddon + hub *ProxyHub + proxy string + start time.Time + + mu sync.Mutex + finished bool + flow Flow + reqSink *traffic.BodySink + respSink *traffic.BodySink + captureErr error + reqCaptured bool + respCaptured bool +} + +func newCaptureState(hub *ProxyHub, f *mitmproxy.Flow) *captureState { + flow := Flow{Timestamp: f.StartTime, ToolID: toolIDOf(f)} + if f.ConnContext != nil && f.ConnContext.ClientConn != nil { + flow.TLS = f.ConnContext.ClientConn.Tls + } + if f.Request != nil { + flow.ID = f.Id.String() + flow.Request = traffic.Request{ + Method: f.Request.Method, + URL: f.Request.URL.String(), + Protocol: f.Request.Proto, + Headers: traffic.PairsFromHTTP(f.Request.Header), + } + flow.Host = f.Request.URL.Hostname() + } + return &captureState{hub: hub, owner: nil, proxy: f.Id.String(), start: f.StartTime, flow: flow} +} + +func (s *captureState) setRequestBody(body []byte) { + if len(body) == 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.finished || s.reqCaptured { + return + } + s.reqCaptured = true + if s.reqSink == nil { + var err error + s.reqSink, err = s.hub.store.bodySink(s.proxy, "req") + if err != nil { + s.captureErr = err + } + } + if s.reqSink != nil { + _, _ = s.reqSink.Write(body) + s.flow.Request.Body = s.reqSink.Preview() + return + } + s.flow.Request.Body = appendPreview(s.flow.Request.Body, body, maxBodySnip) +} + +func (s *captureState) setResponseMeta(status int, headers http.Header) { + s.mu.Lock() + defer s.mu.Unlock() + if s.finished { + return + } + var body []byte + var bodyRef *traffic.BodyRef + if s.flow.Response != nil { + body = s.flow.Response.Body + bodyRef = s.flow.Response.BodyRef + } + s.flow.Response = &traffic.Response{ + StatusCode: status, Headers: traffic.PairsFromHTTP(headers), + Body: body, BodyRef: bodyRef, + } + s.flow.ContentType = headers.Get("Content-Type") +} + +func (s *captureState) setResponseBody(body []byte) { + if len(body) == 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.finished || s.respCaptured { + return + } + s.respCaptured = true + if s.respSink == nil { + var err error + s.respSink, err = s.hub.store.bodySink(s.proxy, "resp") + if err != nil { + s.captureErr = err + } + } + if s.respSink != nil { + _, _ = s.respSink.Write(body) + if s.flow.Response == nil { + s.flow.Response = &traffic.Response{} + } + s.flow.Response.Body = s.respSink.Preview() + return + } + if s.flow.Response == nil { + s.flow.Response = &traffic.Response{} + } + s.flow.Response.Body = appendPreview(s.flow.Response.Body, body, maxBodySnip) +} + +func (s *captureState) requestReader(in io.Reader) io.Reader { + s.mu.Lock() + if s.reqCaptured { + s.mu.Unlock() + return in + } + if s.reqSink == nil { + var err error + s.reqSink, err = s.hub.store.bodySink(s.proxy, "req") + if err != nil { + s.captureErr = err + } + } + s.reqCaptured = true + sink := s.reqSink + s.mu.Unlock() + if sink == nil { + return &previewReader{src: in, add: func(p []byte) { + s.mu.Lock() + s.flow.Request.Body = appendPreview(s.flow.Request.Body, p, maxBodySnip) + s.mu.Unlock() + }} + } + return sink.Reader(in) +} + +func (s *captureState) responseReader(in io.Reader) io.Reader { + s.mu.Lock() + if s.respCaptured { + s.mu.Unlock() + return in + } + if s.respSink == nil { + var err error + s.respSink, err = s.hub.store.bodySink(s.proxy, "resp") + if err != nil { + s.captureErr = err + } + } + s.respCaptured = true + sink := s.respSink + s.mu.Unlock() + if sink == nil { + return &finishReader{src: &previewReader{src: in, add: func(p []byte) { + s.mu.Lock() + if s.flow.Response == nil { + s.flow.Response = &traffic.Response{} + } + s.flow.Response.Body = appendPreview(s.flow.Response.Body, p, maxBodySnip) + s.mu.Unlock() + }}, done: func(err error) { s.finish(err) }} + } + return &finishReader{src: sink.Reader(in), done: func(err error) { s.finish(err) }} +} + +func (s *captureState) finish(err error) { + s.mu.Lock() + if s.finished { + s.mu.Unlock() + return + } + s.finished = true + complete := err == nil && s.flow.Response != nil && s.flow.Response.StatusCode != 0 + if err == nil && s.captureErr != nil { + err = s.captureErr + } + if s.reqSink != nil { + ref, closeErr := s.reqSink.Close(complete) + s.flow.Request.BodyRef = &ref + s.flow.Request.Body = s.reqSink.Preview() + if closeErr != nil && err == nil { + err = closeErr + } + } + if s.respSink != nil { + ref, closeErr := s.respSink.Close(complete) + if s.flow.Response == nil { + s.flow.Response = &traffic.Response{} + } + s.flow.Response.BodyRef = &ref + s.flow.Response.Body = s.respSink.Preview() + if closeErr != nil && err == nil { + err = closeErr + } + } + if err != nil { + // A body write/close failure is part of the observation outcome. Do not + // publish a complete exchange whose file-backed payload is incomplete. + complete = false + } + if err != nil { + s.flow.Error = err.Error() + } + s.flow.Complete = complete + s.flow.Duration = time.Since(s.start) + flow := s.flow + s.mu.Unlock() + s.hub.ingest(flow) + // Keep the pending map bounded even when mitmproxy does not issue a later + // lifecycle callback for a failed/streaming connection. + if s.owner != nil { + s.owner.pending.Delete(s.proxy) + } +} + +func (s *captureState) discard() { + s.mu.Lock() + if s.reqSink != nil { + _ = s.reqSink.Discard() + } + if s.respSink != nil { + _ = s.respSink.Discard() } - sort.Strings(names) - out := make([]traffic.Pair, 0, len(headers)) - for _, name := range names { - for _, value := range headers[name] { - out = append(out, traffic.Pair{Name: name, Value: value}) + s.finished = true + s.mu.Unlock() + if s.owner != nil { + s.owner.pending.Delete(s.proxy) + } +} + +func appendPreview(dst, src []byte, max int) []byte { + if len(dst) >= max || len(src) == 0 { + return dst + } + if len(src) > max-len(dst) { + src = src[:max-len(dst)] + } + return append(dst, src...) +} + +type previewReader struct { + src io.Reader + add func([]byte) +} + +func (r *previewReader) Read(p []byte) (int, error) { + n, err := r.src.Read(p) + if n > 0 { + r.add(p[:n]) + } + return n, err +} + +type finishReader struct { + src io.Reader + done func(error) + once sync.Once +} + +func (r *finishReader) Read(p []byte) (int, error) { + n, err := r.src.Read(p) + if err != nil { + finishErr := err + if err == io.EOF { + finishErr = nil } + r.once.Do(func() { r.done(finishErr) }) } - return out + return n, err } // --------------------------------------------------------------------------- @@ -299,6 +584,17 @@ type Flow struct { TLS bool } +// bodySink creates a file-backed capture when the runner configured a body +// directory. Tests and embedded users can leave it empty and retain the +// bounded in-memory preview behavior. +func (s *FlowStore) bodySink(proxyID, side string) (*traffic.BodySink, error) { + dir := s.BodyDir() + if dir == "" { + return nil, nil + } + return traffic.NewBodySink(filepath.Join(dir, "body"), proxyID+"."+side, maxBodySnip) +} + type QueryOpts struct { Host string Status string @@ -307,41 +603,246 @@ type QueryOpts struct { } type FlowStore struct { - mu sync.RWMutex - flows []Flow - seq int - cap int + mu sync.RWMutex + flows []Flow + head int + size int + seq int + cap int + bodyDir string + indexPath string + indexFile *os.File + indexMu sync.Mutex + indexErr error } func NewFlowStore(cap int) *FlowStore { if cap <= 0 { cap = 10000 } - return &FlowStore{flows: make([]Flow, 0, 256), cap: cap} + return &FlowStore{flows: make([]Flow, cap), cap: cap} +} + +// SetBodyDir enables disk-backed request/response bodies for flows captured by +// this store. The directory is intentionally configured by the runner rather +// than by the traffic protocol, keeping the storage policy local to the tool. +func (s *FlowStore) SetBodyDir(dir string) error { + if dir == "" { + return nil + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + indexPath := filepath.Join(dir, "flows.jsonl") + if err := s.loadIndex(indexPath); err != nil { + return err + } + file, err := os.OpenFile(indexPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600) + if err != nil { + return fmt.Errorf("proxy flow store: open metadata index: %w", err) + } + s.mu.Lock() + if s.indexFile != nil { + s.indexMu.Lock() + _ = s.indexFile.Close() + s.indexMu.Unlock() + } + s.bodyDir = dir + s.indexPath = indexPath + s.indexFile = file + s.mu.Unlock() + return nil +} + +func (s *FlowStore) BodyDir() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.bodyDir +} + +// IndexError reports a metadata append failure. The in-memory/ring capture is +// still usable when the optional index cannot be written, but callers can +// surface this diagnostic instead of mistaking the index for durable storage. +func (s *FlowStore) IndexError() error { + s.indexMu.Lock() + defer s.indexMu.Unlock() + return s.indexErr +} + +// Sequence returns the newest assigned flow id. It does not change when the +// ring is cleared, so a reconnecting consumer can safely use it as a cursor. +func (s *FlowStore) Sequence() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.seq +} + +// after returns the ordered flows whose ids are greater than id. The ring is +// deliberately the source of truth for replay; callers that ask for an id +// older than the retained window receive the oldest retained flow onward. +func (s *FlowStore) after(id int) []Flow { + s.mu.RLock() + defer s.mu.RUnlock() + result := make([]Flow, 0, s.size) + for n := 0; n < s.size; n++ { + idx := (s.head + n) % s.cap + flow := s.flows[idx] + if flowSequence(flow.ID) > id { + result = append(result, flow) + } + } + return result } +// After returns a replay window for callers that need to recover from a +// reconnect. The returned slice is ordered by the store's monotonic id. +func (s *FlowStore) After(id int) []Flow { return s.after(id) } + // Add stores f, assigns it a monotonic ID, and returns the stored copy so the // caller can fan the ID-bearing flow out to subscribers. func (s *FlowStore) Add(f Flow) Flow { s.mu.Lock() - defer s.mu.Unlock() s.seq++ f.ID = strconv.Itoa(s.seq) - if len(s.flows) >= s.cap { - copy(s.flows, s.flows[1:]) - s.flows[len(s.flows)-1] = f + idx := (s.head + s.size) % s.cap + if s.size == s.cap { + idx = s.head + s.head = (s.head + 1) % s.cap } else { - s.flows = append(s.flows, f) + s.size++ } + s.flows[idx] = f + s.mu.Unlock() + s.appendIndex(f) return f } +func (s *FlowStore) appendIndex(f Flow) { + s.mu.RLock() + file := s.indexFile + s.mu.RUnlock() + if file == nil { + return + } + record := map[string]any{ + "id": f.ID, "tool_id": f.ToolID, "timestamp": f.Timestamp, + "host": f.Host, "content_type": f.ContentType, "duration": int64(f.Duration), + "tls": f.TLS, "exchange": f.Exchange, + } + if f.Request.BodyRef != nil { + record["request_body_ref"] = f.Request.BodyRef + } + if f.Response != nil && f.Response.BodyRef != nil { + record["response_body_ref"] = f.Response.BodyRef + } + line, err := json.Marshal(record) + if err == nil { + line = append(line, '\n') + s.indexMu.Lock() + defer s.indexMu.Unlock() + if _, err = file.Write(line); err == nil { + err = file.Sync() + } + if err != nil && s.indexErr == nil { + s.indexErr = err + } + } +} + +func (s *FlowStore) loadIndex(path string) error { + file, err := os.Open(path) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return fmt.Errorf("proxy flow store: open metadata index: %w", err) + } + defer file.Close() + decoder := json.NewDecoder(file) + decoder.UseNumber() + for { + var record map[string]json.RawMessage + err := decoder.Decode(&record) + if err == io.EOF { + break + } + if err != nil { + // A torn final append must not make the whole capture unreadable. + break + } + var flow Flow + if err := decodeIndexRecord(record, &flow); err != nil { + continue + } + s.mu.Lock() + s.putLocked(flow) + s.mu.Unlock() + } + return nil +} + +func decodeIndexRecord(record map[string]json.RawMessage, flow *Flow) error { + decode := func(key string, dst any) error { + raw, ok := record[key] + if !ok { + return fmt.Errorf("missing %s", key) + } + return json.Unmarshal(raw, dst) + } + if err := decode("id", &flow.ID); err != nil { + return err + } + if err := decode("exchange", &flow.Exchange); err != nil { + return err + } + _ = decode("tool_id", &flow.ToolID) + _ = decode("timestamp", &flow.Timestamp) + _ = decode("host", &flow.Host) + _ = decode("content_type", &flow.ContentType) + var duration int64 + if decode("duration", &duration) == nil { + flow.Duration = time.Duration(duration) + } + _ = decode("tls", &flow.TLS) + if raw, ok := record["request_body_ref"]; ok { + var ref traffic.BodyRef + if json.Unmarshal(raw, &ref) == nil { + flow.Request.BodyRef = &ref + } + } + if raw, ok := record["response_body_ref"]; ok && flow.Response != nil { + var ref traffic.BodyRef + if json.Unmarshal(raw, &ref) == nil { + flow.Response.BodyRef = &ref + } + } + return nil +} + +func (s *FlowStore) putLocked(f Flow) { + if f.ID == "" { + return + } + if seq := flowSequence(f.ID); seq > s.seq { + s.seq = seq + } + idx := (s.head + s.size) % s.cap + if s.size == s.cap { + idx = s.head + s.head = (s.head + 1) % s.cap + } else { + s.size++ + } + s.flows[idx] = f +} + func (s *FlowStore) Query(opts QueryOpts) []Flow { s.mu.RLock() defer s.mu.RUnlock() - var result []Flow - for i := range s.flows { - f := &s.flows[i] + result := make([]Flow, 0, s.size) + for n := 0; n < s.size; n++ { + idx := (s.head + n) % s.cap + f := &s.flows[idx] if opts.Host != "" && !strings.Contains(strings.ToLower(f.Host), strings.ToLower(opts.Host)) { continue } @@ -363,28 +864,73 @@ func (s *FlowStore) Query(opts QueryOpts) []Flow { func (s *FlowStore) Get(id int) *Flow { s.mu.RLock() - defer s.mu.RUnlock() want := strconv.Itoa(id) - for i := range s.flows { - if s.flows[i].ID == want { - f := s.flows[i] + for n := 0; n < s.size; n++ { + idx := (s.head + n) % s.cap + if s.flows[idx].ID == want { + f := s.flows[idx] + s.mu.RUnlock() + f.Exchange = f.Clone() + _ = f.HydrateBodies() return &f } } + s.mu.RUnlock() return nil } func (s *FlowStore) Clear() { s.mu.Lock() - defer s.mu.Unlock() - s.flows = s.flows[:0] - s.seq = 0 + bodyDir := s.bodyDir + indexFile := s.indexFile + indexPath := s.indexPath + s.indexFile = nil + for i := range s.flows { + s.flows[i] = Flow{} + } + s.head = 0 + s.size = 0 + s.mu.Unlock() + s.indexMu.Lock() + s.indexErr = nil + s.indexMu.Unlock() + if indexFile != nil { + s.indexMu.Lock() + _ = indexFile.Close() + s.indexMu.Unlock() + } + if bodyDir != "" { + _ = os.RemoveAll(bodyDir) + _ = os.MkdirAll(bodyDir, 0o755) + if indexPath != "" { + if file, err := os.OpenFile(indexPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600); err == nil { + s.mu.Lock() + s.indexFile = file + s.mu.Unlock() + } + } + } +} + +// Close releases the append-only metadata handle. Body files are deliberately +// retained so a caller can inspect a capture after the proxy listener stops. +func (s *FlowStore) Close() error { + s.mu.Lock() + file := s.indexFile + s.indexFile = nil + s.mu.Unlock() + if file == nil { + return nil + } + s.indexMu.Lock() + defer s.indexMu.Unlock() + return file.Close() } func (s *FlowStore) Count() int { s.mu.RLock() defer s.mu.RUnlock() - return len(s.flows) + return s.size } func matchStatus(code int, pattern string) bool { diff --git a/tools/proxy/mitm_test.go b/tools/proxy/mitm_test.go index 4f46bc66..dafa34e4 100644 --- a/tools/proxy/mitm_test.go +++ b/tools/proxy/mitm_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "net/url" "runtime" + "strconv" "strings" "sync" "sync/atomic" @@ -32,6 +33,71 @@ func startTestTarget(bodySize int) *httptest.Server { })) } +func TestLargeResponseIsStreamedToBodyFileAndHydratedOnGet(t *testing.T) { + body := strings.Repeat("streamed-body-", 32*1024) + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/plain") + _, _ = io.WriteString(w, body) + })) + defer target.Close() + hub := startHub(t, true) + + resp, err := hubClient(t, hub, "large-body").Get(target.URL) + if err != nil { + t.Fatal(err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + deadline := time.Now().Add(3 * time.Second) + var flows []Flow + for time.Now().Before(deadline) { + flows = hub.Store().Query(QueryOpts{}) + if len(flows) > 0 { + break + } + time.Sleep(10 * time.Millisecond) + } + if len(flows) != 1 { + t.Fatalf("captured flows = %d, want 1", len(flows)) + } + if got := len(flows[0].Response.Body); got > maxBodySnip { + t.Fatalf("preview size = %d, want <= %d", got, maxBodySnip) + } + id, err := strconv.Atoi(flows[0].ID) + if err != nil { + t.Fatal(err) + } + full := hub.Store().Get(id) + if full == nil || full.Response == nil { + t.Fatal("hydrated flow missing response") + } + if got := string(full.Response.Body); got != body { + t.Fatalf("hydrated body length/content mismatch: got %d want %d", len(got), len(body)) + } + wire := flowToProto(&flows[0]) + if got := string(wire.GetResponse().GetBody()); got != body { + t.Fatalf("wire body length/content mismatch: got %d want %d", len(got), len(body)) + } +} + +func TestCaptureFilterRunsBeforeStore(t *testing.T) { + target := startTestTarget(32) + defer target.Close() + hub := startHub(t, true) + hub.SetCaptureFilter(&traffic.FlowFilter{Status: "404"}) + resp, err := hubClient(t, hub, "filter").Get(target.URL) + if err != nil { + t.Fatal(err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + time.Sleep(100 * time.Millisecond) + if got := hub.Store().Count(); got != 0 { + t.Fatalf("filtered flow count = %d, want 0", got) + } +} + // newCapturingHub wraps a store in a recording ProxyHub so a bare captureAddon // can route flows through hub.ingest in tests without starting the hub's own // listener (the test attaches the addon to its own proxy). diff --git a/tools/proxy/traffic_handler.go b/tools/proxy/traffic_handler.go index cb6d33c1..5cfee309 100644 --- a/tools/proxy/traffic_handler.go +++ b/tools/proxy/traffic_handler.go @@ -70,6 +70,7 @@ func (h *TrafficHandler) handleConfigure(ctx context.Context, env *aop.Envelope, if cap := cfg.GetCapture(); cap != nil && cap.GetMode() != traffic.CaptureMode_CAPTURE_MODE_UNSPECIFIED { record := cap.GetMode() == traffic.CaptureMode_CAPTURE_MODE_RECORD h.infra.Hub.SetCapture(record, cap.GetDecryptHttps()) + h.infra.Hub.SetCaptureFilter(cap.GetFilter()) if record && cap.GetStream() { h.startStream(ctx, env.Id, send) } else {