From 43ccaea10fcf05e0384618b1841513f86a54fb49 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Mon, 17 Aug 2026 01:03:16 +0800 Subject: [PATCH 1/3] ci: reduce GitHub Actions usage --- .actrc | 4 + .github/dependabot.yml | 14 +- .github/workflows/ci.yml | 440 +++-------------------- .github/workflows/go-release.yml | 25 -- .github/workflows/nightly.yml | 71 ---- .github/workflows/scanner-regression.yml | 33 +- agent/tmux/manager_test.go | 12 +- 7 files changed, 104 insertions(+), 495 deletions(-) create mode 100644 .actrc delete mode 100644 .github/workflows/nightly.yml 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/.github/dependabot.yml b/.github/dependabot.yml index 22529caa..f0817478 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -7,19 +7,16 @@ updates: day: monday time: "04:00" timezone: Asia/Shanghai - open-pull-requests-limit: 10 + open-pull-requests-limit: 3 labels: - dependencies - go commit-message: prefix: deps groups: - golang-x: + all-go-dependencies: patterns: - - "golang.org/x/*" - chainreactors: - patterns: - - "github.com/chainreactors/*" + - "*" - package-ecosystem: github-actions directory: "/" @@ -28,8 +25,13 @@ updates: day: monday time: "04:30" timezone: Asia/Shanghai + open-pull-requests-limit: 1 labels: - dependencies - github-actions commit-message: prefix: deps + groups: + all-github-actions: + patterns: + - "*" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e5653a77..68d2631d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,76 +15,42 @@ concurrency: cancel-in-progress: true jobs: - # ── Fast gates (independent, no deps) ────────────────────────── - - lint: + checks: runs-on: ubuntu-22.04 steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true - - - - name: Run golangci-lint - uses: golangci/golangci-lint-action@v9.2.1 - with: - version: v2.12.2 - args: --timeout=5m --build-tags "re2_cgo re2_static" - - tidy: - runs-on: ubuntu-22.04 - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Check go mod tidy run: | - cp go.mod go.mod.orig - cp go.sum go.sum.orig - go mod tidy - if ! diff -q go.mod go.mod.orig >/dev/null 2>&1; then - echo "::error::go.mod is not tidy. Run 'go mod tidy' and commit the result." - diff go.mod.orig go.mod || true - exit 1 - fi - if ! diff -q go.sum go.sum.orig >/dev/null 2>&1; then - echo "::error::go.sum is not tidy. Run 'go mod tidy' and commit the result." - diff go.sum.orig go.sum | head -30 || true + cp go.mod "$RUNNER_TEMP/go.mod.before" + cp go.sum "$RUNNER_TEMP/go.sum.before" + for attempt in 1 2 3; do + if go mod tidy; then + break + fi + if [[ "$attempt" == "3" ]]; then + exit 1 + fi + echo "go mod tidy failed, retrying ($attempt/3)..." + sleep $((attempt * 5)) + done + if ! cmp -s go.mod "$RUNNER_TEMP/go.mod.before" || \ + ! cmp -s go.sum "$RUNNER_TEMP/go.sum.before"; then + echo "::error::go.mod or go.sum is not tidy. Run 'go mod tidy' and commit the result." + diff -u "$RUNNER_TEMP/go.mod.before" go.mod || true + diff -u "$RUNNER_TEMP/go.sum.before" go.sum | head -30 || true exit 1 fi - quality: - runs-on: ubuntu-22.04 - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - name: Check dependency layers, registered skips, and repository debt run: go test -count=1 ./core/deps @@ -96,44 +62,51 @@ jobs: exit 1 fi - - name: Run go vet - run: go vet ./... - - name: Check whitespace and submodule pins run: | git diff --check HEAD - git diff --exit-code --submodule=diff - git submodule foreach --recursive 'test -z "$(git status --porcelain)"' + git diff --ignore-submodules=dirty --exit-code -- \ + .gitmodules templates web/frontend/cyber-ui - # ── Unit tests (depends on tidy) ────────────────────────────── + - name: Run golangci-lint + uses: golangci/golangci-lint-action@v9.2.1 + with: + version: v2.12.2 + args: --timeout=8m --build-tags "re2_cgo re2_static" + skip-cache: ${{ env.ACT == 'true' }} test: runs-on: ubuntu-22.04 - needs: [tidy, quality] + needs: checks steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true - + cache: ${{ env.ACT != 'true' }} - name: Generate embedded resources run: go generate ./core/resources/... - name: Run unit tests with coverage run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m \ + test_args=(-timeout 5m) + if [[ "${ACT:-}" == "true" ]]; then + test_args=(-timeout 10m -p 4) + fi + go test -tags "re2_cgo re2_static" -race -count=1 "${test_args[@]}" \ -coverprofile=coverage.out \ -covermode=atomic \ ./... + - name: Check generated resources are committed + run: git diff --exit-code -- core/resources/template.go + - name: Display coverage summary if: always() run: | @@ -146,70 +119,27 @@ jobs: fi - name: Upload coverage artifact - if: always() + if: ${{ always() && env.ACT != 'true' }} uses: actions/upload-artifact@v7 with: name: coverage-report path: coverage.out retention-days: 14 - # ── Proxy & TMux tool tests (depends on tidy) ───────────────── - - tool-tests: - runs-on: ubuntu-22.04 - needs: tidy - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - - name: Run proxy tool tests - run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ - ./tools/proxy/ - - - name: Run tmux command tests - run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ - -run 'Tmux|BashProxy' \ - ./pkg/commands/ - - - name: Run PTY interactive session tests - run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ - -run 'MultiRound|SendCtrlC' \ - ./agent/tmux/ - - - name: Run agent tmux integration tests - run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ - -run 'AgentTmux' \ - ./agent/ - windows-test: runs-on: windows-2022 - needs: [tidy, quality] + needs: test steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Set up mingw for libcstx shell: bash @@ -225,52 +155,20 @@ jobs: env: CGO_ENABLED: "1" - race-stress: - runs-on: ubuntu-22.04 - needs: [tidy, quality] - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - name: Repeat agent and runner concurrency tests - run: | - go test -race -count=20 -timeout 15m \ - -run 'Test(ConcurrentEmitWhileRegistering|SetProviderRaceWithRun|ResetDoesNotAllowConcurrentPrompt|StreamingProviderEmitsMessageUpdates)$' \ - ./agent/... - go test -race -count=20 -timeout 15m \ - -run 'Test(StdioSameSessionFIFOOrder|StdioSessionsRunConcurrently|StdioDrainWaitsForInFlightAndQueued|RuntimeSessionDirectLoopUsesSessionScheduler|RuntimeSessionRejectsRequestsPastPendingLimit|SessionContextCancellationStopsActiveRun|ActiveRunSteersAsyncInputWithoutSecondLifecycle)$' \ - ./pkg/runner/... - - - name: Repeat web SSE, cancellation, and reload concurrency tests - run: | - go test -race -count=20 -timeout 20m \ - -run 'Test(BroadcastAOPEventPersistsRawEnvelope|ServeSSEWithSnapshotSubscribesBeforeReadingSnapshot|ServeSSEWithSnapshotDropsQueuedSnapshotDuplicates|SessionEventsReplayHasNoSideEffects|SessionEventsResumesAfterLastEventID|CancelRemoteScanStopsAgentAndPreservesCanceledStatus|CancelQueuedScanDoesNotWaitForConcurrencySlot|CancelTaskUsesControlChannelWhenTaskQueueIsFull|CancelTaskWaitsForSaturatedControlChannel|CompleteJobCannotOverwriteCanceledScan|BroadcastConfigReload|BroadcastConfigReloadWaitsBehindCancellationFrames|HandleConfigReloadResultUpdatesAgentStatus|SaveConfigBuildFailureKeepsCommittedConfigAndCurrentApp|SaveConfigCommitFailureClosesCandidateAndKeepsCurrentApp|SaveConfigSerializesConcurrentCandidates)$' \ - ./pkg/web - scanner-functional: runs-on: ubuntu-22.04 - needs: tidy + needs: test steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Run scanner functional regressions run: | @@ -303,20 +201,19 @@ jobs: headless-record-replay-e2e: runs-on: ubuntu-22.04 - needs: tidy + needs: test timeout-minutes: 10 steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Set up Chrome uses: browser-actions/setup-chrome@v2 @@ -347,57 +244,26 @@ jobs: -run '^TestE2EKatanaDeepRendersAuthenticatedSPA$' \ ./tools/scan - # ── Generated templates tests (depends on tidy) ─────────────── - - generated-test: - runs-on: ubuntu-22.04 - needs: tidy - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - - name: Run go generate for templates - run: go generate ./core/resources/... - - - name: Run resources tests - run: | - go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m \ - ./core/resources/... - - - name: Check generated resources are committed - run: git diff --exit-code - - protobuf-generated: + e2e: runs-on: ubuntu-22.04 - needs: tidy + needs: test steps: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Set up Node.js uses: actions/setup-node@v6 with: node-version: 22 - cache: npm + cache: ${{ env.ACT != 'true' && 'npm' || '' }} cache-dependency-path: web/frontend/package-lock.json - name: Set up protoc 35.1 @@ -413,46 +279,18 @@ jobs: run: go run ./cmd/gen - name: Check generated protobuf bindings are committed - run: git diff --exit-code - - # ── E2E tests (depends on test) ─────────────────────────────── - - e2e: - runs-on: ubuntu-22.04 - needs: test - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - name: Set up Node.js - uses: actions/setup-node@v6 - with: - node-version: 22 - cache: npm - cache-dependency-path: web/frontend/package-lock.json + run: | + git diff --exit-code -- \ + pkg/rpc \ + pkg/types \ + web/frontend/src/gen \ + web/frontend/cyber-ui/packages/aop/src/gen/aop - name: Build embedded frontend run: | - npm --prefix web/frontend ci npm --prefix web/frontend run build test -s web/static/index.html - - name: Upload embedded frontend - uses: actions/upload-artifact@v7 - with: - name: embedded-frontend - path: web/static - retention-days: 1 - - name: Install Playwright Chromium working-directory: web/frontend run: npx playwright install --with-deps chromium @@ -467,177 +305,9 @@ jobs: corepack pnpm install --frozen-lockfile corepack pnpm --filter @cyber/viewer test - - name: Run e2e tests + - name: Run backend E2E tests run: | go test -race -count=1 -timeout 10m \ -tags "e2e re2_cgo re2_static" \ -v \ ./pkg/web - - # ── Standard build: pure Go, no libcstx/CGO ────────────────── - - build-standard: - runs-on: ubuntu-22.04 - needs: test - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - name: Generate embedded resources - run: go generate ./core/resources - - - name: Build all standard platforms - run: | - for target in linux/amd64 linux/arm64 darwin/amd64 darwin/arm64 windows/amd64 windows/arm64; do - IFS='/' read -r goos goarch <<< "$target" - echo " compile ${goos}/${goarch}" - suffix="" - [[ "$goos" == "windows" ]] && suffix=".exe" - CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" \ - go build -trimpath -tags "forceposix emptytemplates noembed osusergo netgo" \ - -ldflags "-s -w" -buildvcs=false \ - -o "dist/standard_${goos}_${goarch}${suffix}" ./cmd/aiscan - done - test -f dist/standard_windows_amd64.exe - test -f dist/standard_windows_arm64.exe - ls -lh dist/ - - - name: Upload standard binaries - uses: actions/upload-artifact@v7 - with: - name: aiscan-standard - path: dist/standard_* - if-no-files-found: error - retention-days: 7 - - # ── Full build: native libcstx/CGO on supported platforms ──── - - build-full: - needs: [test, e2e] - runs-on: ${{ matrix.runner }} - defaults: - run: - shell: bash - strategy: - fail-fast: false - matrix: - include: - - id: linux-amd64 - runner: ubuntu-22.04 - goos: linux - goarch: amd64 - - id: linux-arm64 - runner: ubuntu-24.04-arm - goos: linux - goarch: arm64 - - id: darwin-amd64 - runner: macos-15-intel - goos: darwin - goarch: amd64 - - id: darwin-arm64 - runner: macos-15 - goos: darwin - goarch: arm64 - - id: windows-amd64 - runner: windows-2022 - goos: windows - goarch: amd64 - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - submodules: recursive - - - name: Set up Go - uses: actions/setup-go@v6 - with: - go-version-file: go.mod - cache: true - - - name: Set up mingw for libcstx - if: runner.os == 'Windows' - run: echo "C:/msys64/mingw64/bin" >> "$GITHUB_PATH" - - - name: Install recorder SDK link dependencies on Linux - if: runner.os == 'Linux' - run: | - sudo apt-get update - sudo apt-get install -y build-essential nasm yasm pkg-config \ - libxcb1-dev libxcb-shm0-dev libxcb-shape0-dev libxcb-xfixes0-dev - - - name: Install recorder SDK link dependencies on Windows - if: runner.os == 'Windows' - run: | - C:/msys64/usr/bin/bash.exe -lc \ - "pacman -S --noconfirm --needed git diffutils make nasm yasm pkgconf mingw-w64-x86_64-toolchain" - - - name: Prepare static FFmpeg and x264 recorder SDK - if: runner.os != 'macOS' - run: | - if [[ "${RUNNER_OS}" == "Windows" ]]; then - platform=windows - C:/msys64/usr/bin/bash.exe -lc "cd '${GITHUB_WORKSPACE}'; make record-native RECORD_ARCH='${{ matrix.goarch }}' || make record-native-source RECORD_ARCH='${{ matrix.goarch }}'" - else - platform=linux - make record-native RECORD_ARCH='${{ matrix.goarch }}' || make record-native-source RECORD_ARCH='${{ matrix.goarch }}' - fi - bash .github/native/sdk.sh env "${platform}" '${{ matrix.goarch }}' >> "${GITHUB_ENV}" - - - name: Verify recorder SDK link environment - if: runner.os != 'macOS' - run: pkg-config --modversion libavcodec - - - name: Download embedded frontend - uses: actions/download-artifact@v7 - with: - name: embedded-frontend - path: web/static - - - name: Generate embedded resources - run: go generate ./core/resources - - - name: Build full ${{ matrix.id }} - run: | - suffix="" - [[ "${{ matrix.goos }}" == "windows" ]] && suffix=".exe" - CGO_ENABLED=1 GOOS="${{ matrix.goos }}" GOARCH="${{ matrix.goarch }}" \ - go build -trimpath \ - -tags "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" \ - -ldflags "-s -w" -buildvcs=false \ - -o "dist/full_${{ matrix.goos }}_${{ matrix.goarch }}${suffix}" ./cmd/aiscan - - - name: Verify recorder libraries are statically linked - if: runner.os != 'macOS' - run: | - set -euo pipefail - binary="$(find dist -maxdepth 1 -type f -name 'full_*' -print -quit)" - test -n "${binary}" - if [[ "${RUNNER_OS}" == "Windows" ]]; then - if objdump -p "${binary}" | grep -Eiq 'DLL Name:.*(libav|x264|libwinpthread)'; then - echo "recorder library remained dynamically linked" >&2 - exit 1 - fi - else - if ldd "${binary}" | grep -Eiq '(libav|libx264)'; then - echo "recorder library remained dynamically linked" >&2 - exit 1 - fi - fi - - - name: Upload full ${{ matrix.id }} binary - uses: actions/upload-artifact@v7 - with: - name: aiscan-full-${{ matrix.id }} - path: dist/full_* - if-no-files-found: error - retention-days: 7 diff --git a/.github/workflows/go-release.yml b/.github/workflows/go-release.yml index 893719a5..7d6cfe58 100644 --- a/.github/workflows/go-release.yml +++ b/.github/workflows/go-release.yml @@ -20,23 +20,6 @@ on: required: false default: false type: boolean - workflow_call: - inputs: - tag: - description: 'Release tag' - required: true - type: string - target: - description: 'Branch or commit to tag if the tag does not exist' - required: false - default: master - type: string - prerelease: - description: 'Publish immediately as a prerelease instead of creating a draft' - required: false - default: false - type: boolean - permissions: contents: write @@ -52,9 +35,6 @@ jobs: # ── Resolve the tag once, share with all jobs ─────────────────── prepare: - # Nightly tag pushes also match v*.*.*. The nightly workflow invokes this - # workflow explicitly with prerelease=true, so ignore the duplicate push. - if: github.event_name != 'push' || !contains(github.ref_name, '-nightly.') runs-on: ubuntu-22.04 outputs: tag: ${{ steps.tag.outputs.tag }} @@ -82,11 +62,6 @@ jobs: echo "Invalid release tag: ${TAG}" >&2 exit 1 fi - if [[ "${TAG}" == *nightly* && "${{ inputs.prerelease }}" != "true" ]]; then - echo "Nightly tags require prerelease=true: ${TAG}" >&2 - exit 1 - fi - git fetch --force --tags origin if git rev-parse -q --verify "refs/tags/${TAG}" >/dev/null; then diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml deleted file mode 100644 index c10d99e5..00000000 --- a/.github/workflows/nightly.yml +++ /dev/null @@ -1,71 +0,0 @@ -name: nightly - -on: - schedule: - - cron: '0 16 * * *' # UTC 16:00 = CST 00:00 - workflow_dispatch: - -permissions: - contents: write - -concurrency: - group: nightly - cancel-in-progress: false - -jobs: - prepare: - runs-on: ubuntu-22.04 - outputs: - tag: ${{ steps.nightly.outputs.tag }} - target: ${{ steps.nightly.outputs.target }} - steps: - - name: Checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - token: ${{ secrets.GITHUB_TOKEN }} - submodules: recursive - - - name: Set nightly tag - id: nightly - run: | - DATE=$(date -u +%Y%m%d) - TAG="v0.0.0-nightly.${DATE}" - echo "TAG=${TAG}" >> "${GITHUB_ENV}" - echo "tag=${TAG}" >> "${GITHUB_OUTPUT}" - echo "target=${GITHUB_SHA}" >> "${GITHUB_OUTPUT}" - - - name: Create nightly tag - run: | - git config user.name "github-actions[bot]" - git config user.email "github-actions[bot]@users.noreply.github.com" - git tag -f "$TAG" - git push --force origin "$TAG" - - release: - needs: prepare - uses: ./.github/workflows/go-release.yml - with: - tag: ${{ needs.prepare.outputs.tag }} - target: ${{ needs.prepare.outputs.target }} - prerelease: true - secrets: inherit - - cleanup: - needs: [prepare, release] - runs-on: ubuntu-22.04 - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - TAG: ${{ needs.prepare.outputs.tag }} - steps: - - name: Delete superseded nightly releases - run: | - gh release list --limit 50 --json tagName \ - | jq -r '.[] | select(.tagName | startswith("v0.0.0-nightly.")) | .tagName' \ - | while read -r tag; do - if [[ "${tag}" == "${TAG}" ]]; then - continue - fi - echo "Deleting release ${tag}" - gh release delete "${tag}" --yes --cleanup-tag || true - done diff --git a/.github/workflows/scanner-regression.yml b/.github/workflows/scanner-regression.yml index b7c255b0..e29fb80d 100644 --- a/.github/workflows/scanner-regression.yml +++ b/.github/workflows/scanner-regression.yml @@ -22,17 +22,46 @@ jobs: - name: Checkout uses: actions/checkout@v6 with: - fetch-depth: 0 submodules: recursive - name: Set up Go uses: actions/setup-go@v6 with: go-version-file: go.mod - cache: true + cache: ${{ env.ACT != 'true' }} - name: Run bounded public scanner regressions run: | go test -tags "full integration re2_cgo re2_static" -count=1 -timeout 8m -v \ -run 'Test(ScannerPublicIntegration|FullScannerPublicIntegration)$' \ ./tools + + race-stress: + runs-on: ubuntu-22.04 + timeout-minutes: 25 + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + submodules: recursive + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: ${{ env.ACT != 'true' }} + + - name: Repeat agent and runner concurrency tests + run: | + go test -race -count=20 -timeout 15m \ + -run 'Test(ConcurrentEmitWhileRegistering|SetProviderRaceWithRun|ResetDoesNotAllowConcurrentPrompt|StreamingProviderEmitsMessageUpdates)$' \ + ./agent/... + go test -race -count=20 -timeout 15m \ + -run 'Test(StdioSameSessionFIFOOrder|StdioSessionsRunConcurrently|StdioDrainWaitsForInFlightAndQueued|RuntimeSessionDirectLoopUsesSessionScheduler|RuntimeSessionRejectsRequestsPastPendingLimit|SessionContextCancellationStopsActiveRun|ActiveRunSteersAsyncInputWithoutSecondLifecycle)$' \ + ./pkg/runner/... + + - name: Repeat web SSE, cancellation, and reload concurrency tests + run: | + go test -race -count=20 -timeout 20m \ + -run 'Test(BroadcastAOPEventPersistsRawEnvelope|ServeSSEWithSnapshotSubscribesBeforeReadingSnapshot|ServeSSEWithSnapshotDropsQueuedSnapshotDuplicates|SessionEventsReplayHasNoSideEffects|SessionEventsResumesAfterLastEventID|CancelRemoteScanStopsAgentAndPreservesCanceledStatus|CancelQueuedScanDoesNotWaitForConcurrencySlot|CancelTaskUsesControlChannelWhenTaskQueueIsFull|CancelTaskWaitsForSaturatedControlChannel|CompleteJobCannotOverwriteCanceledScan|BroadcastConfigReload|BroadcastConfigReloadWaitsBehindCancellationFrames|HandleConfigReloadResultUpdatesAgentStatus|SaveConfigBuildFailureKeepsCommittedConfigAndCurrentApp|SaveConfigCommitFailureClosesCandidateAndKeepsCurrentApp|SaveConfigSerializesConcurrentCandidates)$' \ + ./pkg/web diff --git a/agent/tmux/manager_test.go b/agent/tmux/manager_test.go index 1be312f2..5c6abc2c 100644 --- a/agent/tmux/manager_test.go +++ b/agent/tmux/manager_test.go @@ -185,7 +185,7 @@ func TestPeekReturnsTail(t *testing.T) { } mgr := NewManager() dir := t.TempDir() - info, err := mgr.Create(dir, "for i in 1 2 3 4 5; do echo line$i; done", "peek-test", 5*time.Second, nil, "") + info, err := mgr.Create(dir, "for i in 1 2 3 4 5; do echo line$i; done; sleep 0.05", "peek-test", 5*time.Second, nil, "") if err != nil { t.Fatalf("Create: %v", err) } @@ -305,7 +305,7 @@ func TestCreateCmd(t *testing.T) { mgr := NewManager() dir := t.TempDir() - info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo from-createcmd"}, "cmd-test", 10*time.Second, nil, "") + info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo from-createcmd; sleep 0.05"}, "cmd-test", 10*time.Second, nil, "") if err != nil { t.Fatalf("CreateCmd: %v", err) } @@ -324,7 +324,7 @@ func TestCreateCmdWithEnv(t *testing.T) { mgr := NewManager() dir := t.TempDir() - info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo $TEST_MAGIC"}, "env-test", 10*time.Second, []string{"TEST_MAGIC=pty_works"}, "") + info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo $TEST_MAGIC; sleep 0.05"}, "env-test", 10*time.Second, []string{"TEST_MAGIC=pty_works"}, "") if err != nil { t.Fatalf("CreateCmd: %v", err) } @@ -370,7 +370,7 @@ func TestPeekNew(t *testing.T) { dir := t.TempDir() payload := strings.Repeat("x", 100) - info, err := mgr.Create(dir, "printf '"+payload+"'", "peeknew-test", 10*time.Second, nil, "") + info, err := mgr.Create(dir, "printf '"+payload+"'; sleep 0.05", "peeknew-test", 10*time.Second, nil, "") if err != nil { t.Fatalf("Create: %v", err) } @@ -473,7 +473,7 @@ func TestExecCommandDirect(t *testing.T) { mgr := NewManager() dir := t.TempDir() - info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo direct"}, "", 5*time.Second, nil, "") + info, err := mgr.CreateCmd(dir, "/bin/sh", []string{"-c", "echo direct; sleep 0.05"}, "", 5*time.Second, nil, "") if err != nil { t.Fatalf("CreateCmd: %v", err) } @@ -530,7 +530,7 @@ func TestPeekBytes(t *testing.T) { t.Skip("unix-only test") } - info, err := mgr.Create(dir, "printf '0123456789'", "peekbytes-test", 5*time.Second, nil, "") + info, err := mgr.Create(dir, "printf '0123456789'; sleep 0.05", "peekbytes-test", 5*time.Second, nil, "") if err != nil { t.Fatal(err) } From b0b1ffb06480b1d6589175454c4ad1f49554f532 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Sat, 22 Aug 2026 16:33:03 +0800 Subject: [PATCH 2/3] feat(proxy): stream MITM bodies into durable capture files --- aop/traffic/body.go | 200 +++++++++ aop/traffic/body_test.go | 29 ++ aop/traffic/exchange.go | 135 ++++++ aop/traffic/exchange_test.go | 15 + tools/curl/client.go | 24 +- tools/proxy/hub.go | 77 +++- tools/proxy/hub_traffic.go | 96 +++- tools/proxy/hub_traffic_test.go | 85 ++++ tools/proxy/mitm.go | 745 +++++++++++++++++++++++++++----- tools/proxy/mitm_test.go | 66 +++ tools/proxy/traffic_handler.go | 1 + 11 files changed, 1349 insertions(+), 124 deletions(-) create mode 100644 aop/traffic/body.go create mode 100644 aop/traffic/body_test.go 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..ae492fb7 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.Exchange.Clone() + _ = copy.Exchange.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..ae9bc767 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,378 @@ 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 { + return value.(*captureState) + } + 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.Exchange.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 + } } - 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.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 } } - return out + 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() + } + 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 n, err } // --------------------------------------------------------------------------- @@ -299,6 +581,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 +600,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 +861,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.Exchange.Clone() + _ = f.Exchange.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 { From a2fbb2be1d01aaf2bb7632edb02a3e5e5157a513 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Sun, 23 Aug 2026 18:46:33 +0800 Subject: [PATCH 3/3] fix(proxy): satisfy lint for durable capture --- tools/proxy/hub_traffic.go | 4 ++-- tools/proxy/mitm.go | 11 +++++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/tools/proxy/hub_traffic.go b/tools/proxy/hub_traffic.go index ae492fb7..1824f096 100644 --- a/tools/proxy/hub_traffic.go +++ b/tools/proxy/hub_traffic.go @@ -139,8 +139,8 @@ func flowToProto(flow *Flow) *traffic.Flow { // 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.Exchange.Clone() - _ = copy.Exchange.HydrateBodies() + copy.Exchange = flow.Clone() + _ = copy.HydrateBodies() message := copy.Proto() message.ToolId = flow.ToolID if !flow.Timestamp.IsZero() { diff --git a/tools/proxy/mitm.go b/tools/proxy/mitm.go index ae9bc767..7e9bfdb2 100644 --- a/tools/proxy/mitm.go +++ b/tools/proxy/mitm.go @@ -289,7 +289,10 @@ func (a *captureAddon) state(f *mitmproxy.Flow) *captureState { return nil } if value, ok := a.pending.Load(f.Id.String()); ok { - return value.(*captureState) + state, ok := value.(*captureState) + if ok { + return state + } } return nil } @@ -316,7 +319,7 @@ func newCaptureState(hub *ProxyHub, f *mitmproxy.Flow) *captureState { flow.TLS = f.ConnContext.ClientConn.Tls } if f.Request != nil { - flow.Exchange.ID = f.Id.String() + flow.ID = f.Id.String() flow.Request = traffic.Request{ Method: f.Request.Method, URL: f.Request.URL.String(), @@ -867,8 +870,8 @@ func (s *FlowStore) Get(id int) *Flow { if s.flows[idx].ID == want { f := s.flows[idx] s.mu.RUnlock() - f.Exchange = f.Exchange.Clone() - _ = f.Exchange.HydrateBodies() + f.Exchange = f.Clone() + _ = f.HydrateBodies() return &f } }